From 564f820efa2bb81ae49b8f55a93e41492ad31917 Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Fri, 24 Apr 2026 07:26:08 +0100
Subject: [PATCH] perf: reuse subagent spawn test module

---
 src/agents/subagent-spawn.context.test.ts | 84 ++++++++++++-----------
 1 file changed, 45 insertions(+), 39 deletions(-)

diff --git a/src/agents/subagent-spawn.context.test.ts b/src/agents/subagent-spawn.context.test.ts
index 95d65f16e1f..c9072fc1d50 100644
--- a/src/agents/subagent-spawn.context.test.ts
+++ b/src/agents/subagent-spawn.context.test.ts
@@ -1,5 +1,5 @@
 import path from "node:path";
-import { beforeEach, describe, expect, it, vi } from "vitest";
+import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
 import {
   loadSubagentSpawnModuleForTest,
   setupAcceptedSubagentGatewayMock,
@@ -8,22 +8,45 @@ import {
 type SessionStore = Record<string, Record<string, unknown>>;
 type GatewayRequest = { method?: string; params?: Record<string, unknown> };
 
-function createPersistentStoreMock(store: SessionStore) {
-  return vi.fn(async (_storePath: unknown, mutator: unknown) => {
-    if (typeof mutator !== "function") {
-      throw new Error("missing session store mutator");
-    }
-    return await mutator(store);
-  });
-}
-
 describe("sessions_spawn context modes", () => {
-  beforeEach(() => {
-    vi.resetModules();
+  const storePath = "/tmp/subagent-context-session-store.json";
+  const callGatewayMock = vi.fn();
+  const updateSessionStoreMock = vi.fn();
+  const forkSessionFromParentMock = vi.fn();
+  const resolveContextEngineMock = vi.fn();
+  let spawnSubagentDirect: Awaited<
+    ReturnType<typeof loadSubagentSpawnModuleForTest>
+  >["spawnSubagentDirect"];
+
+  beforeAll(async () => {
+    ({ spawnSubagentDirect } = await loadSubagentSpawnModuleForTest({
+      callGatewayMock,
+      updateSessionStoreMock,
+      forkSessionFromParentMock,
+      resolveContextEngineMock,
+      sessionStorePath: storePath,
+    }));
   });
 
+  beforeEach(() => {
+    callGatewayMock.mockReset();
+    updateSessionStoreMock.mockReset();
+    forkSessionFromParentMock.mockReset();
+    resolveContextEngineMock.mockReset();
+    setupAcceptedSubagentGatewayMock(callGatewayMock);
+    resolveContextEngineMock.mockResolvedValue({});
+  });
+
+  function usePersistentStoreMock(store: SessionStore) {
+    updateSessionStoreMock.mockImplementation(async (_storePath: unknown, mutator: unknown) => {
+      if (typeof mutator !== "function") {
+        throw new Error("missing session store mutator");
+      }
+      return await mutator(store);
+    });
+  }
+
   it("forks the requester transcript when context=fork", async () => {
-    const storePath = "/tmp/subagent-context-session-store.json";
     const store: SessionStore = {
       main: {
         sessionId: "parent-session-id",
@@ -32,20 +55,13 @@ describe("sessions_spawn context modes", () => {
         totalTokens: 1200,
       },
     };
-    const callGatewayMock = vi.fn();
-    setupAcceptedSubagentGatewayMock(callGatewayMock);
-    const forkSessionFromParentMock = vi.fn(async () => ({
+    usePersistentStoreMock(store);
+    forkSessionFromParentMock.mockImplementation(async () => ({
       sessionId: "forked-session-id",
       sessionFile: "/tmp/forked-session.jsonl",
     }));
     const prepareSubagentSpawn = vi.fn(async () => undefined);
-    const { spawnSubagentDirect } = await loadSubagentSpawnModuleForTest({
-      callGatewayMock,
-      updateSessionStoreMock: createPersistentStoreMock(store),
-      forkSessionFromParentMock,
-      resolveContextEngineMock: vi.fn(async () => ({ prepareSubagentSpawn })),
-      sessionStorePath: storePath,
-    });
+    resolveContextEngineMock.mockResolvedValue({ prepareSubagentSpawn });
 
     const result = await spawnSubagentDirect(
       { task: "inspect the current thread", context: "fork" },
@@ -79,16 +95,9 @@ describe("sessions_spawn context modes", () => {
     const store: SessionStore = {
       main: { sessionId: "parent-session-id", updatedAt: 1 },
     };
-    const callGatewayMock = vi.fn();
-    setupAcceptedSubagentGatewayMock(callGatewayMock);
-    const forkSessionFromParentMock = vi.fn();
+    usePersistentStoreMock(store);
     const prepareSubagentSpawn = vi.fn(async () => undefined);
-    const { spawnSubagentDirect } = await loadSubagentSpawnModuleForTest({
-      callGatewayMock,
-      updateSessionStoreMock: createPersistentStoreMock(store),
-      forkSessionFromParentMock,
-      resolveContextEngineMock: vi.fn(async () => ({ prepareSubagentSpawn })),
-    });
+    resolveContextEngineMock.mockResolvedValue({ prepareSubagentSpawn });
 
     const result = await spawnSubagentDirect({ task: "clean worker" }, { agentSessionKey: "main" });
 
@@ -107,20 +116,17 @@ describe("sessions_spawn context modes", () => {
     const store: SessionStore = {
       main: { sessionId: "parent-session-id", updatedAt: 1 },
     };
+    usePersistentStoreMock(store);
     const rollback = vi.fn(async () => undefined);
-    const callGatewayMock = vi.fn(async (requestUnknown: unknown) => {
+    callGatewayMock.mockImplementation(async (requestUnknown: unknown) => {
       const request = requestUnknown as GatewayRequest;
       if (request.method === "agent") {
         throw new Error("agent start failed");
       }
       return { ok: true };
     });
-    const { spawnSubagentDirect } = await loadSubagentSpawnModuleForTest({
-      callGatewayMock,
-      updateSessionStoreMock: createPersistentStoreMock(store),
-      resolveContextEngineMock: vi.fn(async () => ({
-        prepareSubagentSpawn: vi.fn(async () => ({ rollback })),
-      })),
+    resolveContextEngineMock.mockResolvedValue({
+      prepareSubagentSpawn: vi.fn(async () => ({ rollback })),
     });
 
     const result = await spawnSubagentDirect({ task: "clean worker" }, { agentSessionKey: "main" });