From 96e5ad3c172500a912c15dc1784b0bd6bb13df8f Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Sat, 11 Jul 2026 09:17:36 +0100 Subject: [PATCH] fix(sdk,webapp): stop chat losing a message sent right after an action --- .changeset/chat-turn-correlation.md | 5 + ...ealtime.v1.sessions.$session.$io.append.ts | 7 +- ...ealtime.v1.sessions.$session.$io.append.ts | 4 +- .../realtime/s2realtimeStreams.server.ts | 12 +- docs/ai-chat/client-protocol.mdx | 10 +- packages/trigger-sdk/src/v3/chat.ts | 65 ++++++----- .../test/chat-turn-correlation.test.ts | 110 ++++++++++++++++++ 7 files changed, 174 insertions(+), 39 deletions(-) create mode 100644 .changeset/chat-turn-correlation.md create mode 100644 packages/trigger-sdk/test/chat-turn-correlation.test.ts diff --git a/.changeset/chat-turn-correlation.md b/.changeset/chat-turn-correlation.md new file mode 100644 index 0000000000..a0d63913c0 --- /dev/null +++ b/.changeset/chat-turn-correlation.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/sdk": patch +--- + +Fix a `chat.agent` message-loss race where sending a message right after an action (such as an undo) could drop the follow-up's response from the UI until a refresh. diff --git a/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts b/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts index ddd4ab4ac5..9a3dec90d4 100644 --- a/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts +++ b/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts @@ -156,10 +156,12 @@ const { action, loader } = createActionApiRoute( ) : true; + let appendSeq: number | undefined; if (wonClaim) { - const [appendError] = await tryCatch( + const [appendError, seq] = await tryCatch( realtimeStream.appendPartToSessionStream(part, partId, addressingKey, params.io) ); + appendSeq = seq ?? undefined; if (appendError) { if (clientPartId) { @@ -228,7 +230,8 @@ const { action, loader } = createActionApiRoute( ); } - return json({ ok: true }, { status: 200 }); + // `seq` lets the client correlate this send to the turn that consumes it. + return json({ ok: true, seq: appendSeq }, { status: 200 }); } ); diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.realtime.v1.sessions.$session.$io.append.ts b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.realtime.v1.sessions.$session.$io.append.ts index 0d15865931..ab318f31c7 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.realtime.v1.sessions.$session.$io.append.ts +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.realtime.v1.sessions.$session.$io.append.ts @@ -101,7 +101,7 @@ export async function action({ request, params }: ActionFunctionArgs) { const part = await request.text(); const partId = request.headers.get("X-Part-Id") ?? nanoid(7); - const [appendError] = await tryCatch( + const [appendError, appendSeq] = await tryCatch( realtimeStream.appendPartToSessionStream(part, partId, addressingKey, io) ); @@ -148,5 +148,5 @@ export async function action({ request, params }: ActionFunctionArgs) { ); } - return json({ ok: true }, { status: 200 }); + return json({ ok: true, seq: appendSeq ?? undefined }, { status: 200 }); } diff --git a/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts b/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts index 37d97bc4f2..787b3f1ab4 100644 --- a/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts +++ b/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts @@ -194,22 +194,20 @@ export class S2RealtimeStreams implements StreamResponder, StreamIngestor { } async appendPart(part: string, partId: string, runId: string, streamId: string): Promise { - return this.#appendPartByName(part, partId, this.toStreamName(runId, streamId)); + await this.#appendPartByName(part, partId, this.toStreamName(runId, streamId)); } - /** - * Append a single record to a `Session`-primitive channel. - */ + /** Append one record to a `Session` channel; returns its seq (same space as `session-in-event-id`). */ async appendPartToSessionStream( part: string, partId: string, friendlyId: string, io: "out" | "in" - ): Promise { + ): Promise { return this.#appendPartByName(part, partId, this.toSessionStreamName(friendlyId, io)); } - async #appendPartByName(part: string, partId: string, s2Stream: string): Promise { + async #appendPartByName(part: string, partId: string, s2Stream: string): Promise { this.logger.debug(`S2 appending to stream`, { part, stream: s2Stream }); const recordBody = JSON.stringify({ data: part, id: partId }); @@ -223,6 +221,8 @@ export class S2RealtimeStreams implements StreamResponder, StreamIngestor { }); this.logger.debug(`S2 append result`, { result }); + + return result.start.seq_num; } getLastChunkIndex(runId: string, streamId: string, clientId: string): Promise { diff --git a/docs/ai-chat/client-protocol.mdx b/docs/ai-chat/client-protocol.mdx index f1c33ad6b2..fd8326e4ba 100644 --- a/docs/ai-chat/client-protocol.mdx +++ b/docs/ai-chat/client-protocol.mdx @@ -583,7 +583,7 @@ Signals that the agent's turn is finished — stop reading and wait for user inp headers: ["trigger-control", "turn-complete"] ["public-access-token", "eyJ..."] // optional, refreshed JWT - ["session-in-event-id", "42"] // optional, agent-internal resume cursor + ["session-in-event-id", "42"] // optional, send-correlation + resume cursor body: "" ``` @@ -591,13 +591,17 @@ body: "" | --- | --- | | `trigger-control: turn-complete` | Always present on this record. | | `public-access-token: ` (optional) | A refreshed JWT with the same session + run scopes. If present, replace your stored token. | -| `session-in-event-id: ` (optional) | Internal cursor used by the agent to resume `.in` across worker boots without replaying already-processed user messages. Custom transports should ignore this header — it carries no client-side meaning. | +| `session-in-event-id: ` (optional) | The agent's committed `.in` cursor for this turn: the seq of the last user record it had consumed when the turn finished. Compare it to the `seq` from your `.in/append` to tell whether this turn-complete is *yours* (see the warning below). Also used internally to resume `.in` across worker boots without replaying processed messages. | When you receive this record: 1. Update `publicAccessToken` if one is included on the headers. 2. Close the stream reader (unless you want to keep it open across turns — see [Resuming a stream](#resuming-a-stream)). 3. Wait for the next user message before sending on `.in`. + + **Correlate turn-completes to your send.** A `turn-complete` on `.out` can belong to an *earlier* turn than the message you just sent, for example a concurrent action (an undo) whose completion lands on the stream first. Only treat a `turn-complete` as terminal for your send if its `session-in-event-id` is greater than or equal to the `seq` returned by that send's `.in/append`; skip any lower one and keep reading. Closing on the wrong turn drops your response until the next reload. The built-in `TriggerChatTransport` does this correlation for you. + + ### `upgrade-required` control record Signals that the agent cannot handle this message on its current version and a new run has been started. Emitted when the agent calls [`chat.requestUpgrade()`](/ai-chat/patterns/version-upgrades). @@ -681,7 +685,7 @@ Content-Type: application/json `{sessionId}` accepts the same friendly-or-external forms as `.out`. The `publicAccessToken` from session-create authorizes both. -The body is a JSON-serialized [`ChatInputChunk`](#chatinputchunk) — a tagged union covering messages, stops, and actions. Send them as raw JSON strings (not wrapped in a `data` field). On success the response is `200 OK` with body `{ "ok": true }`; on failure it's `4xx`/`5xx` with `{ "ok": false, "error": "" }`. Common failures: +The body is a JSON-serialized [`ChatInputChunk`](#chatinputchunk), a tagged union covering messages, stops, and actions. Send them as raw JSON strings (not wrapped in a `data` field). On success the response is `200 OK` with body `{ "ok": true, "seq": }`, where `seq` is the appended record's `.in` sequence number. Use it to correlate this send to the turn that consumes it (see [`turn-complete` control record](#turn-complete-control-record)). On failure it's `4xx`/`5xx` with `{ "ok": false, "error": "" }`. Common failures: | Status | When | | --- | --- | diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index f625155189..2c6e369ba7 100644 --- a/packages/trigger-sdk/src/v3/chat.ts +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -730,11 +730,10 @@ export class TriggerChatTransport implements ChatTransport { // and the server-side dedupe sees one logical append. const partId = crypto.randomUUID(); const serializedBody = this.serializeInputChunk({ kind: "message", payload: wirePayload }); - const sendChatMessage = async (token: string) => { - await this.appendInputChunk(chatId, token, serializedBody, partId); - }; + const sendChatMessage = (token: string) => + this.appendInputChunk(chatId, token, serializedBody, partId); - await this.sendWithEvents( + const inSeq = await this.sendWithEvents( chatId, trigger, { @@ -755,7 +754,7 @@ export class TriggerChatTransport implements ChatTransport { state.isStreaming = true; this.notifySessionChange(chatId, state); - return this.subscribeToSessionStream(state, abortSignal, chatId); + return this.subscribeToSessionStream(state, abortSignal, chatId, { sinceInSeq: inSeq }); }; /** @@ -1124,12 +1123,13 @@ export class TriggerChatTransport implements ChatTransport { const body = this.serializeInputChunk({ kind: "message", payload: wirePayload }); const partId = crypto.randomUUID(); - const send = async (token: string) => { - await this.appendInputChunk(chatId, token, body, partId); - }; + const send = (token: string) => this.appendInputChunk(chatId, token, body, partId); - await this.sendWithEvents(chatId, "action", { partId, bodyBytes: byteLength(body) }, () => - this.callWithAuthRetry(chatId, state, send) + const inSeq = await this.sendWithEvents( + chatId, + "action", + { partId, bodyBytes: byteLength(body) }, + () => this.callWithAuthRetry(chatId, state, send) ); // Supersede any in-flight reader before subscribing — same as @@ -1147,7 +1147,7 @@ export class TriggerChatTransport implements ChatTransport { state.isStreaming = true; this.notifySessionChange(chatId, state); - return this.subscribeToSessionStream(state, undefined, chatId); + return this.subscribeToSessionStream(state, undefined, chatId, { sinceInSeq: inSeq }); }; // ------------------------------------------------------------------------- @@ -1191,15 +1191,15 @@ export class TriggerChatTransport implements ChatTransport { } /** Run a send op, emitting the terminal message-sent / message-send-failed event. */ - private async sendWithEvents( + private async sendWithEvents( chatId: string, source: ChatTransportSendSource, extras: { messageId?: string; partId?: string; bodyBytes?: number }, - op: () => Promise - ): Promise { + op: () => Promise + ): Promise { const startedAt = Date.now(); try { - await op(); + const result = await op(); const turnProducing = source === "submit-message" || source === "regenerate-message" || source === "head-start"; if (turnProducing) { @@ -1213,6 +1213,7 @@ export class TriggerChatTransport implements ChatTransport { durationMs: Date.now() - startedAt, ...extras, }); + return result; } catch (error) { const status = (error as { status?: unknown }).status; this.emitEvent({ @@ -1384,7 +1385,7 @@ export class TriggerChatTransport implements ChatTransport { token: string, body: string, partId?: string - ): Promise { + ): Promise { const ctx: ChatTransportEndpointContext = { endpoint: "in", chatId }; const url = `${this.resolveBaseURL(ctx)}/realtime/v1/sessions/${encodeURIComponent(chatId)}/in/append`; // extraHeaders first so the fixed headers below win — a transport-wide @@ -1407,24 +1408,26 @@ export class TriggerChatTransport implements ChatTransport { err.status = response.status; throw err; } + // The appended record's `.in` seq, for correlating the response stream to + // this send. Omitted by older webapps / a lost idempotency claim. + const data = (await response.json().catch(() => undefined)) as { seq?: unknown } | undefined; + return typeof data?.seq === "number" ? data.seq : undefined; } - private async callWithAuthRetry( + private async callWithAuthRetry( chatId: string, state: ChatSessionState, - op: (token: string) => Promise - ): Promise { + op: (token: string) => Promise + ): Promise { // 1) Try with the current PAT. try { - await op(state.publicAccessToken); - return; + return await op(state.publicAccessToken); } catch (err) { if (isSessionNotFoundError(err)) { // The cached PAT authenticated but the session doesn't exist here — // recreate it and retry. await this.recreateSession(chatId, state); - await op(state.publicAccessToken); - return; + return await op(state.publicAccessToken); } if (!isAuthError(err)) throw err; } @@ -1435,8 +1438,7 @@ export class TriggerChatTransport implements ChatTransport { state.publicAccessToken = fresh; this.notifySessionChange(chatId, state); try { - await op(fresh); - return; + return await op(fresh); } catch (err) { if (!isSessionNotFoundError(err)) throw err; } @@ -1445,7 +1447,7 @@ export class TriggerChatTransport implements ChatTransport { // state is stale (created in a different environment, or before the // sessions upgrade). Recreate the session and retry once. await this.recreateSession(chatId, state); - await op(state.publicAccessToken); + return await op(state.publicAccessToken); } /** @@ -1493,6 +1495,8 @@ export class TriggerChatTransport implements ChatTransport { sendStopOnAbort?: boolean; peekSettled?: boolean; resumed?: boolean; + /** `.in` seq of the send that opened this stream; skip turn-completes below it (earlier turns). */ + sinceInSeq?: number; } ): ReadableStream { const internalAbort = new AbortController(); @@ -1750,6 +1754,15 @@ export class TriggerChatTransport implements ChatTransport { } if (controlValue === TRIGGER_CONTROL_SUBTYPE.TURN_COMPLETE) { + // Skip a turn-complete from an earlier turn (committed `.in` cursor + // below this send's seq), e.g. an undo action that raced this send. + if (options?.sinceInSeq !== undefined) { + const cursorRaw = headerValue(value.headers, SESSION_IN_EVENT_ID_HEADER); + const cursor = cursorRaw !== undefined ? Number.parseInt(cursorRaw, 10) : NaN; + if (!Number.isNaN(cursor) && cursor < options.sinceInSeq) { + continue; + } + } const refreshedToken = headerValue(value.headers, PUBLIC_ACCESS_TOKEN_HEADER) ?? legacyChunk?.publicAccessToken; diff --git a/packages/trigger-sdk/test/chat-turn-correlation.test.ts b/packages/trigger-sdk/test/chat-turn-correlation.test.ts new file mode 100644 index 0000000000..7295ecba79 --- /dev/null +++ b/packages/trigger-sdk/test/chat-turn-correlation.test.ts @@ -0,0 +1,110 @@ +import { describe, expect, it } from "vitest"; +import type { UIMessage } from "ai"; +import { TriggerChatTransport, type TriggerChatTransportOptions } from "../src/v3/chat.js"; + +// A send's `.out` stream must close on the turn that consumed its own appended +// record, not an earlier turn-complete (e.g. a racing undo action). The seq +// comes back from `/in/append`; correlation headers ride the v2 batch wire. + +function user(text: string, id: string): UIMessage { + return { id, role: "user", parts: [{ type: "text", text }] }; +} + +type BatchRecord = { + body: string; + seq_num: number; + timestamp: number; + headers?: Array<[string, string]>; +}; + +function batchResponse(records: BatchRecord[]): Response { + const frames = records + .map((r) => `event: batch\ndata: ${JSON.stringify({ records: [r] })}\n\n`) + .join(""); + return new Response(frames, { + status: 200, + headers: { "Content-Type": "text/event-stream", "X-Stream-Version": "v2" }, + }); +} + +/** A turn-complete control record whose committed `.in` cursor is `inCursor`. */ +function turnComplete(seqNum: number, inCursor: number): BatchRecord { + return { + body: "", + seq_num: seqNum, + timestamp: seqNum, + headers: [ + ["trigger-control", "turn-complete"], + ["session-in-event-id", String(inCursor)], + ], + }; +} + +function textDelta(seqNum: number, text: string): BatchRecord { + return { + body: JSON.stringify({ data: { type: "text-delta", id: "t1", delta: text }, id: "m1" }), + seq_num: seqNum, + timestamp: seqNum, + headers: [], + }; +} + +function inResponse(seq?: number): Response { + return new Response(JSON.stringify(seq === undefined ? { ok: true } : { ok: true, seq }), { + status: 200, + }); +} + +async function readDeltas(stream: ReadableStream): Promise { + const out: string[] = []; + const reader = stream.getReader(); + while (true) { + const next = await reader.read(); + if (next.done) return out; + const chunk = next.value as { type?: string; delta?: string }; + if (chunk?.type === "text-delta" && typeof chunk.delta === "string") out.push(chunk.delta); + } +} + +function makeTransport(out: Response, inSeq: number | undefined) { + const options: TriggerChatTransportOptions = { + task: "test-task", + accessToken: async () => "tok_test", + sessions: { c1: { publicAccessToken: "tok_test", isStreaming: false } }, + fetch: async (_url, _init, ctx) => (ctx.endpoint === "in" ? inResponse(inSeq) : out), + }; + return new TriggerChatTransport(options); +} + +async function submit(transport: TriggerChatTransport): Promise { + const stream = await transport.sendMessages({ + trigger: "submit-message", + chatId: "c1", + messageId: undefined, + messages: [user("hi", "u-1")], + abortSignal: undefined, + }); + return readDeltas(stream); +} + +describe("transport turn correlation", () => { + it("skips an earlier turn's turn-complete and closes on its own", async () => { + // Append seq 5; the undo turn's complete (cursor 4) must be skipped. + const out = batchResponse([turnComplete(10, 4), textDelta(11, "56"), turnComplete(12, 5)]); + const deltas = await submit(makeTransport(out, 5)); + expect(deltas).toEqual(["56"]); + }); + + it("does not skip when the turn-complete is at the send's own seq", async () => { + const out = batchResponse([textDelta(10, "56"), turnComplete(11, 5)]); + const deltas = await submit(makeTransport(out, 5)); + expect(deltas).toEqual(["56"]); + }); + + it("without an append seq, closes on the first turn-complete (legacy webapp)", async () => { + // No seq => no baseline => old behavior: close on the first turn-complete. + const out = batchResponse([turnComplete(10, 4), textDelta(11, "56"), turnComplete(12, 5)]); + const deltas = await submit(makeTransport(out, undefined)); + expect(deltas).toEqual([]); + }); +});