Files
2026-07-13 13:32:57 +08:00

320 lines
13 KiB
TypeScript

// Plan F.3: integration test for the crash-recovery boot path. The
// scenario it locks down:
//
// 1. Run A streams chunks to `session.out` and `onTurnComplete` fires.
// 2. Run A crashes BEFORE `writeChatSnapshot` lands the post-turn
// blob (or the write fails silently — both have the same effect).
// 3. Run B boots: `readChatSnapshot` returns `undefined` (no snapshot
// yet, or stale-from-prior-turn). Replay then drains
// `session.out` from the snapshot's `lastOutEventId` (or seq 0)
// and reduces the chunks back into UIMessage[].
// 4. The accumulator is consistent — Run A's completed chunks reach
// Run B's run loop without losing data.
//
// Plan section H.1 / H.4 spell out the "snapshot didn't make it before
// crash" path; this test is the integration safety net behind the
// unit tests in `packages/trigger-sdk/test/replay-session-out.test.ts`.
//
// We exercise the SDK's `__replaySessionOutTailProductionPathForTests`
// against a stubbed `apiClient.readSessionStreamRecords` — the new
// non-SSE records endpoint introduced in plan task #22. The replay path
// is a single GET that returns whatever's already on the stream; no
// long-poll. MinIO is provisioned to keep parity with
// `chat-snapshot-integration.test.ts` (the snapshot read path runs
// through it), even though the replay path itself doesn't read from S3.
import { postgresAndMinioTest } from "@internal/testcontainers";
import { apiClientManager } from "@trigger.dev/core/v3";
import {
__readChatSnapshotProductionPathForTests as readChatSnapshot,
__replaySessionOutTailProductionPathForTests as replaySessionOutTail,
type ChatSnapshotV1,
} from "@trigger.dev/sdk/ai";
import type { UIMessageChunk } from "ai";
import { afterEach, describe, expect, vi } from "vitest";
import { env } from "~/env.server";
import { chatSnapshotStoragePathForSession } from "~/services/realtime/chatSnapshot.server";
import { generatePresignedUrl } from "~/v3/objectStore.server";
vi.setConfig({ testTimeout: 60_000 });
// ── Helpers ────────────────────────────────────────────────────────────
function textTurn(id: string, text: string): UIMessageChunk[] {
return [
{ type: "start", messageId: id, messageMetadata: { role: "assistant" } } 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()` so:
* - `getPayloadUrl` / `createUploadPayloadUrl` mint MinIO presigned URLs
* via the webapp's real `generatePresignedUrl` (so snapshot reads
* hit a real S3-compatible backend).
* - `readSessionStreamRecords` returns the canonical
* `{ records: [{ data, id, seqNum }] }` shape. `data` is the parsed
* chunk OBJECT — the SDK writer puts the chunk object directly into
* the record envelope and the webapp route forwards it as-is, so
* the schema now declares `data: z.unknown()` and consumers use it
* without an extra `JSON.parse` step.
*/
function stubApiClient(opts: { projectRef: string; envSlug: string; sessionOutChunks: unknown[] }) {
const records = opts.sessionOutChunks.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({
async getChatSnapshotUrl(sessionId: string) {
const key = chatSnapshotStoragePathForSession(sessionId);
const result = await generatePresignedUrl(opts.projectRef, opts.envSlug, key, "GET");
if (!result.success) throw new Error(result.error);
return { presignedUrl: result.url };
},
async createChatSnapshotUploadUrl(sessionId: string) {
const key = chatSnapshotStoragePathForSession(sessionId);
const result = await generatePresignedUrl(opts.projectRef, opts.envSlug, key, "PUT");
if (!result.success) throw new Error(result.error);
return { presignedUrl: result.url };
},
readSessionStreamRecords: readRecordsSpy,
} as never);
return readRecordsSpy;
}
let warnSpy: ReturnType<typeof vi.spyOn>;
afterEach(() => {
vi.restoreAllMocks();
warnSpy?.mockRestore();
});
// ── Tests ──────────────────────────────────────────────────────────────
describe("replay after crash (MinIO + SDK helpers)", () => {
postgresAndMinioTest(
"boot reconstructs accumulator from session.out replay when no snapshot exists",
async ({ minioConfig }) => {
env.OBJECT_STORE_BASE_URL = minioConfig.baseUrl;
env.OBJECT_STORE_ACCESS_KEY_ID = minioConfig.accessKeyId;
env.OBJECT_STORE_SECRET_ACCESS_KEY = minioConfig.secretAccessKey;
env.OBJECT_STORE_REGION = minioConfig.region;
env.OBJECT_STORE_DEFAULT_PROTOCOL = undefined;
warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {});
// The crashed run's session.out: two completed assistant turns, no
// snapshot ever written. Boot must recover both via replay.
const chunks = [...textTurn("a-1", "first turn"), ...textTurn("a-2", "second turn")];
stubApiClient({
projectRef: "proj_replay_crash",
envSlug: "dev",
sessionOutChunks: chunks,
});
// Step 1: read snapshot — returns undefined (fresh boot, no snap).
const snapshot = await readChatSnapshot("sess_no_snap");
expect(snapshot).toBeUndefined();
// Step 2: replay tail.
const replayed = await replaySessionOutTail("sess_no_snap");
expect(replayed).toHaveLength(2);
expect(replayed.map((m) => m.id)).toEqual(["a-1", "a-2"]);
const texts = replayed.flatMap((m) =>
(m.parts as Array<{ type: string; text?: string }>)
.filter((p) => p.type === "text")
.map((p) => p.text)
);
expect(texts).toEqual(["first turn", "second turn"]);
}
);
postgresAndMinioTest(
"boot replays only chunks AFTER snapshot.lastOutEventId (resume cursor)",
async ({ minioConfig }) => {
env.OBJECT_STORE_BASE_URL = minioConfig.baseUrl;
env.OBJECT_STORE_ACCESS_KEY_ID = minioConfig.accessKeyId;
env.OBJECT_STORE_SECRET_ACCESS_KEY = minioConfig.secretAccessKey;
env.OBJECT_STORE_REGION = minioConfig.region;
env.OBJECT_STORE_DEFAULT_PROTOCOL = undefined;
// The replay helper accepts the snapshot's `lastEventId` cursor
// and forwards it as `afterEventId` on the records endpoint —
// that's the cursor field name on the new non-SSE route. Here we
// feed only the post-snapshot chunks (modeling what the server
// returns for `afterEventId=evt-snapped`) and verify the helper
// threads the cursor through.
const readRecordsSpy = stubApiClient({
projectRef: "proj_replay_resume",
envSlug: "dev",
sessionOutChunks: textTurn("a-after-snap", "post-snapshot turn"),
});
const result = await replaySessionOutTail("sess_resume", { lastEventId: "evt-snapped" });
expect(readRecordsSpy).toHaveBeenCalledWith(
"sess_resume",
"out",
expect.objectContaining({ afterEventId: "evt-snapped" })
);
expect(result).toHaveLength(1);
expect(result[0]!.id).toBe("a-after-snap");
}
);
postgresAndMinioTest(
"boot returns [] when session.out is empty (first-ever turn, no snapshot)",
async ({ minioConfig }) => {
env.OBJECT_STORE_BASE_URL = minioConfig.baseUrl;
env.OBJECT_STORE_ACCESS_KEY_ID = minioConfig.accessKeyId;
env.OBJECT_STORE_SECRET_ACCESS_KEY = minioConfig.secretAccessKey;
env.OBJECT_STORE_REGION = minioConfig.region;
env.OBJECT_STORE_DEFAULT_PROTOCOL = undefined;
warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {});
stubApiClient({
projectRef: "proj_replay_empty",
envSlug: "dev",
sessionOutChunks: [],
});
const snapshot = await readChatSnapshot("sess_empty");
expect(snapshot).toBeUndefined();
const replayed = await replaySessionOutTail("sess_empty");
expect(replayed).toEqual([]);
}
);
postgresAndMinioTest(
"boot drops orphaned trailing tool parts (cleanupAbortedParts) — partial crash",
async ({ minioConfig }) => {
// Simulates a true mid-turn crash: assistant finished one turn,
// then started a tool-call but the run died before resolution.
// Replay must surface the completed turn but NOT include the
// orphaned tool part in `input-streaming` state.
env.OBJECT_STORE_BASE_URL = minioConfig.baseUrl;
env.OBJECT_STORE_ACCESS_KEY_ID = minioConfig.accessKeyId;
env.OBJECT_STORE_SECRET_ACCESS_KEY = minioConfig.secretAccessKey;
env.OBJECT_STORE_REGION = minioConfig.region;
env.OBJECT_STORE_DEFAULT_PROTOCOL = undefined;
stubApiClient({
projectRef: "proj_replay_partial",
envSlug: "dev",
sessionOutChunks: [
...textTurn("a-complete", "I finished step 1"),
// Partial tool turn — no tool-input-end, no finish.
{
type: "start",
messageId: "a-orphan",
messageMetadata: { role: "assistant" },
} as UIMessageChunk,
{ type: "tool-input-start", id: "tc-cut", toolName: "search" } as UIMessageChunk,
{ type: "tool-input-delta", id: "tc-cut", delta: '{"q":"x"}' } as UIMessageChunk,
],
});
const replayed = await replaySessionOutTail("sess_partial_crash");
// Completed turn always present.
expect(replayed.find((m) => m.id === "a-complete")).toBeTruthy();
// Orphaned tool-call never surfaces in `input-streaming` state.
const orphan = replayed.find((m) => m.id === "a-orphan");
if (orphan) {
const stillStreaming = (
orphan.parts as Array<{ toolCallId?: string; state?: string }>
).find((p) => p.toolCallId === "tc-cut" && p.state === "input-streaming");
expect(stillStreaming).toBeUndefined();
}
}
);
postgresAndMinioTest(
"snapshot+replay merge: snapshot supplies user msgs, replay supplies assistants",
async ({ minioConfig }) => {
// The boot orchestration calls
// `mergeByIdReplaceWins(snapshot.messages, replayed)`. The runtime
// contract is that user messages live in snapshot only (session.in
// never goes through replay) and assistants come from replay
// (which carries the freshest representation). Here we simulate
// the realistic split: snapshot has [u-1, a-1-stale], replay has
// [a-1-fresh, a-2-new]. After merge the accumulator should reflect
// the fresh assistant + new assistant, with the user message
// preserved.
//
// Note: this is a pre-merge round-trip — we drive the read and
// replay through real MinIO + stubbed S2 to confirm both arrive
// intact for the orchestration to merge.
env.OBJECT_STORE_BASE_URL = minioConfig.baseUrl;
env.OBJECT_STORE_ACCESS_KEY_ID = minioConfig.accessKeyId;
env.OBJECT_STORE_SECRET_ACCESS_KEY = minioConfig.secretAccessKey;
env.OBJECT_STORE_REGION = minioConfig.region;
env.OBJECT_STORE_DEFAULT_PROTOCOL = undefined;
// Pre-write a snapshot to MinIO via real apiClient stub.
const sessionId = "sess_merge_round_trip";
const snapshot: ChatSnapshotV1 = {
version: 1,
savedAt: 1_700_000_000_000,
messages: [
{ id: "u-1", role: "user", parts: [{ type: "text", text: "hi" }] },
{ id: "a-1", role: "assistant", parts: [{ type: "text", text: "stale-assistant" }] },
],
lastOutEventId: "evt-prev",
};
// Use the SDK's own writer to lay the snapshot down, then swap
// the stub to also serve replay chunks for the read path.
stubApiClient({
projectRef: "proj_merge",
envSlug: "dev",
sessionOutChunks: [],
});
const { __writeChatSnapshotProductionPathForTests: writeSnapshot } =
await import("@trigger.dev/sdk/ai");
await writeSnapshot(sessionId, snapshot);
// Restubbing for the boot phase: replay tail carries the fresh
// assistant for `a-1` plus a brand-new `a-2`. The orchestration's
// merge would replace `a-1` and append `a-2` after `u-1`.
vi.restoreAllMocks();
stubApiClient({
projectRef: "proj_merge",
envSlug: "dev",
sessionOutChunks: [
...textTurn("a-1", "fresh-assistant"),
...textTurn("a-2", "next-assistant"),
],
});
const readBack = await readChatSnapshot(sessionId);
expect(readBack?.messages.map((m) => m.id)).toEqual(["u-1", "a-1"]);
const replayed = await replaySessionOutTail(sessionId, {
lastEventId: readBack?.lastOutEventId,
});
expect(replayed.map((m) => m.id)).toEqual(["a-1", "a-2"]);
// Replay's `a-1` carries the fresh content — when merge runs in
// the runtime, this version would replace the snapshot's stale
// `a-1`.
const replayedA1Text = (replayed[0]!.parts as Array<{ type: string; text?: string }>)
.filter((p) => p.type === "text")
.map((p) => p.text)
.join("");
expect(replayedA1Text).toBe("fresh-assistant");
}
);
});