diff --git a/packages/server/src/server/agent/agent-registry.test.ts b/packages/server/src/server/agent/agent-registry.test.ts index ad2fbed73..a94d8a55b 100644 --- a/packages/server/src/server/agent/agent-registry.test.ts +++ b/packages/server/src/server/agent/agent-registry.test.ts @@ -1,37 +1,84 @@ -import { describe, expect, test, beforeEach, afterEach, vi } from "vitest"; +import { describe, expect, test, beforeEach, afterEach } from "vitest"; import os from "node:os"; import path from "node:path"; import { mkdtempSync, rmSync, writeFileSync, readFileSync } from "node:fs"; import { AgentRegistry } from "./agent-registry.js"; -import type { AgentSnapshot } from "./agent-manager.js"; +import type { ManagedAgent } from "./agent-manager.js"; +import type { + AgentPermissionRequest, + AgentSession, + AgentSessionConfig, +} from "./agent-sdk-types.js"; -function createSnapshot(overrides?: Partial): AgentSnapshot { - const now = new Date(); - return { - id: "agent-test", - provider: "claude", - cwd: "/tmp/project", - model: null, - createdAt: now, - updatedAt: now, - lastUserMessageAt: null, - status: "idle", - sessionId: null, - capabilities: { - supportsStreaming: true, - supportsSessionPersistence: true, - supportsDynamicModes: true, - supportsMcpServers: true, - supportsReasoningStream: true, - supportsToolInvocations: true, - }, - currentModeId: null, - availableModes: [], - pendingPermissions: [], - persistence: null, - ...overrides, +type ManagedAgentOverrides = Omit< + Partial, + "config" | "pendingPermissions" | "session" | "pendingRun" +> & { + config?: Partial; + pendingPermissions?: Map; + session?: AgentSession | null; + pendingRun?: ManagedAgent["pendingRun"]; +}; + +function createManagedAgent( + overrides: ManagedAgentOverrides = {} +): ManagedAgent { + const now = overrides.updatedAt ?? new Date("2025-01-01T00:00:00.000Z"); + const provider = overrides.provider ?? "claude"; + const cwd = overrides.cwd ?? "/tmp/project"; + const lifecycle = overrides.lifecycle ?? "idle"; + const configOverrides = overrides.config ?? {}; + const config: AgentSessionConfig = { + provider, + cwd, + modeId: configOverrides.modeId ?? "plan", + model: configOverrides.model ?? "gpt-5.1", + extra: configOverrides.extra ?? { claude: { maxThinkingTokens: 1024 } }, }; + const session = + lifecycle === "closed" + ? null + : overrides.session ?? ({} as AgentSession); + const pendingRun = + overrides.pendingRun ?? + (lifecycle === "running" ? (async function* noop() {})() : null); + + const agent: ManagedAgent = { + id: overrides.id ?? "agent-test", + provider, + cwd, + session, + sessionId: overrides.sessionId ?? "session-123", + capabilities: + overrides.capabilities ?? + { + supportsStreaming: true, + supportsSessionPersistence: true, + supportsDynamicModes: true, + supportsMcpServers: true, + supportsReasoningStream: true, + supportsToolInvocations: true, + }, + config, + lifecycle, + createdAt: overrides.createdAt ?? now, + updatedAt: overrides.updatedAt ?? now, + availableModes: overrides.availableModes ?? [], + currentModeId: overrides.currentModeId ?? config.modeId ?? null, + pendingPermissions: + overrides.pendingPermissions ?? + new Map(), + pendingRun, + timeline: overrides.timeline ?? [], + persistence: overrides.persistence ?? null, + historyPrimed: overrides.historyPrimed ?? true, + lastUserMessageAt: overrides.lastUserMessageAt ?? now, + lastUsage: overrides.lastUsage, + lastError: overrides.lastError, + }; + + return agent; } describe("AgentRegistry", () => { @@ -49,27 +96,21 @@ describe("AgentRegistry", () => { rmSync(tmpDir, { recursive: true, force: true }); }); - test("persists configs and snapshot metadata", async () => { + test("applySnapshot persists configs and snapshot metadata", async () => { await registry.applySnapshot( - createSnapshot({ + createManagedAgent({ id: "agent-1", cwd: "/tmp/project", currentModeId: "coding", - status: "idle", + lifecycle: "idle", + config: { + modeId: "coding", + model: "gpt-5.1", + extra: { claude: { maxThinkingTokens: 1024 } }, + }, }) ); - await registry.recordConfig( - "agent-1", - "claude", - "/tmp/project", - { - modeId: "coding", - model: "gpt-5.1", - extra: { claude: { maxThinkingTokens: 1024 } }, - } - ); - const records = await registry.list(); expect(records).toHaveLength(1); const [record] = records; @@ -87,7 +128,7 @@ describe("AgentRegistry", () => { test("stores titles independently of snapshots", async () => { await registry.applySnapshot( - createSnapshot({ + createManagedAgent({ id: "agent-2", provider: "codex", cwd: "/tmp/second", @@ -103,21 +144,30 @@ describe("AgentRegistry", () => { expect(persisted?.title).toBe("Fix Login Bug"); }); - test("recordConfig warns if no snapshot exists", async () => { - const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); - await registry.recordConfig( - "agent-3", - "claude", - "/tmp/project", - { - modeId: "plan", - } + test("applySnapshot preserves custom titles while updating metadata", async () => { + const agentId = "agent-3"; + await registry.applySnapshot( + createManagedAgent({ + id: agentId, + lifecycle: "idle", + currentModeId: "plan", + }) + ); + await registry.setTitle(agentId, "Important Bug Fix"); + + await registry.applySnapshot( + createManagedAgent({ + id: agentId, + lifecycle: "running", + currentModeId: "build", + updatedAt: new Date("2025-01-02T00:00:00.000Z"), + }) ); - const record = await registry.get("agent-3"); - expect(record).toBeNull(); - expect(warnSpy).toHaveBeenCalled(); - warnSpy.mockRestore(); + const record = await registry.get(agentId); + expect(record?.title).toBe("Important Bug Fix"); + expect(record?.lastModeId).toBe("build"); + expect(record?.lastStatus).toBe("running"); }); test("recovers from trailing garbage in agents.json", async () => { diff --git a/packages/server/src/server/agent/agent-registry.ts b/packages/server/src/server/agent/agent-registry.ts index 031db7757..bcba509ea 100644 --- a/packages/server/src/server/agent/agent-registry.ts +++ b/packages/server/src/server/agent/agent-registry.ts @@ -6,8 +6,9 @@ import { fileURLToPath } from "node:url"; import { z } from "zod"; import { AgentStatusSchema } from "../messages.js"; -import type { AgentSnapshot } from "./agent-manager.js"; -import type { AgentProvider, AgentSessionConfig } from "./agent-sdk-types.js"; +import { toStoredAgentRecord } from "./agent-projections.js"; +import type { ManagedAgent } from "./agent-manager.js"; +import type { AgentSessionConfig } from "./agent-sdk-types.js"; const SERIALIZABLE_CONFIG_SCHEMA = z .object({ @@ -103,77 +104,13 @@ export class AgentRegistry { await this.flush(); } - async recordConfig( - agentId: string, - provider: AgentProvider, - cwd: string, - config?: SerializableAgentConfig - ): Promise { + async applySnapshot(agent: ManagedAgent): Promise { await this.load(); - const existing = this.cache.get(agentId); - if (!existing) { - console.warn( - `[AgentRegistry] Cannot record config for ${agentId} because no snapshot has been persisted yet` - ); - return; - } - - const now = new Date().toISOString(); - const sanitizedConfig = config ? sanitizeConfig(config) : existing.config; - const nextModeId = - config?.modeId ?? existing.lastModeId ?? sanitizedConfig?.modeId ?? null; - - const updated: StoredAgentRecord = { - ...existing, - provider, - cwd, - updatedAt: now, - lastModeId: nextModeId, - config: sanitizedConfig, - }; - this.cache.set(agentId, updated); - await this.flush(); - } - - async applySnapshot(snapshot: AgentSnapshot): Promise { - await this.load(); - const now = new Date().toISOString(); - const existing = this.cache.get(snapshot.id); - if (!existing) { - const record: StoredAgentRecord = { - id: snapshot.id, - provider: snapshot.provider, - cwd: snapshot.cwd, - createdAt: now, - updatedAt: now, - lastActivityAt: snapshot.updatedAt.toISOString(), - lastUserMessageAt: snapshot.lastUserMessageAt - ? snapshot.lastUserMessageAt.toISOString() - : null, - title: null, - lastStatus: snapshot.status, - lastModeId: snapshot.currentModeId ?? null, - config: null, - persistence: snapshot.persistence ?? null, - }; - this.cache.set(snapshot.id, record); - await this.flush(); - return; - } - const updated: StoredAgentRecord = { - ...existing, - provider: snapshot.provider, - cwd: snapshot.cwd, - updatedAt: now, - lastActivityAt: snapshot.updatedAt.toISOString(), - lastUserMessageAt: snapshot.lastUserMessageAt - ? snapshot.lastUserMessageAt.toISOString() - : existing.lastUserMessageAt ?? null, - lastStatus: snapshot.status, - lastModeId: snapshot.currentModeId ?? null, - persistence: snapshot.persistence ?? existing.persistence ?? null, - }; - this.cache.set(snapshot.id, updated); + const existing = this.cache.get(agent.id); + const record = toStoredAgentRecord(agent, { + title: existing?.title ?? null, + }); + this.cache.set(agent.id, record); await this.flush(); } @@ -250,19 +187,6 @@ export class AgentRegistry { } } -function sanitizeConfig( - config: SerializableAgentConfig | undefined -): SerializableAgentConfig | undefined { - if (!config) { - return undefined; - } - const cleaned: SerializableAgentConfig = {}; - if (config.modeId) cleaned.modeId = config.modeId; - if (config.model) cleaned.model = config.model; - if (config.extra) cleaned.extra = JSON.parse(JSON.stringify(config.extra)); - return cleaned; -} - function resolveServerPackageRoot(): string { let currentDir = path.dirname(fileURLToPath(import.meta.url)); while (true) { diff --git a/packages/server/src/server/persistence-hooks.test.ts b/packages/server/src/server/persistence-hooks.test.ts index c8bf47d39..c2d212955 100644 --- a/packages/server/src/server/persistence-hooks.test.ts +++ b/packages/server/src/server/persistence-hooks.test.ts @@ -1,38 +1,85 @@ import { describe, expect, test, vi } from "vitest"; -import type { AgentSnapshot } from "./agent/agent-manager.js"; +import type { ManagedAgent } from "./agent/agent-manager.js"; import type { StoredAgentRecord } from "./agent/agent-registry.js"; import { attachAgentRegistryPersistence, restorePersistedAgents, } from "./persistence-hooks.js"; +import type { + AgentPermissionRequest, + AgentSession, + AgentSessionConfig, +} from "./agent/agent-sdk-types.js"; -function createSnapshot(overrides?: Partial): AgentSnapshot { - const now = new Date(); - return { - id: "agent-1", - provider: "claude", - cwd: "/tmp/project", - model: null, - createdAt: now, - updatedAt: now, - lastUserMessageAt: null, - status: "idle", - sessionId: null, - capabilities: { - supportsStreaming: true, - supportsSessionPersistence: true, - supportsDynamicModes: true, - supportsMcpServers: true, - supportsReasoningStream: true, - supportsToolInvocations: true, - }, - currentModeId: "plan", - availableModes: [], - pendingPermissions: [], - persistence: null, - ...overrides, +type ManagedAgentOverrides = Omit< + Partial, + "config" | "pendingPermissions" | "session" | "pendingRun" +> & { + config?: Partial; + pendingPermissions?: Map; + session?: AgentSession | null; + pendingRun?: ManagedAgent["pendingRun"]; +}; + +function createManagedAgent( + overrides: ManagedAgentOverrides = {} +): ManagedAgent { + const now = overrides.updatedAt ?? new Date("2025-01-01T00:00:00.000Z"); + const provider = overrides.provider ?? "claude"; + const cwd = overrides.cwd ?? "/tmp/project"; + const lifecycle = overrides.lifecycle ?? "idle"; + const configOverrides = overrides.config ?? {}; + const config: AgentSessionConfig = { + provider, + cwd, + modeId: configOverrides.modeId ?? "plan", + model: configOverrides.model ?? "claude-3.5-sonnet", + extra: configOverrides.extra ?? { claude: { tone: "focused" } }, }; + const session = + lifecycle === "closed" + ? null + : overrides.session ?? ({} as AgentSession); + const pendingRun = + overrides.pendingRun ?? + (lifecycle === "running" ? (async function* noop() {})() : null); + + const agent: ManagedAgent = { + id: overrides.id ?? "agent-1", + provider, + cwd, + session, + sessionId: overrides.sessionId ?? "session-123", + capabilities: + overrides.capabilities ?? + { + supportsStreaming: true, + supportsSessionPersistence: true, + supportsDynamicModes: true, + supportsMcpServers: true, + supportsReasoningStream: true, + supportsToolInvocations: true, + }, + config, + lifecycle, + createdAt: overrides.createdAt ?? now, + updatedAt: overrides.updatedAt ?? now, + availableModes: overrides.availableModes ?? [], + currentModeId: overrides.currentModeId ?? config.modeId ?? null, + pendingPermissions: + overrides.pendingPermissions ?? + new Map(), + pendingRun, + timeline: overrides.timeline ?? [], + persistence: overrides.persistence ?? null, + historyPrimed: overrides.historyPrimed ?? true, + lastUserMessageAt: overrides.lastUserMessageAt ?? now, + lastUsage: overrides.lastUsage, + lastError: overrides.lastError, + }; + + return agent; } function createRecord( @@ -138,13 +185,13 @@ describe("persistence hooks", () => { } as any); expect(agentManager.subscribe).toHaveBeenCalledTimes(1); - const snapshot = createSnapshot(); - subscriber({ type: "agent_state", agent: snapshot }); - expect(applySnapshot).toHaveBeenCalledWith(snapshot); + const agent = createManagedAgent(); + subscriber({ type: "agent_state", agent }); + expect(applySnapshot).toHaveBeenCalledWith(agent); subscriber({ type: "agent_stream", - agentId: snapshot.id, + agentId: agent.id, event: { type: "timeline", item: { type: "assistant_message", text: "hi" } }, }); expect(applySnapshot).toHaveBeenCalledTimes(1); diff --git a/packages/server/src/server/session.ts b/packages/server/src/server/session.ts index 3e0a57d07..f54938c24 100644 --- a/packages/server/src/server/session.ts +++ b/packages/server/src/server/session.ts @@ -1237,24 +1237,6 @@ export class Session { ); const snapshot = await this.agentManager.createAgent(sessionConfig); this.setCachedTitle(snapshot.id, null); - try { - await this.agentRegistry.recordConfig( - snapshot.id, - snapshot.provider, - snapshot.cwd, - { - modeId: sessionConfig.modeId, - model: sessionConfig.model, - extra: sessionConfig.extra, - } - ); - } catch (registryError) { - console.error( - `[Session ${this.clientId}] Failed to record agent config for ${snapshot.id}:`, - registryError - ); - } - await this.forwardAgentState(snapshot); const trimmedPrompt = initialPrompt?.trim();