fix(sdk,webapp): stop chat losing a message sent right after an action

Author: ericallam

.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.

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 });
   }
 );
 

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 });
 }

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<void> {
-    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<void> {
+  ): Promise<number> {
     return this.#appendPartByName(part, partId, this.toSessionStreamName(friendlyId, io));
   }
 
-  async #appendPartByName(part: string, partId: string, s2Stream: string): Promise<void> {
+  async #appendPartByName(part: string, partId: string, s2Stream: string): Promise<number> {
     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<number> {

docs/ai-chat/client-protocol.mdx

@@ -591,13 +591,17 @@ body: ""
 | --- | --- |
 | `trigger-control: turn-complete` | Always present on this record. |
 | `public-access-token: <jwt>` (optional) | A refreshed JWT with the same session + run scopes. If present, replace your stored token. |
-| `session-in-event-id: <seq>` (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: <seq>` (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`.
 
+<Warning>
+  **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.
+</Warning>
+
 ### `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": "<message>" }`. 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": <number> }`, 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": "<message>" }`. Common failures:
 
 | Status | When |
 | --- | --- |

packages/trigger-sdk/src/v3/chat.ts

@@ -730,11 +730,10 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
     // 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<UIMessage> {
     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<UIMessage> {
 
     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<UIMessage> {
     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<UIMessage> {
   }
 
   /** Run a send op, emitting the terminal message-sent / message-send-failed event. */
-  private async sendWithEvents(
+  private async sendWithEvents<T>(
     chatId: string,
     source: ChatTransportSendSource,
     extras: { messageId?: string; partId?: string; bodyBytes?: number },
-    op: () => Promise<void>
-  ): Promise<void> {
+    op: () => Promise<T>
+  ): Promise<T> {
     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<UIMessage> {
         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<UIMessage> {
     token: string,
     body: string,
     partId?: string
-  ): Promise<void> {
+  ): Promise<number | undefined> {
     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<UIMessage> {
       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<T>(
     chatId: string,
     state: ChatSessionState,
-    op: (token: string) => Promise<void>
-  ): Promise<void> {
+    op: (token: string) => Promise<T>
+  ): Promise<T> {
     // 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<UIMessage> {
     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<UIMessage> {
     //    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<UIMessage> {
       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<UIMessageChunk> {
     const internalAbort = new AbortController();
@@ -1750,6 +1754,15 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
             }
 
             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;

packages/trigger-sdk/test/chat-turn-correlation.test.ts

@@ -0,0 +1,118 @@
+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<unknown>): Promise<string[]> {
+  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<string[]> {
+  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([]);
+  });
+});