From 7aedb6d4420c05dfffc0312b10feae098625a728 Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Sat, 14 Mar 2026 02:21:13 +0000
Subject: [PATCH] test: share subagent gateway mock setup

---
 src/agents/subagent-spawn.workspace.test.ts | 77 ++++++++-------------
 src/agents/test-helpers/subagent-gateway.ts |  9 +++
 2 files changed, 38 insertions(+), 48 deletions(-)
 create mode 100644 src/agents/test-helpers/subagent-gateway.ts

diff --git a/src/agents/subagent-spawn.workspace.test.ts b/src/agents/subagent-spawn.workspace.test.ts
index fef6bc7515c..9955e587c89 100644
--- a/src/agents/subagent-spawn.workspace.test.ts
+++ b/src/agents/subagent-spawn.workspace.test.ts
@@ -1,5 +1,6 @@
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { spawnSubagentDirect } from "./subagent-spawn.js";
+import { installAcceptedSubagentGatewayMock } from "./test-helpers/subagent-gateway.js";
 
 type TestAgentConfig = {
   id?: string;
@@ -100,20 +101,7 @@ function createConfigOverride(overrides?: Record<string, unknown>) {
 }
 
 function setupGatewayMock() {
-  hoisted.callGatewayMock.mockImplementation(
-    async (opts: { method?: string; params?: Record<string, unknown> }) => {
-      if (opts.method === "sessions.patch") {
-        return { ok: true };
-      }
-      if (opts.method === "sessions.delete") {
-        return { ok: true };
-      }
-      if (opts.method === "agent") {
-        return { runId: "run-1" };
-      }
-      return {};
-    },
-  );
+  installAcceptedSubagentGatewayMock(hoisted.callGatewayMock);
 }
 
 function getRegisteredRun() {
@@ -122,6 +110,27 @@ function getRegisteredRun() {
     | undefined;
 }
 
+async function expectAcceptedWorkspace(params: { agentId: string; expectedWorkspaceDir: string }) {
+  const result = await spawnSubagentDirect(
+    {
+      task: "inspect workspace",
+      agentId: params.agentId,
+    },
+    {
+      agentSessionKey: "agent:main:main",
+      agentChannel: "telegram",
+      agentAccountId: "123",
+      agentTo: "456",
+      workspaceDir: "/tmp/requester-workspace",
+    },
+  );
+
+  expect(result.status).toBe("accepted");
+  expect(getRegisteredRun()).toMatchObject({
+    workspaceDir: params.expectedWorkspaceDir,
+  });
+}
+
 describe("spawnSubagentDirect workspace inheritance", () => {
   beforeEach(() => {
     hoisted.callGatewayMock.mockClear();
@@ -149,44 +158,16 @@ describe("spawnSubagentDirect workspace inheritance", () => {
       },
     });
 
-    const result = await spawnSubagentDirect(
-      {
-        task: "inspect workspace",
-        agentId: "ops",
-      },
-      {
-        agentSessionKey: "agent:main:main",
-        agentChannel: "telegram",
-        agentAccountId: "123",
-        agentTo: "456",
-        workspaceDir: "/tmp/requester-workspace",
-      },
-    );
-
-    expect(result.status).toBe("accepted");
-    expect(getRegisteredRun()).toMatchObject({
-      workspaceDir: "/tmp/workspace-ops",
+    await expectAcceptedWorkspace({
+      agentId: "ops",
+      expectedWorkspaceDir: "/tmp/workspace-ops",
     });
   });
 
   it("preserves the inherited workspace for same-agent spawns", async () => {
-    const result = await spawnSubagentDirect(
-      {
-        task: "inspect workspace",
-        agentId: "main",
-      },
-      {
-        agentSessionKey: "agent:main:main",
-        agentChannel: "telegram",
-        agentAccountId: "123",
-        agentTo: "456",
-        workspaceDir: "/tmp/requester-workspace",
-      },
-    );
-
-    expect(result.status).toBe("accepted");
-    expect(getRegisteredRun()).toMatchObject({
-      workspaceDir: "/tmp/requester-workspace",
+    await expectAcceptedWorkspace({
+      agentId: "main",
+      expectedWorkspaceDir: "/tmp/requester-workspace",
     });
   });
 });
diff --git a/src/agents/test-helpers/subagent-gateway.ts b/src/agents/test-helpers/subagent-gateway.ts
new file mode 100644
index 00000000000..9491d971c33
--- /dev/null
+++ b/src/agents/test-helpers/subagent-gateway.ts
@@ -0,0 +1,9 @@
+export function installAcceptedSubagentGatewayMock(mock: {
+  mockImplementation: (
+    impl: (opts: { method?: string; params?: unknown }) => Promise<unknown>,
+  ) => unknown;
+}) {
+  mock.mockImplementation(async ({ method }) =>
+    method === "agent" ? { runId: "run-1" } : method?.startsWith("sessions.") ? { ok: true } : {},
+  );
+}