test(sdk): cover pipeAndCapture status/partial-capture and writeTurnComplete lastEventId

Adds unit tests for the new chat.pipeAndCapture result shape (complete /
aborted-with-partial / error) and for chat.writeTurnComplete returning a
resume cursor. Extends the in-memory session test harness so its .out
writeControl projects a synthetic seq_num and its pipe propagates
source-stream errors (mirroring production) instead of swallowing them.

Author: matt-aitken

packages/trigger-sdk/src/v3/test/test-session-handle.ts

@@ -124,6 +124,7 @@ export class TestSessionOutputChannel extends SessionOutputChannel {
   ): PipeStreamResult<T> {
     const state = this.state;
     const readChunks: T[] = [];
+    let pipeError: unknown;
     let resolveDone!: () => void;
     const done = new Promise<void>((resolve) => {
       resolveDone = resolve;
@@ -135,10 +136,15 @@ export class TestSessionOutputChannel extends SessionOutputChannel {
       try {
         while (true) {
           const { done: d, value: v } = await reader.read();
-          if (d) return;
+          if (d) break;
           readChunks.push(v as T);
           notify(state, v);
         }
+      } catch (err) {
+        // Mirror production: a source-stream error rejects waitUntilComplete
+        // instead of being silently swallowed, so callers (e.g.
+        // chat.pipeAndCapture) can observe the failure.
+        pipeError = err;
       } finally {
         try {
           reader.releaseLock();
@@ -147,9 +153,7 @@ export class TestSessionOutputChannel extends SessionOutputChannel {
         }
         resolveDone();
       }
-    })().catch(() => {
-      resolveDone();
-    });
+    })();
 
     const replayStream = new ReadableStream<T>({
       async start(controller) {
@@ -167,6 +171,7 @@ export class TestSessionOutputChannel extends SessionOutputChannel {
       },
       waitUntilComplete: async () => {
         await done;
+        if (pipeError) throw pipeError;
         return emptyResult;
       },
     };
@@ -262,7 +267,11 @@ export class TestSessionOutputChannel extends SessionOutputChannel {
       }
     }
     notify(this.state, synthetic);
-    return {};
+    // Project a synthetic monotonic seq_num as the ack's `lastEventId`,
+    // mirroring what S2 returns in production (there it's the control
+    // record's seq_num). Using the running `.out` record count lets
+    // `chat.writeTurnComplete()` surface a real resume cursor in tests.
+    return { lastEventId: String(this.state.chunks.length) };
   }
 
   /**

packages/trigger-sdk/test/chat-pipe-and-capture.test.ts

@@ -0,0 +1,216 @@
+// Import the test harness FIRST — this installs the resource catalog so
+// `chat.customAgent()` calls below register their task functions correctly.
+import { mockChatAgent } from "../src/v3/test/index.js";
+
+import { describe, expect, it } from "vitest";
+import type { UIMessage } from "ai";
+import { simulateReadableStream, streamText } from "ai";
+import { MockLanguageModelV3 } from "ai/test";
+import type { LanguageModelV3StreamPart } from "@ai-sdk/provider";
+import { chat } from "../src/v3/ai.js";
+import type { PipeAndCaptureResult } from "../src/v3/ai.js";
+
+// ── Helpers ────────────────────────────────────────────────────────────
+
+function userMessage(text: string, id: string): UIMessage {
+  return { id, role: "user", parts: [{ type: "text", text }] };
+}
+
+function textChunks(text: string, opts?: { split?: boolean }): LanguageModelV3StreamPart[] {
+  const deltas = opts?.split ? text.split(" ").map((w, i) => (i === 0 ? w : ` ${w}`)) : [text];
+  return [
+    { type: "text-start", id: "t1" },
+    ...deltas.map((delta) => ({ type: "text-delta" as const, id: "t1", delta })),
+    { type: "text-end", id: "t1" },
+    {
+      type: "finish",
+      finishReason: { unified: "stop", raw: "stop" },
+      usage: {
+        inputTokens: { total: 5, noCache: 5, cacheRead: undefined, cacheWrite: undefined },
+        outputTokens: { total: 5, text: 5, reasoning: undefined },
+      },
+    },
+  ];
+}
+
+/** Model that streams `text` in one fast pass with a `stop` finish. */
+function fastModel(text: string) {
+  return new MockLanguageModelV3({
+    doStream: async () => ({ stream: simulateReadableStream({ chunks: textChunks(text) }) }),
+  });
+}
+
+/** Model that streams `text` word-by-word with a wide gap before the final
+ *  chunk, leaving a window to abort mid-stream after the first delta. */
+function slowModel(text: string) {
+  return new MockLanguageModelV3({
+    doStream: async () => ({
+      stream: simulateReadableStream({
+        chunks: textChunks(text, { split: true }),
+        initialDelayInMs: 0,
+        chunkDelayInMs: 500,
+      }),
+    }),
+  });
+}
+
+function extractText(message: UIMessage | undefined): string {
+  if (!message) return "";
+  return (message.parts as Array<{ type: string; text?: string }>)
+    .filter((p) => p.type === "text")
+    .map((p) => p.text ?? "")
+    .join("");
+}
+
+async function waitFor(check: () => boolean, timeoutMs = 5_000) {
+  const start = Date.now();
+  while (Date.now() - start < timeoutMs) {
+    if (check()) return;
+    await new Promise((r) => setTimeout(r, 20));
+  }
+  throw new Error("waitFor timed out");
+}
+
+function deltaCount(harness: { allChunks: unknown[] }): number {
+  return (harness.allChunks as { type?: string }[]).filter((c) => c.type === "text-delta").length;
+}
+
+// ── Tests ──────────────────────────────────────────────────────────────
+
+describe("chat.pipeAndCapture", () => {
+  it("returns status 'complete' with the message and finish reason on a normal turn", async () => {
+    const captures: PipeAndCaptureResult[] = [];
+    const lastEventIds: Array<string | undefined> = [];
+
+    const agent = chat.customAgent({
+      id: "pipe-capture.complete",
+      run: async () => {
+        const conversation = new chat.MessageAccumulator();
+        const next = await chat.messages.waitWithIdleTimeout({
+          idleTimeoutInSeconds: 60,
+          timeout: "1h",
+        });
+        if (!next.ok) return;
+        const wire = next.output as { message?: UIMessage; trigger: string };
+        const incoming = wire.message ? [wire.message] : [];
+        const messages = await conversation.addIncoming(incoming, wire.trigger, 0);
+        const result = streamText({ model: fastModel("hello world"), messages });
+        const captured = await chat.pipeAndCapture(result);
+        captures.push(captured);
+        if (captured.message) await conversation.addResponse(captured.message);
+        const { lastEventId } = await chat.writeTurnComplete();
+        lastEventIds.push(lastEventId);
+      },
+    });
+
+    const harness = mockChatAgent(agent, { chatId: "pc-complete" });
+    try {
+      await harness.sendMessage(userMessage("hi", "u-1"));
+      await waitFor(() => captures.length >= 1 && lastEventIds.length >= 1);
+
+      expect(captures[0]!.status).toBe("complete");
+      expect(extractText(captures[0]!.message)).toBe("hello world");
+      expect(captures[0]!.finishReason).toBe("stop");
+      expect(captures[0]!.error).toBeUndefined();
+      // chat.writeTurnComplete() surfaces the resume cursor for the next turn.
+      expect(typeof lastEventIds[0]).toBe("string");
+      expect(lastEventIds[0]!.length).toBeGreaterThan(0);
+    } finally {
+      await harness.close();
+    }
+  });
+
+  it("returns status 'error' with the thrown error and does not throw when the stream fails", async () => {
+    const captures: PipeAndCaptureResult[] = [];
+    let runThrew = false;
+
+    // Synthetic source whose UI stream errors after emitting a partial. This
+    // deterministically drives the pipe-failure path without depending on the
+    // AI SDK's model-error handling.
+    const erroringSource = {
+      toUIMessageStream() {
+        return new ReadableStream({
+          start(controller) {
+            controller.enqueue({ type: "start", messageId: "a-err" });
+            controller.enqueue({ type: "text-start", id: "t1" });
+            controller.enqueue({ type: "text-delta", id: "t1", delta: "partial" });
+            controller.error(new Error("boom"));
+          },
+        });
+      },
+    };
+
+    const agent = chat.customAgent({
+      id: "pipe-capture.error",
+      run: async () => {
+        const next = await chat.messages.waitWithIdleTimeout({
+          idleTimeoutInSeconds: 60,
+          timeout: "1h",
+        });
+        if (!next.ok) return;
+        try {
+          captures.push(await chat.pipeAndCapture(erroringSource as never));
+        } catch {
+          runThrew = true;
+        }
+        await chat.writeTurnComplete();
+      },
+    });
+
+    const harness = mockChatAgent(agent, { chatId: "pc-error" });
+    try {
+      await harness.sendMessage(userMessage("go", "u-1"));
+      await waitFor(() => captures.length >= 1);
+
+      expect(runThrew).toBe(false);
+      expect(captures[0]!.status).toBe("error");
+      expect(captures[0]!.error).toBeInstanceOf(Error);
+      expect((captures[0]!.error as Error).message).toBe("boom");
+    } finally {
+      await harness.close();
+    }
+  });
+
+  it("returns status 'aborted' and preserves the partial message on a mid-stream stop", async () => {
+    const captures: PipeAndCaptureResult[] = [];
+
+    const agent = chat.customAgent({
+      id: "pipe-capture.aborted",
+      run: async () => {
+        const stop = chat.createStopSignal();
+        const conversation = new chat.MessageAccumulator();
+        const next = await chat.messages.waitWithIdleTimeout({
+          idleTimeoutInSeconds: 60,
+          timeout: "1h",
+        });
+        if (!next.ok) return;
+        const wire = next.output as { message?: UIMessage; trigger: string };
+        const incoming = wire.message ? [wire.message] : [];
+        const messages = await conversation.addIncoming(incoming, wire.trigger, 0);
+        const result = streamText({
+          model: slowModel("one two three four"),
+          messages,
+          abortSignal: stop.signal,
+        });
+        captures.push(await chat.pipeAndCapture(result, { signal: stop.signal }));
+        await chat.writeTurnComplete();
+        stop.cleanup();
+      },
+    });
+
+    const harness = mockChatAgent(agent, { chatId: "pc-aborted" });
+    try {
+      void harness.sendMessage(userMessage("hi", "u-1"));
+      // Stop once the first delta has streamed but before the turn finishes.
+      await waitFor(() => deltaCount(harness) >= 1);
+      await harness.sendStop();
+      await waitFor(() => captures.length >= 1);
+
+      expect(captures[0]!.status).toBe("aborted");
+      // The partial that streamed before the stop is preserved.
+      expect(extractText(captures[0]!.message).startsWith("one")).toBe(true);
+    } finally {
+      await harness.close();
+    }
+  });
+});