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: sending a message immediately after an action (for example an undo) could close the response stream on the earlier turn's completion, so the follow-up's answer never appeared on screen (it came back on refresh). The transport now correlates each send to the turn that consumes it, so the stream only closes on its own turn.

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

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([]);
+  });
+});