// Import the test entry point first so the resource catalog is installed. import "../src/v3/test/index.js"; import type { UIMessageChunk } from "ai"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { apiClientManager } from "@trigger.dev/core/v3"; import { __replaySessionOutTailProductionPathForTests as replaySessionOutTail } from "../src/v3/ai.js"; // ── Helpers ──────────────────────────────────────────────────────────── /** * Build the canonical chunk sequence the AI SDK emits for a single text * turn from message `id`. Includes a trailing `finish` so the segment is * marked closed (i.e. NOT subject to `cleanupAbortedParts`). */ function textTurn(id: string, text: string, role: "assistant" = "assistant"): UIMessageChunk[] { return [ { type: "start", messageId: id, messageMetadata: { role } } as UIMessageChunk, { type: "text-start", id: `${id}.t1` } as UIMessageChunk, { type: "text-delta", id: `${id}.t1`, delta: text } as UIMessageChunk, { type: "text-end", id: `${id}.t1` } as UIMessageChunk, { type: "finish" } as UIMessageChunk, ]; } /** * Stub `apiClientManager.clientOrThrow().readSessionStreamRecords` so the * helper sees a `{ records: StreamRecord[] }` response. Each StreamRecord * is `{ data, id, seqNum }` — `data` is the parsed chunk OBJECT (the wire * writer puts chunks directly into the record envelope; the route * forwards them as-is; the schema declares `data: z.unknown()`). * * Pass either a chunk OBJECT (used as `data` directly) or a string * (used as `data` directly — for tests that need deliberately-malformed * bodies; the consumer filters non-objects out). * * Captures the `afterEventId` argument for resume-from-cursor assertions. */ function stubReadRecordsWithChunks(chunks: unknown[]) { const records = chunks.map((chunk, i) => ({ data: chunk, id: `evt-${i + 1}`, seqNum: i + 1, })); const readRecordsSpy = vi.fn( async (_id: string, _io: "in" | "out", _options?: { afterEventId?: string }) => ({ records, }) ); vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ readSessionStreamRecords: readRecordsSpy, } as never); return readRecordsSpy; } // ── Tests ────────────────────────────────────────────────────────────── describe("replaySessionOutTail", () => { let warnSpy: ReturnType; beforeEach(() => { warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); }); afterEach(() => { vi.restoreAllMocks(); warnSpy.mockRestore(); }); it("returns [] for an empty session.out stream", async () => { stubReadRecordsWithChunks([]); const result = await replaySessionOutTail("empty-session"); expect(result).toEqual([]); }); it("reduces a single text turn into one assistant UIMessage", async () => { stubReadRecordsWithChunks(textTurn("a-1", "hello world")); const result = await replaySessionOutTail("text-session"); expect(result).toHaveLength(1); expect(result[0]).toMatchObject({ id: "a-1", role: "assistant" }); const text = (result[0]!.parts as Array<{ type: string; text?: string }>) .filter((p) => p.type === "text") .map((p) => p.text) .join(""); expect(text).toBe("hello world"); }); it("reduces multiple sequential turns into multiple UIMessages", async () => { stubReadRecordsWithChunks([ ...textTurn("a-1", "first"), ...textTurn("a-2", "second"), ...textTurn("a-3", "third"), ]); const result = await replaySessionOutTail("multi-session"); expect(result).toHaveLength(3); expect(result.map((m) => m.id)).toEqual(["a-1", "a-2", "a-3"]); }); it("filters out `trigger:*` control chunks (turn-complete, etc.)", async () => { stubReadRecordsWithChunks([ ...textTurn("a-1", "hello"), { type: "trigger:turn-complete", lastEventId: "evt-1", lastEventTimestamp: 1 }, { type: "trigger:upgrade-required" }, ...textTurn("a-2", "second"), ]); const result = await replaySessionOutTail("control-session"); // Two assistant messages reduced — the trigger:* records are dropped // before reaching the reducer. expect(result).toHaveLength(2); expect(result.map((m) => m.id)).toEqual(["a-1", "a-2"]); }); it("never emits user-role messages (session.out is assistant-only)", async () => { // session.out conceptually only carries assistant chunks (the user's // messages live on session.in). Even if a user-role start somehow // landed there, the reducer wouldn't surface a user message via this // helper's contract. stubReadRecordsWithChunks(textTurn("a-1", "ok")); const result = await replaySessionOutTail("assistant-only"); expect(result.every((m) => m.role !== "user")).toBe(true); }); it("passes `lastEventId` through as `afterEventId` to readSessionStreamRecords", async () => { // The replay helper accepts `lastEventId` from the caller (matching // the snapshot's persisted cursor name) and forwards it as // `afterEventId` on the records endpoint — that's the field name on // the new non-SSE route. const readRecordsSpy = stubReadRecordsWithChunks(textTurn("a-1", "ok")); await replaySessionOutTail("resume-session", { lastEventId: "evt-99" }); expect(readRecordsSpy).toHaveBeenCalledWith( "resume-session", "out", expect.objectContaining({ afterEventId: "evt-99" }) ); }); it("uses the non-SSE records endpoint (drain-and-close, no long-poll)", async () => { // Replay no longer subscribes to the SSE stream — that imposed a ~1s // long-poll tax on every fresh chat boot. The new path hits // `readSessionStreamRecords` (one synchronous GET that returns // whatever's already in the stream) and returns immediately when // empty. Lock the call site down so a regression to SSE shows up // here. const readRecordsSpy = stubReadRecordsWithChunks([]); const result = await replaySessionOutTail("drain-session"); expect(readRecordsSpy).toHaveBeenCalledWith("drain-session", "out", expect.any(Object)); expect(result).toEqual([]); }); it("strips orphaned in-flight tool parts from a partial trailing assistant", async () => { // The runtime applies `cleanupAbortedParts` only on the trailing // segment when its closure flag is `false` (no `finish` chunk // received). The cleanup removes tool parts that never reached a // terminal state — `input-streaming`, `output-pending`, etc. — // because those represent partial in-flight work that won't resolve. // // Text parts with already-streamed content are preserved (the user // already saw them), so we test the tool-part path specifically. stubReadRecordsWithChunks([ ...textTurn("a-1", "previous-turn-finished"), // Trailing turn: starts a tool call but never resolves it. { type: "start", messageId: "a-2", messageMetadata: { role: "assistant" } } as UIMessageChunk, { type: "tool-input-start", toolCallId: "tc-cut", toolName: "search" } as UIMessageChunk, { type: "tool-input-delta", toolCallId: "tc-cut", inputTextDelta: '{"q":"x"}', } as UIMessageChunk, // No tool-input-end, no tool-call, no finish → orphaned. ]); const result = await replaySessionOutTail("partial-tool-session"); // The closed turn survives. expect(result.find((m) => m.id === "a-1")).toBeTruthy(); // Trailing message either gets dropped (cleanup empties it) or its // orphaned tool part is stripped to a terminal state. Either way, // no `tc-cut` part should be left in `input-streaming` state — that // would represent a tool the next turn would re-process. const trailing = result.find((m) => m.id === "a-2"); if (trailing) { const orphanedToolPart = ( trailing.parts as Array<{ type: string; toolCallId?: string; state?: string }> ).find((p) => p.toolCallId === "tc-cut" && p.state === "input-streaming"); expect(orphanedToolPart).toBeUndefined(); } }); it("drops a trailing message whose only parts are stripped by cleanup", async () => { // Trailing turn whose ONLY content is an orphaned tool — after // cleanup the message has no parts left, so the helper drops it // entirely (it never reached the next turn's accumulator). stubReadRecordsWithChunks([ ...textTurn("a-1", "complete"), { type: "start", messageId: "a-orphan", messageMetadata: { role: "assistant" }, } as UIMessageChunk, { type: "tool-input-start", toolCallId: "tc-orph", toolName: "search" } as UIMessageChunk, // No tool-input-end, no tool-call, no finish. ]); const result = await replaySessionOutTail("dropped-trailing"); expect(result).toHaveLength(1); expect(result[0]!.id).toBe("a-1"); }); it("preserves a complete trailing assistant (cleanup is a no-op)", async () => { // Trailing turn that DID end with `finish` is closed — cleanupAbortedParts // doesn't fire. Use this to lock down that closed segments survive // unchanged. stubReadRecordsWithChunks(textTurn("a-1", "fully-finished")); const result = await replaySessionOutTail("closed-session"); expect(result).toHaveLength(1); const text = (result[0]!.parts as Array<{ type: string; text?: string }>) .filter((p) => p.type === "text") .map((p) => p.text) .join(""); expect(text).toBe("fully-finished"); }); it("skips records whose data is a string (the wire delivers objects)", async () => { // The writer puts chunk objects directly into the record envelope; // the route forwards them as-is. A string body is malformed — the // consumer drops it defensively rather than JSON.parsing. stubReadRecordsWithChunks(["not-an-object", ...textTurn("a-1", "survived")]); const result = await replaySessionOutTail("garbage-session"); expect(result).toHaveLength(1); expect(result[0]!.id).toBe("a-1"); }); it("skips records whose data is not an object", async () => { // The consumer requires `chunk` to be a non-null object with a // string `type` field. Records that arrive as primitives // (number, null, string) are dropped silently. stubReadRecordsWithChunks([42, null, "just-a-string", ...textTurn("a-1", "survived")]); const result = await replaySessionOutTail("primitive-data-session"); expect(result).toHaveLength(1); expect(result[0]!.id).toBe("a-1"); }); it("ignores chunks missing a `type` field", async () => { stubReadRecordsWithChunks([{ foo: "bar" }, { type: 42 }, ...textTurn("a-1", "valid")]); const result = await replaySessionOutTail("typeless-session"); expect(result).toHaveLength(1); expect(result[0]!.id).toBe("a-1"); }); it("recovers from a malformed segment by skipping it (logs a warn)", async () => { // The reducer for one segment throws (e.g. invalid chunk sequence). // The helper logs the warning and proceeds with the next segment — // a single corrupt segment shouldn't sink the entire replay. stubReadRecordsWithChunks([ // Malformed: text-end with no preceding text-start. { type: "start", messageId: "bad-1", messageMetadata: { role: "assistant" }, } as UIMessageChunk, { type: "text-end", id: "no-such-text" } as UIMessageChunk, { type: "finish" } as UIMessageChunk, ...textTurn("a-1", "after-bad"), ]); const result = await replaySessionOutTail("recovery-session"); // The valid turn after the malformed one must still surface. expect(result.find((m) => m.id === "a-1")).toBeTruthy(); }); });