diff --git a/packages/server/src/server/agent/create-agent/create.test.ts b/packages/server/src/server/agent/create-agent/create.test.ts index 89cf7bab6..e8547d27e 100644 --- a/packages/server/src/server/agent/create-agent/create.test.ts +++ b/packages/server/src/server/agent/create-agent/create.test.ts @@ -22,6 +22,22 @@ function createRealAgentManager(storage: AgentStorage): AgentManager { }); } +async function removeRealAgentManagerWorkdir({ + agentManager, + storage, + workdir, +}: { + agentManager: AgentManager; + storage: AgentStorage; + workdir: string; +}): Promise { + agentManager.prepareForShutdown(); + await Promise.all(agentManager.listAgents().map((agent) => agentManager.closeAgent(agent.id))); + await agentManager.flushForShutdown(); + await storage.flush(); + rmSync(workdir, { recursive: true, force: true }); +} + // Creates a worktree directory under repoRoot and reports it back as a fresh // workspace so the command can stamp the agent with it (mirrors the production // worktree service). @@ -238,7 +254,7 @@ test("session create stamps the requested workspaceId when no worktree setup run const stored = await storage.get(snapshot.id); expect(stored?.workspaceId).toBe("ws-source"); } finally { - rmSync(workdir, { recursive: true, force: true }); + await removeRealAgentManagerWorkdir({ agentManager, storage, workdir }); } }); @@ -273,7 +289,7 @@ test("session create stamps the new worktree's workspaceId when a setup continua const stored = await storage.get(snapshot.id); expect(stored?.workspaceId).toBe("ws-new-worktree"); } finally { - rmSync(workdir, { recursive: true, force: true }); + await removeRealAgentManagerWorkdir({ agentManager, storage, workdir }); } }); @@ -324,7 +340,7 @@ test("mcp create stamps the new worktree's workspaceId, not the parent's", async expect(storedChild?.workspaceId).toBe("ws-new-worktree"); expect(child.cwd).toBe(join(workdir, "worktree", "packages", "app")); } finally { - rmSync(workdir, { recursive: true, force: true }); + await removeRealAgentManagerWorkdir({ agentManager, storage, workdir }); } }); @@ -376,7 +392,7 @@ test("mcp create exposes the created worktree before dispatching the initial pro expect(observed).toEqual({ createdWorktree, lifecycle: "idle" }); } finally { - rmSync(workdir, { recursive: true, force: true }); + await removeRealAgentManagerWorkdir({ agentManager, storage, workdir }); } }); @@ -414,7 +430,7 @@ test("session create keeps the prompt title after the initial prompt settles", a const settled = await storage.get(snapshot.id); expect(settled?.title).toBe(title); } finally { - rmSync(workdir, { recursive: true, force: true }); + await removeRealAgentManagerWorkdir({ agentManager, storage, workdir }); } }); @@ -452,6 +468,6 @@ test("session create keeps an explicit title after the initial prompt settles", const settled = await storage.get(snapshot.id); expect(settled?.title).toBe(title); } finally { - rmSync(workdir, { recursive: true, force: true }); + await removeRealAgentManagerWorkdir({ agentManager, storage, workdir }); } }); diff --git a/packages/server/src/server/hub/test-utils/relationship-harness.ts b/packages/server/src/server/hub/test-utils/relationship-harness.ts index f18bc4215..a74ebe068 100644 --- a/packages/server/src/server/hub/test-utils/relationship-harness.ts +++ b/packages/server/src/server/hub/test-utils/relationship-harness.ts @@ -203,7 +203,10 @@ class ControlledAgentClient implements AgentClient { resumes = 0; createdConfigs: AgentSessionConfig[] = []; - constructor(private readonly client: AgentClient) { + constructor( + private readonly client: AgentClient, + private readonly internalClient: AgentClient, + ) { this.provider = client.provider; this.capabilities = client.capabilities; } @@ -230,6 +233,9 @@ class ControlledAgentClient implements AgentClient { launchContext?: AgentLaunchContext, options?: AgentCreateSessionOptions, ): Promise { + if (config.internal) { + return this.internalClient.createSession(config, launchContext, options); + } this.creations++; this.createdConfigs.push({ ...config }); this.creationObserved.resolve(); @@ -242,6 +248,9 @@ class ControlledAgentClient implements AgentClient { overrides?: Partial, launchContext?: AgentLaunchContext, ): Promise { + if (overrides?.internal) { + return this.internalClient.resumeSession(handle, overrides, launchContext); + } this.resumes++; return this.client.resumeSession(handle, overrides, launchContext); } @@ -438,6 +447,7 @@ export class HubRelationshipHarness { this.providerPrompts.push(prompt); }, }).codex, + createTestAgentClients({ supportsMcpServers: mcpEnabled }).codex, ); }