From 54f121f8432ebb2db539ddee2e3626dd6b6dcfe8 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Sat, 18 Apr 2026 17:41:35 +0100 Subject: [PATCH] test: speed up subagent runtime tests --- ...agents.sessions-spawn-depth-limits.test.ts | 4 +- ...s.subagents.sessions-spawn.test-harness.ts | 57 ++--- src/agents/pi-bundle-mcp-runtime.test.ts | 235 +++++++++--------- src/agents/pi-bundle-mcp-runtime.ts | 12 +- .../pi-bundle-mcp-tools.materialize.test.ts | 128 ++++------ .../subagent-registry.persistence.test.ts | 92 +++++-- .../subagent-spawn.model-session.test.ts | 7 +- 7 files changed, 290 insertions(+), 245 deletions(-) diff --git a/src/agents/openclaw-tools.subagents.sessions-spawn-depth-limits.test.ts b/src/agents/openclaw-tools.subagents.sessions-spawn-depth-limits.test.ts index 15158e26c02..566274abdcb 100644 --- a/src/agents/openclaw-tools.subagents.sessions-spawn-depth-limits.test.ts +++ b/src/agents/openclaw-tools.subagents.sessions-spawn-depth-limits.test.ts @@ -106,7 +106,7 @@ describe("sessions_spawn depth + child limits", () => { return { runId: "run-depth" }; } if (req.method === "agent.wait") { - return { status: "running" }; + return { status: "pending" }; } return {}; }); @@ -330,7 +330,7 @@ describe("sessions_spawn depth + child limits", () => { return { runId: "run-depth" }; } if (req.method === "agent.wait") { - return { status: "running" }; + return { status: "pending" }; } return {}; }); diff --git a/src/agents/openclaw-tools.subagents.sessions-spawn.test-harness.ts b/src/agents/openclaw-tools.subagents.sessions-spawn.test-harness.ts index 1001bfbcd05..2e82f3ac9f3 100644 --- a/src/agents/openclaw-tools.subagents.sessions-spawn.test-harness.ts +++ b/src/agents/openclaw-tools.subagents.sessions-spawn.test-harness.ts @@ -267,37 +267,32 @@ vi.mock("../../gateway/call.js", () => ({ callGateway: (opts: unknown) => hoisted.callGatewayMock(opts), })); -vi.mock("../config/config.js", async () => { - const actual = await vi.importActual("../config/config.js"); - return { - ...actual, - loadConfig: () => hoisted.state.configOverride, - resolveGatewayPort: () => 18789, - }; -}); +vi.mock("../config/config.js", () => ({ + loadConfig: () => hoisted.state.configOverride, + resolveGatewayPort: () => 18789, +})); -vi.mock("../config/sessions.js", async () => { - const actual = - await vi.importActual("../config/sessions.js"); - return { - ...actual, - loadSessionStore: () => hoisted.sessionStore, - resolveStorePath: () => "/tmp/openclaw-sessions-spawn-test-store.json", - updateSessionStore: async ( - _storePath: string, - mutator: (store: typeof hoisted.sessionStore) => void | Promise, - ) => { - await mutator(hoisted.sessionStore); - }, - }; -}); +vi.mock("../config/sessions.js", () => ({ + loadSessionStore: () => hoisted.sessionStore, + mergeSessionEntry: (existing: object | undefined, patch: object) => ({ + ...existing, + ...patch, + }), + resolveAgentMainSessionKey: (params: { + cfg?: { session?: { mainKey?: string } }; + agentId: string; + }) => `agent:${params.agentId}:${params.cfg?.session?.mainKey ?? "main"}`, + resolveStorePath: () => "/tmp/openclaw-sessions-spawn-test-store.json", + updateSessionStore: async ( + _storePath: string, + mutator: (store: typeof hoisted.sessionStore) => void | Promise, + ) => { + await mutator(hoisted.sessionStore); + }, +})); // Same module, different specifier (used by tools under src/agents/tools/*). -vi.mock("../../config/config.js", async () => { - const actual = await vi.importActual("../config/config.js"); - return { - ...actual, - loadConfig: () => hoisted.state.configOverride, - resolveGatewayPort: () => 18789, - }; -}); +vi.mock("../../config/config.js", () => ({ + loadConfig: () => hoisted.state.configOverride, + resolveGatewayPort: () => 18789, +})); diff --git a/src/agents/pi-bundle-mcp-runtime.test.ts b/src/agents/pi-bundle-mcp-runtime.test.ts index 2f8dd091593..8e33cce68df 100644 --- a/src/agents/pi-bundle-mcp-runtime.test.ts +++ b/src/agents/pi-bundle-mcp-runtime.test.ts @@ -16,54 +16,60 @@ import { } from "./pi-bundle-mcp-tools.js"; import type { SessionMcpRuntime } from "./pi-bundle-mcp-types.js"; +type RuntimeFactoryOptions = NonNullable< + Parameters[0] +>; +type RuntimeFactory = NonNullable; + +function makeRuntime( + tools: Array<{ toolName: string; description: string }>, + serverName = "bundleProbe", +): SessionMcpRuntime { + return { + sessionId: "session-colliding-tools", + workspaceDir: "/tmp", + configFingerprint: "fingerprint", + createdAt: 0, + lastUsedAt: 0, + markUsed: () => {}, + getCatalog: async () => ({ + version: 1, + generatedAt: 0, + servers: { + [serverName]: { + serverName, + launchSummary: serverName, + toolCount: tools.length, + }, + }, + tools: tools.map((tool) => ({ + serverName, + safeServerName: serverName, + toolName: tool.toolName, + description: tool.description, + inputSchema: { + type: "object", + properties: { + toolName: { type: "string", const: tool.toolName }, + }, + }, + fallbackDescription: tool.description, + })), + }), + callTool: async (_serverName, toolName) => ({ + content: [{ type: "text", text: toolName }], + isError: false, + }), + dispose: async () => {}, + }; +} + afterEach(async () => { await cleanupBundleMcpHarness(); }); describe("session MCP runtime", () => { it("keeps colliding sanitized tool definitions stable across catalog order changes", async () => { - function makeRuntime( - tools: Array<{ toolName: string; description: string }>, - ): SessionMcpRuntime { - return { - sessionId: "session-colliding-tools", - workspaceDir: "/tmp", - configFingerprint: "fingerprint", - createdAt: 0, - lastUsedAt: 0, - markUsed: () => {}, - getCatalog: async () => ({ - version: 1, - generatedAt: 0, - servers: { - collision: { - serverName: "collision", - launchSummary: "collision", - toolCount: tools.length, - }, - }, - tools: tools.map((tool) => ({ - serverName: "collision", - safeServerName: "collision", - toolName: tool.toolName, - description: tool.description, - inputSchema: { - type: "object", - properties: { - toolName: { type: "string", const: tool.toolName }, - }, - }, - fallbackDescription: tool.description, - })), - }), - callTool: async (_serverName, toolName) => ({ - content: [{ type: "text", text: toolName }], - isError: false, - }), - dispose: async () => {}, - }; - } - const catalogA = [ { toolName: "alpha?", description: "question" }, { toolName: "alpha!", description: "bang" }, @@ -71,10 +77,10 @@ describe("session MCP runtime", () => { const catalogB = catalogA.toReversed(); const materializedA = await materializeBundleMcpToolsForRun({ - runtime: makeRuntime(catalogA), + runtime: makeRuntime(catalogA, "collision"), }); const materializedB = await materializeBundleMcpToolsForRun({ - runtime: makeRuntime(catalogB), + runtime: makeRuntime(catalogB, "collision"), }); const summarizeTools = (runtime: Awaited>) => @@ -110,31 +116,29 @@ describe("session MCP runtime", () => { }); it("reuses repeated materialization and recreates after explicit disposal", async () => { - const workspaceDir = await makeTempDir("openclaw-bundle-mcp-tools-"); - const startupCounterPath = path.join(workspaceDir, "bundle-starts.txt"); - const pluginRoot = path.join(workspaceDir, ".openclaw", "extensions", "bundle-probe"); - const serverScriptPath = path.join(pluginRoot, "servers", "bundle-probe.mjs"); - await writeBundleProbeMcpServer(serverScriptPath, { startupCounterPath }); - await writeClaudeBundle({ pluginRoot, serverScriptPath }); - const cfg = { - plugins: { - entries: { - "bundle-probe": { enabled: true }, - }, - }, + const created: SessionMcpRuntime[] = []; + const createRuntime: RuntimeFactory = (params) => { + const runtime = makeRuntime([{ toolName: "bundle_probe", description: "Bundle MCP probe" }]); + created.push(runtime); + return { + ...runtime, + sessionId: params.sessionId, + sessionKey: params.sessionKey, + workspaceDir: params.workspaceDir, + configFingerprint: params.configFingerprint ?? "fingerprint", + }; }; + const manager = __testing.createSessionMcpRuntimeManager({ createRuntime }); - const runtimeA = await getOrCreateSessionMcpRuntime({ + const runtimeA = await manager.getOrCreate({ sessionId: "session-a", sessionKey: "agent:test:session-a", - workspaceDir, - cfg, + workspaceDir: "/workspace", }); - const runtimeB = await getOrCreateSessionMcpRuntime({ + const runtimeB = await manager.getOrCreate({ sessionId: "session-a", sessionKey: "agent:test:session-a", - workspaceDir, - cfg, + workspaceDir: "/workspace", }); const materializedA = await materializeBundleMcpToolsForRun({ runtime: runtimeA }); @@ -146,39 +150,51 @@ describe("session MCP runtime", () => { expect(runtimeA).toBe(runtimeB); expect(materializedA.tools.map((tool) => tool.name)).toEqual(["bundleProbe__bundle_probe"]); expect(materializedB.tools.map((tool) => tool.name)).toEqual(["bundleProbe__bundle_probe"]); - expect(await fs.readFile(startupCounterPath, "utf8")).toBe("1"); - expect(__testing.getCachedSessionIds()).toEqual(["session-a"]); + expect(created).toHaveLength(1); + expect(manager.listSessionIds()).toEqual(["session-a"]); - await disposeSessionMcpRuntime("session-a"); + await manager.disposeSession("session-a"); - const runtimeC = await getOrCreateSessionMcpRuntime({ + const runtimeC = await manager.getOrCreate({ sessionId: "session-a", sessionKey: "agent:test:session-a", - workspaceDir, - cfg, + workspaceDir: "/workspace", }); await materializeBundleMcpToolsForRun({ runtime: runtimeC }); expect(runtimeC).not.toBe(runtimeA); - expect(await fs.readFile(startupCounterPath, "utf8")).toBe("2"); + expect(created).toHaveLength(2); }); it("recreates the session runtime when MCP config changes", async () => { - const workspaceDir = await makeTempDir("openclaw-bundle-mcp-tools-"); - const startupCounterPath = path.join(workspaceDir, "bundle-starts.txt"); - const serverScriptPath = path.join(workspaceDir, "servers", "configured-probe.mjs"); - await writeBundleProbeMcpServer(serverScriptPath, { startupCounterPath }); + const createRuntime: RuntimeFactory = (params) => { + const probeText = String( + params.cfg?.mcp?.servers?.configuredProbe?.env?.BUNDLE_PROBE_TEXT ?? "FROM-CONFIG", + ); + return { + ...makeRuntime([{ toolName: "bundle_probe", description: "Bundle MCP probe" }]), + sessionId: params.sessionId, + sessionKey: params.sessionKey, + workspaceDir: params.workspaceDir, + configFingerprint: params.configFingerprint ?? "fingerprint", + callTool: async () => ({ + content: [{ type: "text", text: probeText }], + isError: false, + }), + }; + }; + const manager = __testing.createSessionMcpRuntimeManager({ createRuntime }); - const runtimeA = await getOrCreateSessionMcpRuntime({ + const runtimeA = await manager.getOrCreate({ sessionId: "session-c", sessionKey: "agent:test:session-c", - workspaceDir, + workspaceDir: "/workspace", cfg: { mcp: { servers: { configuredProbe: { command: "node", - args: [serverScriptPath], + args: ["server-a.mjs"], env: { BUNDLE_PROBE_TEXT: "FROM-CONFIG-A", }, @@ -195,16 +211,16 @@ describe("session MCP runtime", () => { undefined, ); - const runtimeB = await getOrCreateSessionMcpRuntime({ + const runtimeB = await manager.getOrCreate({ sessionId: "session-c", sessionKey: "agent:test:session-c", - workspaceDir, + workspaceDir: "/workspace", cfg: { mcp: { servers: { configuredProbe: { command: "node", - args: [serverScriptPath], + args: ["server-b.mjs"], env: { BUNDLE_PROBE_TEXT: "FROM-CONFIG-B", }, @@ -224,7 +240,6 @@ describe("session MCP runtime", () => { expect(runtimeA).not.toBe(runtimeB); expect(resultA.content[0]).toMatchObject({ type: "text", text: "FROM-CONFIG-A" }); expect(resultB.content[0]).toMatchObject({ type: "text", text: "FROM-CONFIG-B" }); - expect(await fs.readFile(startupCounterPath, "utf8")).toBe("2"); }); it("disposes startup-in-flight runtimes without leaking MCP processes", async () => { @@ -275,61 +290,51 @@ describe("session MCP runtime", () => { }); it("materialized disposal can retire a manager-owned runtime", async () => { - const workspaceDir = await makeTempDir("openclaw-bundle-mcp-tools-"); - const startupCounterPath = path.join(workspaceDir, "bundle-starts.txt"); - const pidPath = path.join(workspaceDir, "bundle.pid"); - const exitMarkerPath = path.join(workspaceDir, "bundle.exit"); - const pluginRoot = path.join(workspaceDir, ".openclaw", "extensions", "bundle-probe"); - const serverScriptPath = path.join(pluginRoot, "servers", "bundle-probe.mjs"); - await writeBundleProbeMcpServer(serverScriptPath, { - startupCounterPath, - pidPath, - exitMarkerPath, - }); - await writeClaudeBundle({ pluginRoot, serverScriptPath }); + const disposed: string[] = []; + const created: SessionMcpRuntime[] = []; + const createRuntime: RuntimeFactory = (params) => { + const runtime = { + ...makeRuntime([{ toolName: "bundle_probe", description: "Bundle MCP probe" }]), + sessionId: params.sessionId, + sessionKey: params.sessionKey, + workspaceDir: params.workspaceDir, + configFingerprint: params.configFingerprint ?? "fingerprint", + dispose: async () => { + disposed.push(params.sessionId); + }, + }; + created.push(runtime); + return runtime; + }; + const manager = __testing.createSessionMcpRuntimeManager({ createRuntime }); - const runtimeA = await getOrCreateSessionMcpRuntime({ + const runtimeA = await manager.getOrCreate({ sessionId: "session-e", sessionKey: "agent:test:session-e", - workspaceDir, - cfg: { - plugins: { - entries: { - "bundle-probe": { enabled: true }, - }, - }, - }, + workspaceDir: "/workspace", }); const materialized = await materializeBundleMcpToolsForRun({ runtime: runtimeA, disposeRuntime: async () => { - await disposeSessionMcpRuntime("session-e"); + await manager.disposeSession("session-e"); }, }); expect(materialized.tools.map((tool) => tool.name)).toEqual(["bundleProbe__bundle_probe"]); - expect(await waitForFileText(pidPath)).toMatch(/^\d+$/); await materialized.dispose(); - expect(await waitForFileText(exitMarkerPath)).toBe("exited"); - expect(__testing.getCachedSessionIds()).not.toContain("session-e"); + expect(disposed).toEqual(["session-e"]); + expect(manager.listSessionIds()).not.toContain("session-e"); - const runtimeB = await getOrCreateSessionMcpRuntime({ + const runtimeB = await manager.getOrCreate({ sessionId: "session-e", sessionKey: "agent:test:session-e", - workspaceDir, - cfg: { - plugins: { - entries: { - "bundle-probe": { enabled: true }, - }, - }, - }, + workspaceDir: "/workspace", }); expect(runtimeB).not.toBe(runtimeA); await materializeBundleMcpToolsForRun({ runtime: runtimeB }); - expect(await fs.readFile(startupCounterPath, "utf8")).toBe("2"); + expect(created).toHaveLength(2); }); }); diff --git a/src/agents/pi-bundle-mcp-runtime.ts b/src/agents/pi-bundle-mcp-runtime.ts index dcf116f7e45..7b488130f2a 100644 --- a/src/agents/pi-bundle-mcp-runtime.ts +++ b/src/agents/pi-bundle-mcp-runtime.ts @@ -30,6 +30,9 @@ type BundleMcpSession = { type LoadedMcpConfig = ReturnType; type ListedTool = Awaited>["tools"][number]; +type CreateSessionMcpRuntime = ( + params: Parameters[0] & { configFingerprint?: string }, +) => SessionMcpRuntime; const SESSION_MCP_RUNTIME_MANAGER_KEY = Symbol.for("openclaw.sessionMcpRuntimeManager"); @@ -289,9 +292,12 @@ export function createSessionMcpRuntime(params: { }; } -function createSessionMcpRuntimeManager(): SessionMcpRuntimeManager { +function createSessionMcpRuntimeManager( + opts: { createRuntime?: CreateSessionMcpRuntime } = {}, +): SessionMcpRuntimeManager { const runtimesBySessionId = new Map(); const sessionIdBySessionKey = new Map(); + const createRuntime = opts.createRuntime ?? createSessionMcpRuntime; const createInFlight = new Map< string, { @@ -338,11 +344,12 @@ function createSessionMcpRuntimeManager(): SessionMcpRuntimeManager { await staleRuntime?.dispose(); } const created = Promise.resolve( - createSessionMcpRuntime({ + createRuntime({ sessionId: params.sessionId, sessionKey: params.sessionKey, workspaceDir: params.workspaceDir, cfg: params.cfg, + configFingerprint: nextFingerprint, }), ).then((runtime) => { runtime.markUsed(); @@ -434,6 +441,7 @@ export async function disposeAllSessionMcpRuntimes(): Promise { } export const __testing = { + createSessionMcpRuntimeManager, async resetSessionMcpRuntimeManager() { await disposeAllSessionMcpRuntimes(); }, diff --git a/src/agents/pi-bundle-mcp-tools.materialize.test.ts b/src/agents/pi-bundle-mcp-tools.materialize.test.ts index ec018044a71..e40d9b35ca5 100644 --- a/src/agents/pi-bundle-mcp-tools.materialize.test.ts +++ b/src/agents/pi-bundle-mcp-tools.materialize.test.ts @@ -1,4 +1,3 @@ -import { createRequire } from "node:module"; import path from "node:path"; import { afterEach, describe, expect, it, vi } from "vitest"; import { @@ -6,44 +5,19 @@ import { makeTempDir, startSseProbeServer, writeBundleProbeMcpServer, - writeClaudeBundle, - writeExecutable, } from "./pi-bundle-mcp-test-harness.js"; import { createBundleMcpToolRuntime, materializeBundleMcpToolsForRun, } from "./pi-bundle-mcp-tools.js"; +import type { McpCatalogTool } from "./pi-bundle-mcp-types.js"; import type { SessionMcpRuntime } from "./pi-bundle-mcp-types.js"; -const require = createRequire(import.meta.url); -const SDK_SERVER_MCP_PATH = require.resolve("@modelcontextprotocol/sdk/server/mcp.js"); -const SDK_SERVER_STDIO_PATH = require.resolve("@modelcontextprotocol/sdk/server/stdio.js"); - afterEach(async () => { await cleanupBundleMcpHarness(); }); -async function createBundleProbeRuntime(params?: { reservedToolNames?: string[] }) { - const workspaceDir = await makeTempDir("openclaw-bundle-mcp-tools-"); - const pluginRoot = path.join(workspaceDir, ".openclaw", "extensions", "bundle-probe"); - const serverScriptPath = path.join(pluginRoot, "servers", "bundle-probe.mjs"); - await writeBundleProbeMcpServer(serverScriptPath); - await writeClaudeBundle({ pluginRoot, serverScriptPath }); - - return await createBundleMcpToolRuntime({ - workspaceDir, - cfg: { - plugins: { - entries: { - "bundle-probe": { enabled: true }, - }, - }, - }, - reservedToolNames: params?.reservedToolNames, - }); -} - -function makeSingleToolRuntime(): SessionMcpRuntime { +function makeToolRuntime(tools?: McpCatalogTool[]): SessionMcpRuntime { return { sessionId: "session-collision", workspaceDir: "/tmp", @@ -61,7 +35,7 @@ function makeSingleToolRuntime(): SessionMcpRuntime { toolCount: 1, }, }, - tools: [ + tools: tools ?? [ { serverName: "bundleProbe", safeServerName: "bundleProbe", @@ -81,28 +55,26 @@ function makeSingleToolRuntime(): SessionMcpRuntime { } describe("createBundleMcpToolRuntime", () => { - it("loads bundle MCP tools and executes them", async () => { - const runtime = await createBundleProbeRuntime(); + it("materializes bundle MCP tools and executes them", async () => { + const runtime = await materializeBundleMcpToolsForRun({ + runtime: makeToolRuntime(), + }); - try { - expect(runtime.tools.map((tool) => tool.name)).toEqual(["bundleProbe__bundle_probe"]); - const result = await runtime.tools[0].execute("call-bundle-probe", {}, undefined, undefined); - expect(result.content[0]).toMatchObject({ - type: "text", - text: "FROM-BUNDLE", - }); - expect(result.details).toEqual({ - mcpServer: "bundleProbe", - mcpTool: "bundle_probe", - }); - } finally { - await runtime.dispose(); - } + expect(runtime.tools.map((tool) => tool.name)).toEqual(["bundleProbe__bundle_probe"]); + const result = await runtime.tools[0].execute("call-bundle-probe", {}, undefined, undefined); + expect(result.content[0]).toMatchObject({ + type: "text", + text: "FROM-BUNDLE", + }); + expect(result.details).toEqual({ + mcpServer: "bundleProbe", + mcpTool: "bundle_probe", + }); }); it("disambiguates bundle MCP tools that collide with existing tool names", async () => { const runtime = await materializeBundleMcpToolsForRun({ - runtime: makeSingleToolRuntime(), + runtime: makeToolRuntime(), reservedToolNames: ["bundleProbe__bundle_probe"], }); @@ -192,41 +164,39 @@ describe("createBundleMcpToolRuntime", () => { }); it("returns tools sorted alphabetically for stable prompt-cache keys", async () => { - const workspaceDir = await makeTempDir("openclaw-bundle-mcp-tools-"); - const serverScriptPath = path.join(workspaceDir, "servers", "multi-tool.mjs"); - // Register tools in non-alphabetical order; runtime must sort them. - await writeExecutable( - serverScriptPath, - `#!/usr/bin/env node -import { McpServer } from ${JSON.stringify(SDK_SERVER_MCP_PATH)}; -import { StdioServerTransport } from ${JSON.stringify(SDK_SERVER_STDIO_PATH)}; -const server = new McpServer({ name: "multi", version: "1.0.0" }); -server.tool("zeta", "z", async () => ({ content: [{ type: "text", text: "z" }] })); -server.tool("alpha", "a", async () => ({ content: [{ type: "text", text: "a" }] })); -server.tool("mu", "m", async () => ({ content: [{ type: "text", text: "m" }] })); -await server.connect(new StdioServerTransport()); -`, - ); - - const runtime = await createBundleMcpToolRuntime({ - workspaceDir, - cfg: { - mcp: { - servers: { - multi: { command: "node", args: [serverScriptPath] }, - }, + const runtime = await materializeBundleMcpToolsForRun({ + runtime: makeToolRuntime([ + { + serverName: "multi", + safeServerName: "multi", + toolName: "zeta", + description: "z", + inputSchema: { type: "object", properties: {} }, + fallbackDescription: "z", }, - }, + { + serverName: "multi", + safeServerName: "multi", + toolName: "alpha", + description: "a", + inputSchema: { type: "object", properties: {} }, + fallbackDescription: "a", + }, + { + serverName: "multi", + safeServerName: "multi", + toolName: "mu", + description: "m", + inputSchema: { type: "object", properties: {} }, + fallbackDescription: "m", + }, + ]), }); - try { - expect(runtime.tools.map((tool) => tool.name)).toEqual([ - "multi__alpha", - "multi__mu", - "multi__zeta", - ]); - } finally { - await runtime.dispose(); - } + expect(runtime.tools.map((tool) => tool.name)).toEqual([ + "multi__alpha", + "multi__mu", + "multi__zeta", + ]); }); }); diff --git a/src/agents/subagent-registry.persistence.test.ts b/src/agents/subagent-registry.persistence.test.ts index d98ac955f25..b2768760c45 100644 --- a/src/agents/subagent-registry.persistence.test.ts +++ b/src/agents/subagent-registry.persistence.test.ts @@ -157,13 +157,18 @@ describe("subagent registry persistence", () => { const flushQueuedRegistryWork = async () => { await Promise.resolve(); await Promise.resolve(); - await new Promise((resolve) => setTimeout(resolve, 25)); }; - const restartRegistryAndFlush = async () => { + const waitForRegistryWork = async (predicate: () => boolean | Promise) => { + await vi.waitFor(async () => expect(await predicate()).toBe(true), { + interval: 1, + timeout: 1_000, + }); + }; + + const restartRegistry = () => { resetSubagentRegistryForTests({ persist: false }); initSubagentRegistry(); - await flushQueuedRegistryWork(); }; beforeEach(() => { @@ -260,9 +265,7 @@ describe("subagent registry persistence", () => { sessionId: "sess-two", }); - resetSubagentRegistryForTests({ persist: false }); - initSubagentRegistry(); - + restartRegistry(); await flushQueuedRegistryWork(); // announce should NOT be called since cleanupHandled was true @@ -385,7 +388,18 @@ describe("subagent registry persistence", () => { const registryPath = await writePersistedRegistry(persisted); announceSpy.mockResolvedValueOnce(false); - await restartRegistryAndFlush(); + restartRegistry(); + await waitForRegistryWork(async () => { + const afterFirst = await readPersistedRun<{ + cleanupHandled?: boolean; + cleanupCompletedAt?: number; + }>(registryPath, "run-3"); + return ( + announceSpy.mock.calls.length === 1 && + afterFirst?.cleanupHandled === false && + afterFirst.cleanupCompletedAt === undefined + ); + }); expect(announceSpy).toHaveBeenCalledTimes(1); const afterFirst = await readPersistedRun<{ @@ -396,7 +410,13 @@ describe("subagent registry persistence", () => { expect(afterFirst?.cleanupCompletedAt).toBeUndefined(); announceSpy.mockResolvedValueOnce(true); - await restartRegistryAndFlush(); + restartRegistry(); + await waitForRegistryWork(async () => { + const afterSecond = await readPersistedRun<{ + cleanupCompletedAt?: number; + }>(registryPath, "run-3"); + return announceSpy.mock.calls.length === 2 && afterSecond?.cleanupCompletedAt != null; + }); expect(announceSpy).toHaveBeenCalledTimes(2); const afterSecond = JSON.parse(await fs.readFile(registryPath, "utf8")) as { @@ -415,7 +435,18 @@ describe("subagent registry persistence", () => { const registryPath = await writePersistedRegistry(persisted); announceSpy.mockRejectedValueOnce(new Error("announce boom")); - await restartRegistryAndFlush(); + restartRegistry(); + await waitForRegistryWork(async () => { + const afterFirst = await readPersistedRun<{ + cleanupHandled?: boolean; + cleanupCompletedAt?: number; + }>(registryPath, "run-reject"); + return ( + announceSpy.mock.calls.length === 1 && + afterFirst?.cleanupHandled === false && + afterFirst.cleanupCompletedAt === undefined + ); + }); expect(announceSpy).toHaveBeenCalledTimes(1); const afterFirst = JSON.parse(await fs.readFile(registryPath, "utf8")) as { @@ -425,7 +456,13 @@ describe("subagent registry persistence", () => { expect(afterFirst.runs["run-reject"].cleanupCompletedAt).toBeUndefined(); announceSpy.mockResolvedValueOnce(true); - await restartRegistryAndFlush(); + restartRegistry(); + await waitForRegistryWork(async () => { + const afterSecond = await readPersistedRun<{ + cleanupCompletedAt?: number; + }>(registryPath, "run-reject"); + return announceSpy.mock.calls.length === 2 && afterSecond?.cleanupCompletedAt != null; + }); expect(announceSpy).toHaveBeenCalledTimes(2); const afterSecond = JSON.parse(await fs.readFile(registryPath, "utf8")) as { @@ -444,14 +481,27 @@ describe("subagent registry persistence", () => { const registryPath = await writePersistedRegistry(persisted); announceSpy.mockResolvedValueOnce(false); - await restartRegistryAndFlush(); + restartRegistry(); + await waitForRegistryWork(async () => { + const afterFirst = await readPersistedRun<{ cleanupHandled?: boolean }>( + registryPath, + "run-4", + ); + return announceSpy.mock.calls.length === 1 && afterFirst?.cleanupHandled === false; + }); expect(announceSpy).toHaveBeenCalledTimes(1); const afterFirst = await readPersistedRun<{ cleanupHandled?: boolean }>(registryPath, "run-4"); expect(afterFirst?.cleanupHandled).toBe(false); announceSpy.mockResolvedValueOnce(true); - await restartRegistryAndFlush(); + restartRegistry(); + await waitForRegistryWork(async () => { + const afterSecond = JSON.parse(await fs.readFile(registryPath, "utf8")) as { + runs?: Record; + }; + return announceSpy.mock.calls.length === 2 && afterSecond.runs?.["run-4"] === undefined; + }); expect(announceSpy).toHaveBeenCalledTimes(2); const afterSecond = JSON.parse(await fs.readFile(registryPath, "utf8")) as { @@ -471,7 +521,13 @@ describe("subagent registry persistence", () => { seedChildSessions: false, }); - await restartRegistryAndFlush(); + restartRegistry(); + await waitForRegistryWork(async () => { + const after = JSON.parse(await fs.readFile(registryPath, "utf8")) as { + runs?: Record; + }; + return after.runs?.["run-orphan-restore"] === undefined; + }); expect(announceSpy).not.toHaveBeenCalled(); const after = JSON.parse(await fs.readFile(registryPath, "utf8")) as { @@ -511,7 +567,15 @@ describe("subagent registry persistence", () => { }; await fs.writeFile(registryPath, `${JSON.stringify(parsed)}\n`, "utf8"); - await restartRegistryAndFlush(); + restartRegistry(); + await waitForRegistryWork(async () => { + try { + await fs.access(attachmentsDir); + return false; + } catch (err) { + return (err as NodeJS.ErrnoException).code === "ENOENT"; + } + }); await expect(fs.access(attachmentsDir)).rejects.toMatchObject({ code: "ENOENT" }); const after = JSON.parse(await fs.readFile(registryPath, "utf8")) as { diff --git a/src/agents/subagent-spawn.model-session.test.ts b/src/agents/subagent-spawn.model-session.test.ts index 2518bcedc8d..06d0b987ca0 100644 --- a/src/agents/subagent-spawn.model-session.test.ts +++ b/src/agents/subagent-spawn.model-session.test.ts @@ -1,5 +1,5 @@ import os from "node:os"; -import { beforeEach, describe, expect, it, vi } from "vitest"; +import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import { createSubagentSpawnTestConfig, expectPersistedRuntimeModel, @@ -16,7 +16,7 @@ let resetSubagentRegistryForTests: typeof import("./subagent-registry.js").reset let spawnSubagentDirect: typeof import("./subagent-spawn.js").spawnSubagentDirect; describe("spawnSubagentDirect runtime model persistence", () => { - beforeEach(async () => { + beforeAll(async () => { ({ resetSubagentRegistryForTests, spawnSubagentDirect } = await loadSubagentSpawnModuleForTest({ callGatewayMock, loadConfig: () => createSubagentSpawnTestConfig(os.tmpdir()), @@ -24,6 +24,9 @@ describe("spawnSubagentDirect runtime model persistence", () => { pruneLegacyStoreKeysMock, workspaceDir: os.tmpdir(), })); + }); + + beforeEach(() => { resetSubagentRegistryForTests(); callGatewayMock.mockReset(); updateSessionStoreMock.mockReset();