fix(chat): cover head-start and watch reconnects

Author: gtremper

packages/core/src/v3/schemas/api.ts

@@ -2181,6 +2181,7 @@ export type CreateStreamResponseBody = z.infer<typeof CreateStreamResponseBody>;
 export const AppendToStreamResponseBody = z.object({
   ok: z.boolean(),
   message: z.string().optional(),
+  seq: z.number().optional(),
 });
 export type AppendToStreamResponseBody = z.infer<typeof AppendToStreamResponseBody>;
 

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

@@ -119,8 +119,8 @@ function createSessionResponse(externalId: string): Response {
   );
 }
 
-function appendOkResponse(): Response {
-  return new Response(JSON.stringify({ ok: true }), {
+function appendOkResponse(seq?: number): Response {
+  return new Response(JSON.stringify(seq === undefined ? { ok: true } : { ok: true, seq }), {
     status: 200,
     headers: { "content-type": "application/json" },
   });
@@ -289,7 +289,7 @@ describe("chat.headStart (route handler)", () => {
         return createSessionResponse("chat-final");
       }
       if (urlStr.includes("/realtime/v1/sessions/") && urlStr.endsWith("/in/append")) {
-        return appendOkResponse();
+        return appendOkResponse(17);
       }
       // Stitched response subscribes to `.out` after handover.
       if (/\/realtime\/v1\/sessions\/[^/]+\/out$/.test(urlStr)) {
@@ -336,6 +336,9 @@ describe("chat.headStart (route handler)", () => {
     // Drain the SSE body so handoverWhenDone observes finishReason.
     const chunks = await readSSEBodyToChunks(res);
     expect(chunks.some((c) => c.type === "text-delta")).toBe(true);
+    expect(chunks).toContainEqual(
+      expect.objectContaining({ type: "trigger:session-state", activeInputSeq: 17 })
+    );
 
     // Give the deferred handoverWhenDone a tick to dispatch.
     await new Promise((r) => setTimeout(r, 30));

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

@@ -662,7 +662,7 @@ async function openHandoverSession(opts: {
     }
   };
 
-  const handover = async (args: {
+  const dispatchHandover = async (args: {
     partialAssistantMessage: ModelMessage[];
     messageId?: string;
     isFinal: boolean;
@@ -673,7 +673,16 @@ async function openHandoverSession(opts: {
       messageId: args.messageId,
       isFinal: args.isFinal,
     };
-    await apiClient.appendToSessionStream(chatId, "in", JSON.stringify(chunk));
+    const result = await apiClient.appendToSessionStream(chatId, "in", JSON.stringify(chunk));
+    return result.seq;
+  };
+
+  const handover = async (args: {
+    partialAssistantMessage: ModelMessage[];
+    messageId?: string;
+    isFinal: boolean;
+  }) => {
+    await dispatchHandover(args);
   };
 
   /**
@@ -707,7 +716,7 @@ async function openHandoverSession(opts: {
   // and dispatches the handover decision. The stitched response stream
   // awaits this to know whether to close (skip) or pull more chunks
   // from session.out (handover).
-  type HandoverDecision = { kind: "handover" | "handover-skip" };
+  type HandoverDecision = { kind: "handover"; activeInputSeq?: number } | { kind: "handover-skip" };
   let resolveDecision!: (decision: HandoverDecision) => void;
   const decisionPromise = new Promise<HandoverDecision>((resolve) => {
     resolveDecision = resolve;
@@ -737,25 +746,26 @@ async function openHandoverSession(opts: {
         // so the agent's `streamText` resumes by executing them
         // before the step-2 LLM call.
         const reshaped = reshapeForHandoverResume(responseMessages);
-        await handover({
+        const activeInputSeq = await dispatchHandover({
           partialAssistantMessage: reshaped,
           messageId: turnMessageId,
           isFinal: false,
         });
+        resolveDecision({ kind: "handover", activeInputSeq });
       } else {
         // Pure-text (or any non-tool-calls) finish — customer's step 1
         // IS the final response. The agent runs the turn-loop hooks
         // (`onChatStart`, `onTurnStart`, `onTurnComplete`, etc.) using
         // this partial as the response, but skips the LLM call. That
         // way persistence (`onTurnComplete` writing to DB), self-
         // review, and any post-turn work all fire normally.
-        await handover({
+        const activeInputSeq = await dispatchHandover({
           partialAssistantMessage: responseMessages,
           messageId: turnMessageId,
           isFinal: true,
         });
+        resolveDecision({ kind: "handover", activeInputSeq });
       }
-      resolveDecision({ kind: "handover" });
     } catch (err) {
       // Dispatch failed before we could send the handover signal.
       // Tell the agent to exit clean (no hooks fire) and close the
@@ -812,6 +822,18 @@ async function openHandoverSession(opts: {
             return;
           }
 
+          // The handover append happens after the HTTP response has started,
+          // so its `.in` sequence cannot be returned in a response header.
+          // Send it as transport-only state before any agent output so a
+          // reload during the remainder of this turn can reject stale
+          // turn-complete records.
+          if (decision.activeInputSeq !== undefined) {
+            controller.enqueue({
+              type: "trigger:session-state",
+              activeInputSeq: decision.activeInputSeq,
+            } as unknown as UIMessageChunk);
+          }
+
           // Phase 2b: agent is taking over. Resume from session.out
           // starting AFTER the customer tee's last write, so we don't
           // re-emit chunks the browser already saw.

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

@@ -35,6 +35,10 @@ function sseEncode(chunks: (UIMessageChunk | Record<string, unknown>)[]): string
       const headers: Array<[string, string]> = [["trigger-control", "turn-complete"]];
       const token = (chunk as { publicAccessToken?: string }).publicAccessToken;
       if (token) headers.push(["public-access-token", token]);
+      const sessionInEventId = (chunk as { sessionInEventId?: string | number }).sessionInEventId;
+      if (sessionInEventId !== undefined) {
+        headers.push(["session-in-event-id", String(sessionInEventId)]);
+      }
       return {
         body: "",
         seq_num: nextSeq++,
@@ -1762,24 +1766,28 @@ describe("TriggerChatTransport", () => {
      * does: `data: <JSON>\n\n` per chunk. The transport's
      * `parseUIMessageSseTransform` parses this back into chunk objects.
      */
-    function handoverSseBody(chunks: UIMessageChunk[]): ReadableStream<Uint8Array> {
+    function handoverSseBody(
+      chunks: (UIMessageChunk | Record<string, unknown>)[],
+      close = true
+    ): ReadableStream<Uint8Array> {
       const encoder = new TextEncoder();
       return new ReadableStream({
         start(controller) {
           for (const chunk of chunks) {
             controller.enqueue(encoder.encode(`data: ${JSON.stringify(chunk)}\n\n`));
           }
-          controller.close();
+          if (close) controller.close();
         },
       });
     }
 
     function handoverResponse(args: {
       chatId: string;
       accessToken: string;
-      chunks: UIMessageChunk[];
+      chunks: (UIMessageChunk | Record<string, unknown>)[];
+      close?: boolean;
     }): Response {
-      return new Response(handoverSseBody(args.chunks), {
+      return new Response(handoverSseBody(args.chunks, args.close), {
         status: 200,
         headers: {
           "content-type": "text/event-stream",
@@ -1914,6 +1922,70 @@ describe("TriggerChatTransport", () => {
       expect(subscribe).toBeDefined();
     });
 
+    it("persists the handover input sequence for a reload while the first turn is active", async () => {
+      global.fetch = vi.fn().mockImplementation(async (url: string | URL) => {
+        const urlStr = typeof url === "string" ? url : url.toString();
+        if (urlStr === "https://my-app.example/api/chat") {
+          return handoverResponse({
+            chatId: "chat-handover-reload",
+            accessToken: "handover-pat-reload",
+            chunks: [
+              { type: "trigger:session-state", activeInputSeq: 7 },
+              { type: "text-delta", id: "part-1", delta: "working" },
+            ],
+            close: false,
+          });
+        }
+        if (isSessionOutSubscribeUrl(urlStr)) {
+          return defaultSseResponse([
+            { type: "trigger:turn-complete", sessionInEventId: 6 },
+            { type: "text-delta", id: "part-2", delta: "current" },
+            { type: "trigger:turn-complete", sessionInEventId: 7 },
+          ]);
+        }
+        throw new Error(`Unexpected URL: ${urlStr}`);
+      });
+
+      const transport = new TriggerChatTransport({
+        task: "my-chat-task",
+        accessToken: () => "pat",
+        headStart: "https://my-app.example/api/chat",
+      });
+
+      const firstTurn = await transport.sendMessages({
+        trigger: "submit-message",
+        chatId: "chat-handover-reload",
+        messageId: "m1",
+        messages: [createUserMessage("first")],
+        abortSignal: undefined,
+      });
+      const firstTurnReader = firstTurn.getReader();
+      await expect(firstTurnReader.read()).resolves.toMatchObject({
+        value: { type: "text-delta", delta: "working" },
+      });
+
+      const persisted = transport.getSession("chat-handover-reload");
+      expect(persisted).toMatchObject({
+        publicAccessToken: "handover-pat-reload",
+        activeInputSeq: 7,
+        isStreaming: true,
+      });
+      await firstTurnReader.cancel();
+
+      const rehydrated = new TriggerChatTransport({
+        task: "my-chat-task",
+        accessToken: () => "pat",
+        sessions: { "chat-handover-reload": persisted! },
+      });
+      const resumed = await rehydrated.reconnectToStream({ chatId: "chat-handover-reload" });
+
+      expect(resumed).not.toBeNull();
+      await expect(drainChunks(resumed!)).resolves.toEqual([
+        { type: "text-delta", id: "part-2", delta: "current" },
+      ]);
+      expect(rehydrated.getSession("chat-handover-reload")?.activeInputSeq).toBeUndefined();
+    });
+
     it("bypasses endpoint when a session is already hydrated (page reload after first turn)", async () => {
       const requests: Array<{ url: string; init?: RequestInit }> = [];
       global.fetch = vi.fn().mockImplementation(async (url: string | URL, init?: RequestInit) => {

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

@@ -981,10 +981,10 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
     //     useChat resume / reconnectToStream path doesn't open a
     //     second `session.out` subscription on top of our stitched
     //     response.
-    //   - On `trigger:session-state`, hydrate `state.lastEventId`
-    //     with the agent's final S2 event id. Without this, turn 2's
-    //     `session.out` subscribe reads from the start and replays
-    //     turn 1's chunks back into the UI.
+    //   - On `trigger:session-state`, hydrate `state.activeInputSeq`
+    //     as soon as the handover append commits and `state.lastEventId`
+    //     when the agent stream finishes. These keep reloads correlated
+    //     to the active turn and keep turn 2 from replaying turn 1.
     //   - On stream end (handover-skip case — no
     //     `trigger:turn-complete` arrives, customer's stream just
     //     ends), also clear `isStreaming` for the same reason.
@@ -993,9 +993,17 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
       this.notifySessionChange(id, state);
     const TRIGGER_TURN_COMPLETE = "trigger:turn-complete";
     const TRIGGER_SESSION_STATE = "trigger:session-state";
+    const clearActiveTurn = () => {
+      const state = sessions.get(chatId);
+      if (state && (state.isStreaming || state.activeInputSeq !== undefined)) {
+        state.activeInputSeq = undefined;
+        state.isStreaming = false;
+        notifyChange(chatId, state);
+      }
+    };
     const clearStreaming = () => {
       const state = sessions.get(chatId);
-      if (state && state.isStreaming) {
+      if (state?.isStreaming) {
         state.isStreaming = false;
         notifyChange(chatId, state);
       }
@@ -1007,6 +1015,13 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
         notifyChange(chatId, state);
       }
     };
+    const setActiveInputSeq = (activeInputSeq: number) => {
+      const state = sessions.get(chatId);
+      if (state) {
+        state.activeInputSeq = activeInputSeq;
+        notifyChange(chatId, state);
+      }
+    };
     const emit = (event: ChatTransportEvent) => this.emitEvent(event);
     const attribution = () => this.turnAttribution(chatId);
     let sawFirstChunk = false;
@@ -1028,7 +1043,7 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
             if (chunk && typeof chunk === "object") {
               const type = (chunk as { type?: unknown }).type;
               if (type === TRIGGER_TURN_COMPLETE) {
-                clearStreaming();
+                clearActiveTurn();
                 emit({
                   type: "turn-completed",
                   chatId,
@@ -1039,10 +1054,17 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
                 return; // drop — not a real UIMessageChunk
               }
               if (type === TRIGGER_SESSION_STATE) {
-                const lastEventId = (chunk as { lastEventId?: unknown }).lastEventId;
+                const sessionState = chunk as {
+                  lastEventId?: unknown;
+                  activeInputSeq?: unknown;
+                };
+                const lastEventId = sessionState.lastEventId;
                 if (typeof lastEventId === "string") {
                   setLastEventId(lastEventId);
                 }
+                if (typeof sessionState.activeInputSeq === "number") {
+                  setActiveInputSeq(sessionState.activeInputSeq);
+                }
                 return; // drop
               }
             }
@@ -1760,6 +1782,7 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
             }) as typeof fetch)
           : undefined;
         let sawFirstChunk = false;
+        let sinceInSeq = options?.sinceInSeq;
 
         const connectSseOnce = async (token: string) => {
           const subscription = new SSEStreamSubscription(streamUrl, {
@@ -1993,10 +2016,10 @@ 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) {
+              if (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) {
+                if (!Number.isNaN(cursor) && cursor < sinceInSeq) {
                   continue;
                 }
               }
@@ -2014,6 +2037,8 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
                 sessionInEventId: headerValue(value.headers, SESSION_IN_EVENT_ID_HEADER),
                 ...this.turnAttribution(chatId),
               });
+              state.activeInputSeq = undefined;
+              sinceInSeq = undefined;
               state.isStreaming = false;
               this.notifySessionChange(chatId, state);
               this.coordinator?.release(chatId);