test: stabilize embedded runner tests

src/agents/pi-embedded-runner.run-embedded-pi-agent.appends-new-user-assistant-after-existing-transcript.test.ts

@@ -78,6 +78,7 @@ vi.mock("@mariozechner/pi-ai", async () => {
               ? buildAssistantErrorMessage(model)
               : buildAssistantMessage(model),
         });
+        stream.end();
       });
       return stream;
     },
@@ -113,6 +114,9 @@ const makeOpenAiConfig = (modelIds: string[]) =>
 const ensureModels = (cfg: ClawdbotConfig, agentDir: string) =>
   ensureClawdbotModelsJson(cfg, agentDir);
 
+const testSessionKey = "agent:test:embedded-ordering";
+const immediateEnqueue = async <T>(task: () => Promise<T>) => task();
+
 const textFromContent = (content: unknown) => {
   if (typeof content === "string") return content;
   if (Array.isArray(content) && content[0]?.type === "text") {
@@ -179,7 +183,7 @@ describe("runEmbeddedPiAgent", () => {
 
     await runEmbeddedPiAgent({
       sessionId: "session:test",
-      sessionKey: "agent:main:main",
+      sessionKey: testSessionKey,
       sessionFile,
       workspaceDir,
       config: cfg,
@@ -188,6 +192,7 @@ describe("runEmbeddedPiAgent", () => {
       model: "mock-1",
       timeoutMs: 5_000,
       agentDir,
+      enqueue: immediateEnqueue,
     });
 
     const messages = await readSessionMessages(sessionFile);
@@ -219,7 +224,7 @@ describe("runEmbeddedPiAgent", () => {
 
     await runEmbeddedPiAgent({
       sessionId: "session:test",
-      sessionKey: "agent:main:main",
+      sessionKey: testSessionKey,
       sessionFile,
       workspaceDir,
       config: cfg,
@@ -228,11 +233,12 @@ describe("runEmbeddedPiAgent", () => {
       model: "mock-1",
       timeoutMs: 5_000,
       agentDir,
+      enqueue: immediateEnqueue,
     });
 
     await runEmbeddedPiAgent({
       sessionId: "session:test",
-      sessionKey: "agent:main:main",
+      sessionKey: testSessionKey,
       sessionFile,
       workspaceDir,
       config: cfg,
@@ -241,6 +247,7 @@ describe("runEmbeddedPiAgent", () => {
       model: "mock-1",
       timeoutMs: 5_000,
       agentDir,
+      enqueue: immediateEnqueue,
     });
 
     const messages = await readSessionMessages(sessionFile);
@@ -306,7 +313,7 @@ describe("runEmbeddedPiAgent", () => {
 
     const result = await runEmbeddedPiAgent({
       sessionId: "session:test",
-      sessionKey: "agent:main:main",
+      sessionKey: testSessionKey,
       sessionFile,
       workspaceDir,
       config: cfg,
@@ -315,6 +322,7 @@ describe("runEmbeddedPiAgent", () => {
       model: "mock-1",
       timeoutMs: 5_000,
       agentDir,
+      enqueue: immediateEnqueue,
     });
 
     expect(result.meta.error).toBeUndefined();
@@ -339,7 +347,7 @@ describe("runEmbeddedPiAgent", () => {
 
     const result = await runEmbeddedPiAgent({
       sessionId: "session:test",
-      sessionKey: "agent:main:main",
+      sessionKey: testSessionKey,
       sessionFile,
       workspaceDir,
       config: cfg,
@@ -348,6 +356,7 @@ describe("runEmbeddedPiAgent", () => {
       model: "mock-1",
       timeoutMs: 5_000,
       agentDir,
+      enqueue: immediateEnqueue,
     });
 
     expect(result.meta.error).toBeUndefined();

src/agents/pi-embedded-runner.run-embedded-pi-agent.writes-models-json-into-provided-agentdir.test.ts

@@ -78,6 +78,7 @@ vi.mock("@mariozechner/pi-ai", async () => {
               ? buildAssistantErrorMessage(model)
               : buildAssistantMessage(model),
         });
+        stream.end();
       });
       return stream;
     },
@@ -113,6 +114,9 @@ const makeOpenAiConfig = (modelIds: string[]) =>
 const ensureModels = (cfg: ClawdbotConfig, agentDir: string) =>
   ensureClawdbotModelsJson(cfg, agentDir);
 
+const testSessionKey = "agent:test:embedded-models";
+const immediateEnqueue = async <T>(task: () => Promise<T>) => task();
+
 const textFromContent = (content: unknown) => {
   if (typeof content === "string") return content;
   if (Array.isArray(content) && content[0]?.type === "text") {
@@ -169,7 +173,7 @@ describe("runEmbeddedPiAgent", () => {
     await expect(
       runEmbeddedPiAgent({
         sessionId: "session:test",
-        sessionKey: "agent:dev:test",
+        sessionKey: testSessionKey,
         sessionFile,
         workspaceDir,
         config: cfg,
@@ -178,6 +182,7 @@ describe("runEmbeddedPiAgent", () => {
         model: "definitely-not-a-model",
         timeoutMs: 1,
         agentDir,
+        enqueue: immediateEnqueue,
       }),
     ).rejects.toThrow(/Unknown model:/);
 
@@ -193,7 +198,7 @@ describe("runEmbeddedPiAgent", () => {
 
     await runEmbeddedPiAgent({
       sessionId: "session:test",
-      sessionKey: "agent:main:main",
+      sessionKey: testSessionKey,
       sessionFile,
       workspaceDir,
       config: cfg,
@@ -202,6 +207,7 @@ describe("runEmbeddedPiAgent", () => {
       model: "mock-1",
       timeoutMs: 5_000,
       agentDir,
+      enqueue: immediateEnqueue,
     });
 
     const messages = await readSessionMessages(sessionFile);
@@ -224,7 +230,7 @@ describe("runEmbeddedPiAgent", () => {
 
     const result = await runEmbeddedPiAgent({
       sessionId: "session:test",
-      sessionKey: "agent:main:main",
+      sessionKey: testSessionKey,
       sessionFile,
       workspaceDir,
       config: cfg,
@@ -233,6 +239,7 @@ describe("runEmbeddedPiAgent", () => {
       model: "mock-error",
       timeoutMs: 5_000,
       agentDir,
+      enqueue: immediateEnqueue,
     });
     expect(result.payloads[0]?.isError).toBe(true);
 

src/agents/pi-embedded-runner.sanitize-session-history.test.ts

@@ -2,7 +2,9 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core";
 import type { SessionManager } from "@mariozechner/pi-coding-agent";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import * as helpers from "./pi-embedded-helpers.js";
-import { sanitizeSessionHistory } from "./pi-embedded-runner/google.js";
+
+type SanitizeSessionHistory = typeof import("./pi-embedded-runner/google.js").sanitizeSessionHistory;
+let sanitizeSessionHistory: SanitizeSessionHistory;
 
 // Mock dependencies
 vi.mock("./pi-embedded-helpers.js", async () => {
@@ -26,7 +28,7 @@ describe("sanitizeSessionHistory", () => {
 
   const mockMessages: AgentMessage[] = [{ role: "user", content: "hello" }];
 
-  beforeEach(() => {
+  beforeEach(async () => {
     vi.resetAllMocks();
     vi.mocked(helpers.sanitizeSessionMessagesImages).mockImplementation(async (msgs) => msgs);
     // Default mock implementation
@@ -34,6 +36,8 @@ describe("sanitizeSessionHistory", () => {
       if (!msgs) return [];
       return [...msgs, { role: "system", content: "downgraded" }];
     });
+    vi.resetModules();
+    ({ sanitizeSessionHistory } = await import("./pi-embedded-runner/google.js"));
   });
 
   it("should downgrade history for Google models if provider is not google-antigravity", async () => {