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

## Summary Sending a chat message immediately after an action (for example an undo) could make the message's response vanish from the UI. The transport opened a response stream that closed on the *earlier* turn's completion instead of waiting for the send's own turn. The agent still produced and persisted the answer, so it reappeared on refresh. Same "disappearing message" class as [#4176](https://github.com/triggerdotdev/trigger.dev/pull/4176), different cause. ## Fix A send's response stream had no way to tell whether a `turn-complete` belonged to its turn. `POST /realtime/v1/sessions/:id/in/append` now returns the appended record's sequence number, and the transport skips any turn-complete whose `session-in-event-id` (the agent's committed `.in` cursor) is below that seq, closing only on its own turn. Older webapps omit the seq, in which case the transport falls back to the previous behavior, so the SDK and server can ship independently. Because the fix spans the SDK and the server, both a webapp deploy and an SDK release are needed for the full effect. Verified end to end with the ai-chat reference app: undo-then-immediate-send loses the follow-up's answer before the fix and streams it inline after, with a revert-the-guard run reproducing the loss on the same script. Unit tests cover the skip and the no-seq fallback.
2026-07-11 12:46:24 +01:00
parent 5d0e9d9dc5
commit 9b3a7bd7b2
7 changed files with 174 additions and 39 deletions
@@ -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.
@@ -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 });
  }
 );

@@ -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 });
 }
@@ -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> {
@@ -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: <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 |
 | --- | --- |
@@ -849,11 +849,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,
      {
@@ -874,7 +873,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 });
  };

  /**
@@ -1243,12 +1242,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
@@ -1266,7 +1266,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 });
  };

  // -------------------------------------------------------------------------
@@ -1310,15 +1310,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) {
@@ -1332,6 +1332,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({
@@ -1503,7 +1504,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
@@ -1526,24 +1527,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;
    }
@@ -1554,8 +1557,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;
    }
@@ -1564,7 +1566,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);
  }

  /**
@@ -1612,6 +1614,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();
@@ -1869,6 +1873,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;
@@ -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<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([]);
+  });
+});