From dfe43210512ce8674e581655dbbfee99fc33be35 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Tue, 2 Dec 2025 10:13:59 +0000 Subject: [PATCH] Ensure agent titles persist synchronously during creation --- .../src/server/agent/agent-manager.test.ts | 112 ++++++++++++++++++ .../server/src/server/agent/agent-manager.ts | 18 +++ .../src/server/agent/agent-registry.test.ts | 16 +++ .../server/src/server/agent/agent-registry.ts | 11 +- .../src/server/agent/agent-sdk-types.ts | 1 + .../src/server/agent/mcp-server.test.ts | 10 +- .../server/src/server/agent/mcp-server.ts | 11 +- packages/server/src/server/index.ts | 1 + packages/server/src/server/messages.ts | 7 ++ .../server/src/server/persistence-hooks.ts | 2 + 10 files changed, 173 insertions(+), 16 deletions(-) create mode 100644 packages/server/src/server/agent/agent-manager.test.ts diff --git a/packages/server/src/server/agent/agent-manager.test.ts b/packages/server/src/server/agent/agent-manager.test.ts new file mode 100644 index 000000000..8c8427f2d --- /dev/null +++ b/packages/server/src/server/agent/agent-manager.test.ts @@ -0,0 +1,112 @@ +import { describe, expect, test } from "vitest"; +import { mkdtempSync } from "node:fs"; +import { join } from "node:path"; +import { tmpdir } from "node:os"; +import { randomUUID } from "node:crypto"; + +import { AgentManager } from "./agent-manager.js"; +import { AgentRegistry } from "./agent-registry.js"; +import type { + AgentClient, + AgentRunResult, + AgentSession, + AgentSessionConfig, + AgentStreamEvent, +} from "./agent-sdk-types.js"; + +const TEST_CAPABILITIES = { + supportsStreaming: false, + supportsSessionPersistence: false, + supportsDynamicModes: false, + supportsMcpServers: false, + supportsReasoningStream: false, + supportsToolInvocations: false, +} as const; + +class TestAgentClient implements AgentClient { + readonly provider = "codex" as const; + readonly capabilities = TEST_CAPABILITIES; + + async createSession(config: AgentSessionConfig): Promise { + return new TestAgentSession(config); + } + + async resumeSession(config?: Partial): Promise { + return new TestAgentSession({ + provider: "codex", + cwd: config?.cwd ?? process.cwd(), + }); + } +} + +class TestAgentSession implements AgentSession { + readonly provider = "codex" as const; + readonly capabilities = TEST_CAPABILITIES; + readonly id = randomUUID(); + + constructor(private readonly config: AgentSessionConfig) {} + + async run(): Promise { + return { + sessionId: this.id ?? this.config.provider, + finalText: "", + timeline: [], + }; + } + + async *stream(): AsyncGenerator {} + + async *streamHistory(): AsyncGenerator {} + + async getAvailableModes() { + return []; + } + + async getCurrentMode() { + return null; + } + + async setMode(): Promise {} + + getPendingPermissions() { + return []; + } + + async respondToPermission(): Promise {} + + describePersistence() { + return null; + } + + async interrupt(): Promise {} + + async close(): Promise {} +} + +describe("AgentManager", () => { + test("createAgent persists provided title before returning", async () => { + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-test-")); + const registryPath = join(workdir, "agents.json"); + const registry = new AgentRegistry(registryPath); + const manager = new AgentManager({ + clients: { + codex: new TestAgentClient(), + }, + registry, + idFactory: () => "agent-with-title", + }); + + const snapshot = await manager.createAgent({ + provider: "codex", + cwd: workdir, + title: "Fix Login Bug", + }); + + expect(snapshot.id).toBe("agent-with-title"); + expect(snapshot.lifecycle).toBe("idle"); + + const persisted = await registry.get("agent-with-title"); + expect(persisted?.title).toBe("Fix Login Bug"); + expect(persisted?.id).toBe("agent-with-title"); + }); +}); diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index 5d1e6faf6..c6272dc57 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -21,6 +21,7 @@ import type { PersistedAgentDescriptor, } from "./agent-sdk-types.js"; import { resolveAgentModel } from "./model-resolver.js"; +import type { AgentRegistry } from "./agent-registry.js"; export const AGENT_LIFECYCLE_STATUSES = [ "initializing", @@ -52,6 +53,7 @@ export type AgentManagerOptions = { clients?: Partial>; maxTimelineItems?: number; idFactory?: () => string; + registry?: AgentRegistry; }; export type WaitForAgentOptions = { @@ -180,11 +182,13 @@ export class AgentManager { private readonly subscribers = new Set(); private readonly maxTimelineItems: number; private readonly idFactory: () => string; + private readonly registry?: AgentRegistry; constructor(options?: AgentManagerOptions) { this.maxTimelineItems = options?.maxTimelineItems ?? DEFAULT_MAX_TIMELINE_ITEMS; this.idFactory = options?.idFactory ?? (() => randomUUID()); + this.registry = options?.registry; if (options?.clients) { for (const [provider, client] of Object.entries(options.clients)) { if (client) { @@ -718,14 +722,28 @@ export class AgentManager { } as ActiveManagedAgent; this.agents.set(agentId, managed); + await this.persistSnapshot(managed, { + title: config.title ?? null, + }); this.emitState(managed); await this.refreshSessionState(managed); managed.lifecycle = "idle"; + await this.persistSnapshot(managed); this.emitState(managed); return this.createAgentView(managed); } + private async persistSnapshot( + agent: ManagedAgent, + options?: { title?: string | null } + ): Promise { + if (!this.registry) { + return; + } + await this.registry.applySnapshot(agent, options); + } + private async refreshSessionState(agent: ActiveManagedAgent): Promise { try { const modes = await agent.session.getAvailableModes(); diff --git a/packages/server/src/server/agent/agent-registry.test.ts b/packages/server/src/server/agent/agent-registry.test.ts index 0604e32af..7f1755693 100644 --- a/packages/server/src/server/agent/agent-registry.test.ts +++ b/packages/server/src/server/agent/agent-registry.test.ts @@ -168,6 +168,22 @@ describe("AgentRegistry", () => { expect(persisted?.title).toBe("Fix Login Bug"); }); + test("setTitle throws when the agent record does not exist", async () => { + await expect(registry.setTitle("missing-agent", "Impossible")) + .rejects.toThrow("Agent missing-agent not found"); + }); + + test("applySnapshot accepts explicit title overrides", async () => { + const agentId = "agent-override"; + await registry.applySnapshot( + createManagedAgent({ id: agentId }), + { title: "Provided Title" } + ); + + const record = await registry.get(agentId); + expect(record?.title).toBe("Provided Title"); + }); + test("applySnapshot preserves custom titles while updating metadata", async () => { const agentId = "agent-3"; await registry.applySnapshot( diff --git a/packages/server/src/server/agent/agent-registry.ts b/packages/server/src/server/agent/agent-registry.ts index 0ca4bd171..90e0e80c2 100644 --- a/packages/server/src/server/agent/agent-registry.ts +++ b/packages/server/src/server/agent/agent-registry.ts @@ -102,11 +102,16 @@ export class AgentRegistry { await this.flush(); } - async applySnapshot(agent: ManagedAgent): Promise { + async applySnapshot( + agent: ManagedAgent, + options?: { title?: string | null } + ): Promise { await this.load(); const existing = this.cache.get(agent.id); + const hasTitleOverride = + options !== undefined && Object.prototype.hasOwnProperty.call(options, "title"); const record = toStoredAgentRecord(agent, { - title: existing?.title ?? null, + title: hasTitleOverride ? options?.title ?? null : existing?.title ?? null, createdAt: existing?.createdAt, }); this.cache.set(agent.id, record); @@ -117,7 +122,7 @@ export class AgentRegistry { await this.load(); const record = this.cache.get(agentId); if (!record) { - return; + throw new Error(`Agent ${agentId} not found`); } this.cache.set(agentId, { ...record, title }); await this.flush(); diff --git a/packages/server/src/server/agent/agent-sdk-types.ts b/packages/server/src/server/agent/agent-sdk-types.ts index d54be5dcb..30f07124f 100644 --- a/packages/server/src/server/agent/agent-sdk-types.ts +++ b/packages/server/src/server/agent/agent-sdk-types.ts @@ -136,6 +136,7 @@ export type AgentSessionConfig = { cwd: string; modeId?: string; model?: string; + title?: string | null; approvalPolicy?: string; sandboxMode?: string; networkAccess?: boolean; diff --git a/packages/server/src/server/agent/mcp-server.test.ts b/packages/server/src/server/agent/mcp-server.test.ts index d59648a0c..db64adec8 100644 --- a/packages/server/src/server/agent/mcp-server.test.ts +++ b/packages/server/src/server/agent/mcp-server.test.ts @@ -71,7 +71,7 @@ describe("create_agent MCP tool", () => { expect(ok.success).toBe(true); }); - it("persists provided titles immediately after agent creation", async () => { + it("passes caller-provided titles directly into createAgent", async () => { const { agentManager, agentRegistry, spies } = createTestDeps(); spies.agentManager.createAgent.mockResolvedValue({ id: "agent-123", @@ -88,9 +88,11 @@ describe("create_agent MCP tool", () => { title: " Fix auth bug ", }); - expect(spies.agentRegistry.setTitle).toHaveBeenCalledWith( - "agent-123", - "Fix auth bug" + expect(spies.agentManager.createAgent).toHaveBeenCalledWith( + expect.objectContaining({ + cwd: "/tmp/repo", + title: "Fix auth bug", + }) ); }); }); diff --git a/packages/server/src/server/agent/mcp-server.ts b/packages/server/src/server/agent/mcp-server.ts index 7dfb5f51b..8e01b0c41 100644 --- a/packages/server/src/server/agent/mcp-server.ts +++ b/packages/server/src/server/agent/mcp-server.ts @@ -193,21 +193,14 @@ export async function createAgentMcpServer( } const provider: AgentProvider = agentType ?? "claude"; + const normalizedTitle = title?.trim() ?? null; const snapshot = await agentManager.createAgent({ provider, cwd: resolvedCwd, modeId: initialMode, + title: normalizedTitle ?? undefined, }); - try { - await agentRegistry.setTitle(snapshot.id, title); - } catch (error) { - console.error( - `[Agent MCP] Failed to persist title for ${snapshot.id}:`, - error - ); - } - if (initialPrompt) { try { agentManager.recordUserMessage(snapshot.id, initialPrompt); diff --git a/packages/server/src/server/index.ts b/packages/server/src/server/index.ts index 25dc98dcc..fb9466d46 100644 --- a/packages/server/src/server/index.ts +++ b/packages/server/src/server/index.ts @@ -103,6 +103,7 @@ async function main() { }), codex: new CodexAgentClient(), }, + registry: agentRegistry, }); attachAgentRegistryPersistence(agentManager, agentRegistry); diff --git a/packages/server/src/server/messages.ts b/packages/server/src/server/messages.ts index e0f576c5e..c39c10b13 100644 --- a/packages/server/src/server/messages.ts +++ b/packages/server/src/server/messages.ts @@ -58,6 +58,13 @@ const AgentSessionConfigSchema = z.object({ cwd: z.string(), modeId: z.string().optional(), model: z.string().optional(), + title: z + .string() + .trim() + .min(1) + .max(40) + .optional() + .nullable(), approvalPolicy: z.string().optional(), sandboxMode: z.string().optional(), networkAccess: z.boolean().optional(), diff --git a/packages/server/src/server/persistence-hooks.ts b/packages/server/src/server/persistence-hooks.ts index 4dc823de8..389764588 100644 --- a/packages/server/src/server/persistence-hooks.ts +++ b/packages/server/src/server/persistence-hooks.ts @@ -42,6 +42,7 @@ export function buildConfigOverrides( cwd: record.cwd, modeId: record.lastModeId ?? record.config?.modeId ?? undefined, model: record.config?.model ?? undefined, + title: record.title ?? undefined, extra: record.config?.extra ?? undefined, }; } @@ -58,6 +59,7 @@ export function buildSessionConfig( cwd: record.cwd, modeId: overrides.modeId, model: overrides.model, + title: overrides.title, extra: overrides.extra, }; }