320 lines
13 KiB
TypeScript
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");
|
|
}
|
|
);
|
|
});
|