test: speed up subagent runtime tests

2026-05-06 07:40:44 +00:00 · 2026-04-18 17:41:35 +01:00
parent fa2f53993a
commit 54f121f843
7 changed files with 290 additions and 245 deletions
--- 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 {};
    });
--- 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<typeof import("../config/config.js")>("../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<typeof import("../config/sessions.js")>("../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<void>,
-    ) => {
-      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<void>,
+  ) => {
+    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<typeof import("../config/config.js")>("../config/config.js");
-  return {
-    ...actual,
-    loadConfig: () => hoisted.state.configOverride,
-    resolveGatewayPort: () => 18789,
-  };
-});
+vi.mock("../../config/config.js", () => ({
+  loadConfig: () => hoisted.state.configOverride,
+  resolveGatewayPort: () => 18789,
+}));
--- 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<typeof __testing.createSessionMcpRuntimeManager>[0]
+>;
+type RuntimeFactory = NonNullable<RuntimeFactoryOptions["createRuntime"]>;
+
+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<ReturnType<typeof materializeBundleMcpToolsForRun>>) =>
@@ -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);
  });
 });
--- 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<typeof loadEmbeddedPiMcpConfig>;
 type ListedTool = Awaited<ReturnType<Client["listTools"]>>["tools"][number];
+type CreateSessionMcpRuntime = (
+  params: Parameters<typeof createSessionMcpRuntime>[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<string, SessionMcpRuntime>();
  const sessionIdBySessionKey = new Map<string, string>();
+  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<void> {
 }

 export const __testing = {
+  createSessionMcpRuntimeManager,
  async resetSessionMcpRuntimeManager() {
    await disposeAllSessionMcpRuntimes();
  },
--- 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",
+    ]);
  });
 });
--- 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<boolean>) => {
+    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<string, unknown>;
+      };
+      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<string, unknown>;
+      };
+      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 {
--- 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();