refactor(tests): dedupe pi embedded test harness

src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.subscribeembeddedpisession.test.ts

@@ -11,10 +11,6 @@ import {
 } from "./pi-embedded-subscribe.e2e-harness.js";
 import { subscribeEmbeddedPiSession } from "./pi-embedded-subscribe.js";
 
-type StubSession = {
-  subscribe: (fn: (evt: unknown) => void) => () => void;
-};
-
 describe("subscribeEmbeddedPiSession", () => {
   function createAgentEventHarness(options?: { runId?: string; sessionKey?: string }) {
     const { session, emit } = createStubSessionHarness();
@@ -41,6 +37,32 @@ describe("subscribeEmbeddedPiSession", () => {
     return { emit, subscription };
   }
 
+  function createSubscribedHarness(
+    options: Omit<Parameters<typeof subscribeEmbeddedPiSession>[0], "session">,
+  ) {
+    const { session, emit } = createStubSessionHarness();
+    subscribeEmbeddedPiSession({
+      session,
+      ...options,
+    });
+    return { emit };
+  }
+
+  function emitAssistantTextDelta(
+    emit: (evt: unknown) => void,
+    delta: string,
+    message: Record<string, unknown> = { role: "assistant" },
+  ) {
+    emit({
+      type: "message_update",
+      message,
+      assistantMessageEvent: {
+        type: "text_delta",
+        delta,
+      },
+    });
+  }
+
   function createWriteFailureHarness(params: {
     runId: string;
     path: string;
@@ -85,19 +107,10 @@ describe("subscribeEmbeddedPiSession", () => {
   it.each(THINKING_TAG_CASES)(
     "streams <%s> reasoning via onReasoningStream without leaking into final text",
     ({ open, close }) => {
-      let handler: ((evt: unknown) => void) | undefined;
-      const session: StubSession = {
-        subscribe: (fn) => {
-          handler = fn;
-          return () => {};
-        },
-      };
-
       const onReasoningStream = vi.fn();
       const onBlockReply = vi.fn();
 
-      subscribeEmbeddedPiSession({
-        session: session as unknown as Parameters<typeof subscribeEmbeddedPiSession>[0]["session"],
+      const { emit } = createSubscribedHarness({
         runId: "run",
         onReasoningStream,
         onBlockReply,
@@ -105,23 +118,8 @@ describe("subscribeEmbeddedPiSession", () => {
         reasoningMode: "stream",
       });
 
-      handler?.({
-        type: "message_update",
-        message: { role: "assistant" },
-        assistantMessageEvent: {
-          type: "text_delta",
-          delta: `${open}\nBecause`,
-        },
-      });
-
-      handler?.({
-        type: "message_update",
-        message: { role: "assistant" },
-        assistantMessageEvent: {
-          type: "text_delta",
-          delta: ` it helps\n${close}\n\nFinal answer`,
-        },
-      });
+      emitAssistantTextDelta(emit, `${open}\nBecause`);
+      emitAssistantTextDelta(emit, ` it helps\n${close}\n\nFinal answer`);
 
       const assistantMessage = {
         role: "assistant",
@@ -133,7 +131,7 @@ describe("subscribeEmbeddedPiSession", () => {
         ],
       } as AssistantMessage;
 
-      handler?.({ type: "message_end", message: assistantMessage });
+      emit({ type: "message_end", message: assistantMessage });
 
       expect(onBlockReply).toHaveBeenCalledTimes(1);
       expect(onBlockReply.mock.calls[0][0].text).toBe("Final answer");
@@ -152,18 +150,9 @@ describe("subscribeEmbeddedPiSession", () => {
   it.each(THINKING_TAG_CASES)(
     "suppresses <%s> blocks across chunk boundaries",
     ({ open, close }) => {
-      let handler: ((evt: unknown) => void) | undefined;
-      const session: StubSession = {
-        subscribe: (fn) => {
-          handler = fn;
-          return () => {};
-        },
-      };
-
       const onBlockReply = vi.fn();
 
-      subscribeEmbeddedPiSession({
-        session: session as unknown as Parameters<typeof subscribeEmbeddedPiSession>[0]["session"],
+      const { emit } = createSubscribedHarness({
         runId: "run",
         onBlockReply,
         blockReplyBreak: "text_end",
@@ -174,29 +163,13 @@ describe("subscribeEmbeddedPiSession", () => {
         },
       });
 
-      handler?.({ type: "message_start", message: { role: "assistant" } });
-
-      handler?.({
-        type: "message_update",
-        message: { role: "assistant" },
-        assistantMessageEvent: {
-          type: "text_delta",
-          delta: `${open}Reasoning chunk that should not leak`,
-        },
-      });
+      emit({ type: "message_start", message: { role: "assistant" } });
+      emitAssistantTextDelta(emit, `${open}Reasoning chunk that should not leak`);
 
       expect(onBlockReply).not.toHaveBeenCalled();
 
-      handler?.({
-        type: "message_update",
-        message: { role: "assistant" },
-        assistantMessageEvent: {
-          type: "text_delta",
-          delta: `${close}\n\nFinal answer`,
-        },
-      });
-
-      handler?.({
+      emitAssistantTextDelta(emit, `${close}\n\nFinal answer`);
+      emit({
         type: "message_update",
         message: { role: "assistant" },
         assistantMessageEvent: { type: "text_end" },
@@ -216,26 +189,17 @@ describe("subscribeEmbeddedPiSession", () => {
   );
 
   it("streams native thinking_delta events and signals reasoning end", () => {
-    let handler: ((evt: unknown) => void) | undefined;
-    const session: StubSession = {
-      subscribe: (fn) => {
-        handler = fn;
-        return () => {};
-      },
-    };
-
     const onReasoningStream = vi.fn();
     const onReasoningEnd = vi.fn();
 
-    subscribeEmbeddedPiSession({
-      session: session as unknown as Parameters<typeof subscribeEmbeddedPiSession>[0]["session"],
+    const { emit } = createSubscribedHarness({
       runId: "run",
       reasoningMode: "stream",
       onReasoningStream,
       onReasoningEnd,
     });
 
-    handler?.({
+    emit({
       type: "message_update",
       message: {
         role: "assistant",
@@ -247,7 +211,7 @@ describe("subscribeEmbeddedPiSession", () => {
       },
     });
 
-    handler?.({
+    emit({
       type: "message_update",
       message: {
         role: "assistant",
@@ -266,36 +230,18 @@ describe("subscribeEmbeddedPiSession", () => {
   });
 
   it("emits reasoning end once when native and tagged reasoning end overlap", () => {
-    let handler: ((evt: unknown) => void) | undefined;
-    const session: StubSession = {
-      subscribe: (fn) => {
-        handler = fn;
-        return () => {};
-      },
-    };
-
     const onReasoningEnd = vi.fn();
 
-    subscribeEmbeddedPiSession({
-      session: session as unknown as Parameters<typeof subscribeEmbeddedPiSession>[0]["session"],
+    const { emit } = createSubscribedHarness({
       runId: "run",
       reasoningMode: "stream",
       onReasoningStream: vi.fn(),
       onReasoningEnd,
     });
 
-    handler?.({ type: "message_start", message: { role: "assistant" } });
-
-    handler?.({
-      type: "message_update",
-      message: { role: "assistant" },
-      assistantMessageEvent: {
-        type: "text_delta",
-        delta: "<think>Checking",
-      },
-    });
-
-    handler?.({
+    emit({ type: "message_start", message: { role: "assistant" } });
+    emitAssistantTextDelta(emit, "<think>Checking");
+    emit({
       type: "message_update",
       message: {
         role: "assistant",
@@ -306,14 +252,7 @@ describe("subscribeEmbeddedPiSession", () => {
       },
     });
 
-    handler?.({
-      type: "message_update",
-      message: { role: "assistant" },
-      assistantMessageEvent: {
-        type: "text_delta",
-        delta: " files</think>\nFinal answer",
-      },
-    });
+    emitAssistantTextDelta(emit, " files</think>\nFinal answer");
 
     expect(onReasoningEnd).toHaveBeenCalledTimes(1);
   });
@@ -374,16 +313,8 @@ describe("subscribeEmbeddedPiSession", () => {
     const { emit, onAgentEvent } = createAgentEventHarness();
 
     emit({ type: "message_start", message: { role: "assistant" } });
-    emit({
-      type: "message_update",
-      message: { role: "assistant" },
-      assistantMessageEvent: { type: "text_delta", delta: "MEDIA:" },
-    });
-    emit({
-      type: "message_update",
-      message: { role: "assistant" },
-      assistantMessageEvent: { type: "text_delta", delta: " https://example.com/a.png\nCaption" },
-    });
+    emitAssistantTextDelta(emit, "MEDIA:");
+    emitAssistantTextDelta(emit, " https://example.com/a.png\nCaption");
 
     const payloads = extractAgentEventPayloads(onAgentEvent.mock.calls);
     expect(payloads).toHaveLength(1);
@@ -394,11 +325,7 @@ describe("subscribeEmbeddedPiSession", () => {
     const { emit, onAgentEvent } = createAgentEventHarness();
 
     emit({ type: "message_start", message: { role: "assistant" } });
-    emit({
-      type: "message_update",
-      message: { role: "assistant" },
-      assistantMessageEvent: { type: "text_delta", delta: "MEDIA: https://example.com/a.png" },
-    });
+    emitAssistantTextDelta(emit, "MEDIA: https://example.com/a.png");
 
     const payloads = extractAgentEventPayloads(onAgentEvent.mock.calls);
     expect(payloads).toHaveLength(1);