diff --git a/apps/webapp/app/components/runs/v3/agent/AgentView.tsx b/apps/webapp/app/components/runs/v3/agent/AgentView.tsx index f2570e3b928..f8c72d4db9b 100644 --- a/apps/webapp/app/components/runs/v3/agent/AgentView.tsx +++ b/apps/webapp/app/components/runs/v3/agent/AgentView.tsx @@ -1,9 +1,10 @@ import type { UIMessage } from "@ai-sdk/react"; -import { ChatSnapshotV1Schema, SSEStreamSubscription } from "@trigger.dev/core/v3"; +import { SSEStreamSubscription } from "@trigger.dev/core/v3"; import { useEffect, useMemo, useRef, useState } from "react"; import { Paragraph } from "~/components/primitives/Paragraph"; import { Spinner } from "~/components/primitives/Spinner"; import { AgentMessageView } from "~/components/runs/v3/agent/AgentMessageView"; +import { seedFromTranscriptSnapshot } from "~/components/runs/v3/agent/transcriptSnapshotSeed"; import { useAutoScrollToBottom } from "~/hooks/useAutoScrollToBottom"; import { useEnvironment } from "~/hooks/useEnvironment"; import { useOrganization } from "~/hooks/useOrganizations"; @@ -392,30 +393,17 @@ function useAgentSessionMessages({ const resp = await fetch(url, { signal: abort.signal }); if (!resp.ok) return undefined; const json = (await resp.json()) as unknown; - const parsed = ChatSnapshotV1Schema.safeParse(json); - if (!parsed.success) return undefined; - const snapshot = parsed.data; - // Preserve the snapshot's array order in the final render by - // giving each message a unique, monotonically increasing - // timestamp from `(savedAt - count + index)`. Real chunk - // timestamps from the SSE path use S2 arrival ms (positive - // numbers in the present), so anything below `savedAt` sorts - // before live chunks while preserving snapshot order among - // themselves. - const count = snapshot.messages.length; - snapshot.messages.forEach((raw, i) => { - const message = raw as UIMessage; - if (!message?.id) return; + const seed = seedFromTranscriptSnapshot(json); + if (!seed) return undefined; + for (const { id, message, timestamp } of seed.messages) { // The snapshot's seed wins over the task-payload seed for any // overlapping ids (the snapshot represents the agent's // canonical accumulator, post-turn). - pendingRef.current.set(message.id, message); - if (!timestampsRef.current.has(message.id)) { - timestampsRef.current.set(message.id, snapshot.savedAt - count + i); - } - }); + pendingRef.current.set(id, message); + timestampsRef.current.set(id, timestamp); + } scheduleFlush.current(); - return snapshot.lastOutEventId; + return seed.lastOutEventId; } catch { // 404 / network / parse / abort — fall back to seq=0 SSE return undefined; diff --git a/apps/webapp/app/components/runs/v3/agent/transcriptSnapshotSeed.test.ts b/apps/webapp/app/components/runs/v3/agent/transcriptSnapshotSeed.test.ts new file mode 100644 index 00000000000..db0a36524ec --- /dev/null +++ b/apps/webapp/app/components/runs/v3/agent/transcriptSnapshotSeed.test.ts @@ -0,0 +1,62 @@ +import { describe, expect, it } from "vitest"; +import { seedFromTranscriptSnapshot } from "./transcriptSnapshotSeed"; + +const user = { id: "u-1", role: "user", parts: [{ type: "text", text: "hello" }] }; +const assistant = { id: "a-1", role: "assistant", parts: [{ type: "text", text: "world" }] }; + +describe("seedFromTranscriptSnapshot", () => { + it("seeds from a version 1 snapshot in array order", () => { + const seed = seedFromTranscriptSnapshot({ + version: 1, + savedAt: 1_000, + messages: [user, assistant], + lastOutEventId: "42", + lastInEventId: "7", + }); + + expect(seed).toBeDefined(); + expect(seed!.lastOutEventId).toBe("42"); + expect(seed!.messages.map((m) => m.id)).toEqual(["u-1", "a-1"]); + expect(seed!.messages.map((m) => m.timestamp)).toEqual([998, 999]); + expect(seed!.messages[1]!.message).toEqual(assistant); + }); + + it("seeds from a version 2 snapshot, unwrapping the message envelope", () => { + const seed = seedFromTranscriptSnapshot({ + version: 2, + savedAt: 1_000, + messages: [ + { id: "u-1", final: true, message: user }, + { id: "a-1", final: false, message: assistant }, + ], + state: { summary: "irrelevant to rendering" }, + lastOutEventId: "42", + lastInEventId: "7", + }); + + expect(seed).toBeDefined(); + expect(seed!.lastOutEventId).toBe("42"); + expect(seed!.messages.map((m) => m.id)).toEqual(["u-1", "a-1"]); + expect(seed!.messages.map((m) => m.timestamp)).toEqual([998, 999]); + expect(seed!.messages[1]!.message).toEqual(assistant); + }); + + it("skips version 1 entries without an id but keeps the others' positions", () => { + const seed = seedFromTranscriptSnapshot({ + version: 1, + savedAt: 1_000, + messages: [{ role: "user", parts: [] }, assistant], + }); + + expect(seed!.messages.map((m) => m.id)).toEqual(["a-1"]); + expect(seed!.messages[0]!.timestamp).toBe(999); + expect(seed!.lastOutEventId).toBeUndefined(); + }); + + it("returns undefined for an unknown version or a non-snapshot body", () => { + expect(seedFromTranscriptSnapshot({ version: 3, savedAt: 1, messages: [] })).toBeUndefined(); + expect(seedFromTranscriptSnapshot({ error: "not found" })).toBeUndefined(); + expect(seedFromTranscriptSnapshot(null)).toBeUndefined(); + expect(seedFromTranscriptSnapshot("[]")).toBeUndefined(); + }); +}); diff --git a/apps/webapp/app/components/runs/v3/agent/transcriptSnapshotSeed.ts b/apps/webapp/app/components/runs/v3/agent/transcriptSnapshotSeed.ts new file mode 100644 index 00000000000..8a7302eebbb --- /dev/null +++ b/apps/webapp/app/components/runs/v3/agent/transcriptSnapshotSeed.ts @@ -0,0 +1,32 @@ +import type { UIMessage } from "@ai-sdk/react"; +import { parseTranscriptSnapshot } from "@trigger.dev/core/v3"; + +export type TranscriptSnapshotSeed = { + messages: Array<{ id: string; message: UIMessage; timestamp: number }>; + lastOutEventId: string | undefined; +}; + +/** + * Turn a fetched chat-snapshot blob into the messages the AgentView seeds + * before it opens the `.out` subscription. + * + * Each message gets a unique, monotonically increasing timestamp from + * `(savedAt - count + index)`. Live chunk timestamps are S2 arrival + * milliseconds in the present, so anything below `savedAt` sorts before + * live chunks while preserving the snapshot's own order. + * + * Reads both snapshot versions through `parseTranscriptSnapshot`. Returns + * `undefined` for anything that is not a snapshot this reader understands; + * the caller then falls back to the seq=0 SSE. + */ +export function seedFromTranscriptSnapshot(json: unknown): TranscriptSnapshotSeed | undefined { + const snapshot = parseTranscriptSnapshot(json); + if (!snapshot) return undefined; + const count = snapshot.messages.length; + const messages = snapshot.messages.map((entry, i) => ({ + id: entry.id, + message: entry.message, + timestamp: snapshot.savedAt - count + i, + })); + return { messages, lastOutEventId: snapshot.lastOutEventId }; +} diff --git a/packages/core/src/v3/sessionStreams/chatSnapshot.ts b/packages/core/src/v3/sessionStreams/chatSnapshot.ts index 84b204a134d..fba0c3942cf 100644 --- a/packages/core/src/v3/sessionStreams/chatSnapshot.ts +++ b/packages/core/src/v3/sessionStreams/chatSnapshot.ts @@ -50,6 +50,98 @@ export const ChatSnapshotV1Schema = z.object({ lastInEventId: z.string().optional(), }); +/** + * One transcript entry in a version 2 snapshot. `id` duplicates + * `message.id` so a reader can address entries without inspecting the + * message body; `final` is false for a partial assistant message captured + * by an errored or stopped turn. + */ +export type TranscriptSnapshotEntry = { + id: string; + final: boolean; + message: TUIMessage; +}; + +/** + * Version 2 of the persisted transcript blob. Entries are ordered by array + * position. `state` is an opaque record the runtime uses for compaction and + * other cross-run bookkeeping; `null` when nothing has been recorded. + * + * Readers must accept version 1 as well; writers only emit version 2. Use + * {@link parseTranscriptSnapshot} to read either. + */ +export type TranscriptSnapshotV2 = { + version: 2; + savedAt: number; + messages: TranscriptSnapshotEntry[]; + state: unknown | null; + lastOutEventId?: string; + lastInEventId?: string; +}; + +export const TranscriptSnapshotV2Schema = z.object({ + version: z.literal(2), + savedAt: z.number(), + messages: z.array( + z.object({ + id: z.string(), + final: z.boolean(), + message: z.unknown(), + }) + ), + state: z.unknown().nullable(), + lastOutEventId: z.string().optional(), + lastInEventId: z.string().optional(), +}); + +/** + * Parse a fetched snapshot blob of any known version into the version 2 + * shape. A version 1 blob is upgraded in memory: every message becomes a + * `final: true` entry keyed by its `id` (entries without a string `id` are + * dropped) and `state` is `null`. Returns `undefined` for an unknown + * version or a body that is not a snapshot; callers treat that as + * "no snapshot". + */ +export function parseTranscriptSnapshot( + input: unknown +): TranscriptSnapshotV2 | undefined { + const v2 = TranscriptSnapshotV2Schema.safeParse(input); + if (v2.success) { + return { + version: 2, + savedAt: v2.data.savedAt, + messages: v2.data.messages.map((entry) => ({ + id: entry.id, + final: entry.final, + message: entry.message as TUIMessage, + })), + state: v2.data.state ?? null, + lastOutEventId: v2.data.lastOutEventId, + lastInEventId: v2.data.lastInEventId, + }; + } + + const v1 = ChatSnapshotV1Schema.safeParse(input); + if (v1.success) { + const messages: TranscriptSnapshotEntry[] = []; + for (const raw of v1.data.messages) { + const id = (raw as { id?: unknown } | null)?.id; + if (typeof id !== "string" || id.length === 0) continue; + messages.push({ id, final: true, message: raw as TUIMessage }); + } + return { + version: 2, + savedAt: v1.data.savedAt, + messages, + state: null, + lastOutEventId: v1.data.lastOutEventId, + lastInEventId: v1.data.lastInEventId, + }; + } + + return undefined; +} + /** * S3 key suffix for a session's snapshot blob. The webapp's presigned * URL routes prefix this with `packets/{projectRef}/{envSlug}/`. diff --git a/packages/core/test/chatSnapshot.test.ts b/packages/core/test/chatSnapshot.test.ts new file mode 100644 index 00000000000..1b5ef58b1ff --- /dev/null +++ b/packages/core/test/chatSnapshot.test.ts @@ -0,0 +1,85 @@ +import { describe, expect, it } from "vitest"; +import { + parseTranscriptSnapshot, + type ChatSnapshotV1, + type TranscriptSnapshotV2, +} from "../src/v3/sessionStreams/chatSnapshot.js"; + +const user = { id: "u-1", role: "user" as const, parts: [{ type: "text" as const, text: "hi" }] }; +const assistant = { + id: "a-1", + role: "assistant" as const, + parts: [{ type: "text" as const, text: "hello" }], +}; + +describe("parseTranscriptSnapshot", () => { + it("returns a version 2 blob unchanged", () => { + const blob: TranscriptSnapshotV2 = { + version: 2, + savedAt: 10, + messages: [ + { id: "u-1", final: true, message: user }, + { id: "a-1", final: false, message: assistant }, + ], + state: { summary: "s", through: "u-1" }, + lastOutEventId: "9", + lastInEventId: "3", + }; + + expect(parseTranscriptSnapshot(blob)).toEqual(blob); + }); + + it("normalises a version 2 blob with no state to state: null", () => { + const parsed = parseTranscriptSnapshot({ + version: 2, + savedAt: 10, + messages: [], + state: undefined, + }); + + expect(parsed?.state).toBeNull(); + }); + + it("upgrades a version 1 blob: every message final, state null, cursors kept", () => { + const blob: ChatSnapshotV1 = { + version: 1, + savedAt: 10, + messages: [user, assistant], + lastOutEventId: "9", + lastInEventId: "3", + }; + + expect(parseTranscriptSnapshot(blob)).toEqual({ + version: 2, + savedAt: 10, + messages: [ + { id: "u-1", final: true, message: user }, + { id: "a-1", final: true, message: assistant }, + ], + state: null, + lastOutEventId: "9", + lastInEventId: "3", + }); + }); + + it("drops version 1 messages that have no string id", () => { + const parsed = parseTranscriptSnapshot({ + version: 1, + savedAt: 10, + messages: [{ role: "user", parts: [] }, { id: 7, role: "user", parts: [] }, assistant, null], + }); + + expect(parsed?.messages.map((m) => m.id)).toEqual(["a-1"]); + }); + + it("returns undefined for unknown versions and non-snapshot bodies", () => { + expect(parseTranscriptSnapshot({ version: 3, savedAt: 1, messages: [] })).toBeUndefined(); + expect( + parseTranscriptSnapshot({ version: 2, savedAt: 1, messages: [{ id: "x" }] }) + ).toBeUndefined(); + expect(parseTranscriptSnapshot({ version: 1, savedAt: "1", messages: [] })).toBeUndefined(); + expect(parseTranscriptSnapshot({ message: "Not Found" })).toBeUndefined(); + expect(parseTranscriptSnapshot(undefined)).toBeUndefined(); + expect(parseTranscriptSnapshot([])).toBeUndefined(); + }); +});