From d4b193b581de7eb041739d5402cca2b699e4e10d Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Sat, 14 Mar 2026 01:06:06 +0000
Subject: [PATCH] test: share embedded workspace attempt helpers

---
 .../run/attempt.spawn-workspace.test.ts       | 90 +++++++------------
 1 file changed, 34 insertions(+), 56 deletions(-)

diff --git a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test.ts b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test.ts
index 53edfbbc6cc..e67bb20d88d 100644
--- a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test.ts
+++ b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test.ts
@@ -290,7 +290,9 @@ async function cleanupTempPaths(tempPaths: string[]) {
   }
 }
 
-function createDefaultEmbeddedSession(): MutableSession {
+function createDefaultEmbeddedSession(params?: {
+  prompt?: (session: MutableSession) => Promise<void>;
+}): MutableSession {
   const session: MutableSession = {
     sessionId: "embedded-session",
     messages: [],
@@ -302,6 +304,10 @@ function createDefaultEmbeddedSession(): MutableSession {
       },
     },
     prompt: async () => {
+      if (params?.prompt) {
+        await params.prompt(session);
+        return;
+      }
       session.messages = [
         ...session.messages,
         { role: "assistant", content: "done", timestamp: 2 },
@@ -315,6 +321,24 @@ function createDefaultEmbeddedSession(): MutableSession {
   return session;
 }
 
+function createContextEngineBootstrapAndAssemble() {
+  return {
+    bootstrap: vi.fn(async (_params: { sessionKey?: string }) => ({ bootstrapped: true })),
+    assemble: vi.fn(async ({ messages }: { messages: AgentMessage[]; sessionKey?: string }) => ({
+      messages,
+      estimatedTokens: 1,
+    })),
+  };
+}
+
+function expectCalledWithSessionKey(mock: ReturnType<typeof vi.fn>, sessionKey: string) {
+  expect(mock).toHaveBeenCalledWith(
+    expect.objectContaining({
+      sessionKey,
+    }),
+  );
+}
+
 const testModel = {
   api: "openai-completions",
   provider: "openai",
@@ -366,16 +390,7 @@ describe("runEmbeddedAttempt sessions_spawn workspace inheritance", () => {
 
     hoisted.createAgentSessionMock.mockImplementation(
       async (params: { customTools: ToolDefinition[] }) => {
-        const session: MutableSession = {
-          sessionId: "embedded-session",
-          messages: [],
-          isCompacting: false,
-          isStreaming: false,
-          agent: {
-            replaceMessages: (messages: unknown[]) => {
-              session.messages = [...messages];
-            },
-          },
+        const session = createDefaultEmbeddedSession({
           prompt: async () => {
             const spawnTool = params.customTools.find((tool) => tool.name === "sessions_spawn");
             expect(spawnTool).toBeDefined();
@@ -390,10 +405,7 @@ describe("runEmbeddedAttempt sessions_spawn workspace inheritance", () => {
               {} as unknown as ExtensionContext,
             );
           },
-          abort: async () => {},
-          dispose: () => {},
-          steer: async () => {},
-        };
+        });
 
         return { session };
       },
@@ -650,13 +662,7 @@ describe("runEmbeddedAttempt context engine sessionKey forwarding", () => {
   }
 
   it("forwards sessionKey to bootstrap, assemble, and afterTurn", async () => {
-    const bootstrap = vi.fn(async (_params: { sessionKey?: string }) => ({ bootstrapped: true }));
-    const assemble = vi.fn(
-      async ({ messages }: { messages: AgentMessage[]; sessionKey?: string }) => ({
-        messages,
-        estimatedTokens: 1,
-      }),
-    );
+    const { bootstrap, assemble } = createContextEngineBootstrapAndAssemble();
     const afterTurn = vi.fn(async (_params: { sessionKey?: string }) => {});
 
     const result = await runAttemptWithContextEngine({
@@ -666,31 +672,13 @@ describe("runEmbeddedAttempt context engine sessionKey forwarding", () => {
     });
 
     expect(result.promptError).toBeNull();
-    expect(bootstrap).toHaveBeenCalledWith(
-      expect.objectContaining({
-        sessionKey,
-      }),
-    );
-    expect(assemble).toHaveBeenCalledWith(
-      expect.objectContaining({
-        sessionKey,
-      }),
-    );
-    expect(afterTurn).toHaveBeenCalledWith(
-      expect.objectContaining({
-        sessionKey,
-      }),
-    );
+    expectCalledWithSessionKey(bootstrap, sessionKey);
+    expectCalledWithSessionKey(assemble, sessionKey);
+    expectCalledWithSessionKey(afterTurn, sessionKey);
   });
 
   it("forwards sessionKey to ingestBatch when afterTurn is absent", async () => {
-    const bootstrap = vi.fn(async (_params: { sessionKey?: string }) => ({ bootstrapped: true }));
-    const assemble = vi.fn(
-      async ({ messages }: { messages: AgentMessage[]; sessionKey?: string }) => ({
-        messages,
-        estimatedTokens: 1,
-      }),
-    );
+    const { bootstrap, assemble } = createContextEngineBootstrapAndAssemble();
     const ingestBatch = vi.fn(
       async (_params: { sessionKey?: string; messages: AgentMessage[] }) => ({ ingestedCount: 1 }),
     );
@@ -702,21 +690,11 @@ describe("runEmbeddedAttempt context engine sessionKey forwarding", () => {
     });
 
     expect(result.promptError).toBeNull();
-    expect(ingestBatch).toHaveBeenCalledWith(
-      expect.objectContaining({
-        sessionKey,
-      }),
-    );
+    expectCalledWithSessionKey(ingestBatch, sessionKey);
   });
 
   it("forwards sessionKey to per-message ingest when ingestBatch is absent", async () => {
-    const bootstrap = vi.fn(async (_params: { sessionKey?: string }) => ({ bootstrapped: true }));
-    const assemble = vi.fn(
-      async ({ messages }: { messages: AgentMessage[]; sessionKey?: string }) => ({
-        messages,
-        estimatedTokens: 1,
-      }),
-    );
+    const { bootstrap, assemble } = createContextEngineBootstrapAndAssemble();
     const ingest = vi.fn(async (_params: { sessionKey?: string; message: AgentMessage }) => ({
       ingested: true,
     }));