@truefoundry/assistant-ui-runtime 0.1.3-rc.1 → 0.1.3-rc.3

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@truefoundry/assistant-ui-runtime",
-  "version": "0.1.3-rc.1",
+  "version": "0.1.3-rc.3",
   "description": "TrueFoundry Gateway agent runtime adapter for assistant-ui",
   "license": "Apache-2.0",
   "repository": {

package/src/streamTurn.test.ts

@@ -227,6 +227,91 @@ describe("streamTurn", () => {
                 },
             );
         });
+
+        it("notifies gateway turn id when turn.done errors with no content yields", async () => {
+            const gatewayTurnId = "01ky6mqzmczwt6ssyd5r02gjjc";
+            const turn = {
+                id: undefined as string | undefined,
+                execute: vi.fn(() =>
+                    (async function* () {
+                        // SDK sets turn.id when turn.created is observed.
+                        turn.id = gatewayTurnId;
+                        yield streamData(1, {
+                            type: "turn.created",
+                            createdAt,
+                            input: [{ type: "user.message", content: "hello" }],
+                        });
+                        yield streamData(2, {
+                            type: "turn.done",
+                            createdAt,
+                            state: {
+                                status: "error",
+                                message:
+                                    "Publisher Model is not servable in region us-central1.",
+                            },
+                        });
+                    })(),
+                ),
+            };
+            const prepareTurn = vi.fn(() => turn);
+            const session = {
+                prepareTurn,
+                cancel: vi.fn().mockResolvedValue(undefined),
+            } as unknown as AgentSession;
+            const onTurnIdAvailable = vi.fn();
+
+            await expect(
+                collectUpdates(
+                    streamTurnContent(
+                        session,
+                        new PeerThreadFoldState(),
+                        { userMessage: "hello" },
+                        new AbortController().signal,
+                        undefined,
+                        onTurnIdAvailable,
+                    ),
+                ),
+            ).rejects.toThrow("Publisher Model is not servable in region us-central1.");
+
+            expect(onTurnIdAvailable).toHaveBeenCalledTimes(1);
+            expect(onTurnIdAvailable).toHaveBeenCalledWith(gatewayTurnId);
+        });
+
+        it("does not notify when an error stream never assigns turn.id", async () => {
+            const execute = vi.fn(() =>
+                (async function* () {
+                    yield streamData(1, {
+                        type: "turn.done",
+                        createdAt,
+                        state: {
+                            status: "error",
+                            message: "boom",
+                        },
+                    });
+                })(),
+            );
+            const prepareTurn = vi.fn(() => ({ id: undefined, execute }));
+            const session = {
+                prepareTurn,
+                cancel: vi.fn().mockResolvedValue(undefined),
+            } as unknown as AgentSession;
+            const onTurnIdAvailable = vi.fn();
+
+            await expect(
+                collectUpdates(
+                    streamTurnContent(
+                        session,
+                        new PeerThreadFoldState(),
+                        { userMessage: "hello" },
+                        new AbortController().signal,
+                        undefined,
+                        onTurnIdAvailable,
+                    ),
+                ),
+            ).rejects.toThrow("boom");
+
+            expect(onTurnIdAvailable).not.toHaveBeenCalled();
+        });
     });
 
     describe("resumeTurnStream", () => {

package/src/streamTurn.ts

@@ -2,6 +2,7 @@ import {
     type AgentSession,
     type Turn,
     type TurnInputItem,
+    TrueFoundryGatewayApi,
 } from "truefoundry-gateway-sdk/agents";
 
 import {
@@ -18,7 +19,7 @@ export type StreamTurnOptions = {
     inputs?: RequiredActionInput[];
     /**
      * Branch anchor for `prepareTurn`. Omit for `"auto"`. Pass `null` for a fresh
-     * root turn (no `previousTurnId` field).
+     * root turn — sent as `previous_turn_id: null` on the wire.
      */
     previousTurnId?: string | null;
     /** Extra headers for the createTurn request (`execute` request options). */
@@ -53,14 +54,26 @@ export async function* streamTurnContent(
     options: StreamTurnOptions,
     abortSignal: AbortSignal,
     groupRootBaseline?: readonly string[],
+    /**
+     * Called once with the gateway-assigned turn ID as soon as it becomes
+     * available (after the first `turn.created` SSE event). Use this to
+     * reconcile the locally-generated optimistic ID with the real gateway ID
+     * so that edit/retry can find the turn in `buildSnapshotBeforeTurn`.
+     */
+    onTurnIdAvailable?: (turnId: string) => void,
 ): AsyncGenerator<TurnStreamUpdate> {
-    const previousTurnId =
+    // When previousTurnId is explicitly null, pass it through to prepareTurn so the
+    // SDK serializer emits `previous_turn_id: null` on the wire (first turn in session).
+    // The SDK type doesn't admit null, but the Fern serializer handles it correctly.
+    const previousTurnId: TrueFoundryGatewayApi.PreviousTurnIdInput | null | undefined =
         options.previousTurnId === null
-            ? undefined
+            ? null
             : (options.previousTurnId ?? "auto");
     const turn = session.prepareTurn({
         input: buildTurnInput(options),
-        ...(previousTurnId != null ? { previousTurnId } : {}),
+        ...(previousTurnId !== undefined
+            ? { previousTurnId: previousTurnId as TrueFoundryGatewayApi.PreviousTurnIdInput }
+            : {}),
     });
 
     const onAbort = bindAbort(session, abortSignal);
@@ -68,8 +81,16 @@ export async function* streamTurnContent(
         return;
     }
 
+    let turnIdNotified = false;
+    const notifyTurnIdIfAvailable = () => {
+        if (!turnIdNotified && turn.id != null) {
+            onTurnIdAvailable?.(turn.id);
+            turnIdNotified = true;
+        }
+    };
+
     try {
-        yield* streamTurnEvents(
+        for await (const update of streamTurnEvents(
             turn.execute(
                 { stream: true },
                 {
@@ -79,8 +100,21 @@ export async function* streamTurnContent(
             ),
             foldState,
             groupRootBaseline,
-        );
+        )) {
+            // After the first `turn.created` event, `turn.id` is set.
+            // Notify BEFORE yielding so the caller can update its tracking
+            // before the snapshot is written with the stream update.
+            notifyTurnIdIfAvailable();
+            yield update;
+        }
+        // Handle streams that complete without yielding any content.
+        notifyTurnIdIfAvailable();
     } catch (error) {
+        // Error streams often throw on `turn.done` (status=error) without ever
+        // yielding content (e.g. model not servable). `turn.id` is still set
+        // after `turn.created` — notify so edit/retry can resolve the turn.
+        // Same for AbortError after create: keep local ids aligned with gateway.
+        notifyTurnIdIfAvailable();
         if (error instanceof Error && error.name === "AbortError") {
             return;
         }

package/src/useTrueFoundryAgentMessages.ts

@@ -353,7 +353,14 @@ export function useTrueFoundryAgentMessages({
     const runStream = useCallback(
         (
             createStream: (signal: AbortSignal) => AsyncGenerator<TurnStreamUpdate>,
-            turnId: string,
+            /**
+             * A mutable ref whose `.current` is the turn ID to use for
+             * `activeStream.turnId`. Callers that capture the gateway turn ID
+             * via `onTurnIdAvailable` update this ref in-place so that both the
+             * pending-update flush and `commitActiveStream` always see the real
+             * gateway ID rather than the locally-generated optimistic one.
+             */
+            turnIdRef: { current: string },
             isContinuation: boolean,
         ): Promise<void> => {
             const streamGeneration = ++streamGenerationRef.current;
@@ -368,7 +375,6 @@ export function useTrueFoundryAgentMessages({
                 // message tree (UI hang). The buffer belongs to this stream only.
                 let pendingStreamUpdate: {
                     update: TurnStreamUpdate;
-                    turnId: string;
                     isContinuation: boolean;
                 } | null = null;
                 let streamUpdateRaf: number | null = null;
@@ -383,12 +389,14 @@ export function useTrueFoundryAgentMessages({
                     ) {
                         return;
                     }
-                    const { update, turnId: pendingTurnId, isContinuation: pendingIsContinuation } =
-                        pending;
+                    const { update, isContinuation: pendingIsContinuation } = pending;
                     setSnapshot((prev) =>
                         replaceSessionSnapshot(prev, {
                             activeStream: {
-                                turnId: pendingTurnId,
+                                // Read from the ref so we always use the latest ID,
+                                // including any gateway ID that arrived after the RAf
+                                // was scheduled.
+                                turnId: turnIdRef.current,
                                 update,
                                 isContinuation: pendingIsContinuation,
                             },
@@ -397,7 +405,7 @@ export function useTrueFoundryAgentMessages({
                 };
 
                 const applyStreamUpdate = (update: TurnStreamUpdate) => {
-                    pendingStreamUpdate = { update, turnId, isContinuation };
+                    pendingStreamUpdate = { update, isContinuation };
                     if (streamUpdateRaf == null) {
                         streamUpdateRaf = requestAnimationFrame(flushPendingStreamUpdate);
                     }
@@ -540,7 +548,7 @@ export function useTrueFoundryAgentMessages({
                             undefined,
                             loadedSnapshot.groupRootBaseline,
                         ),
-                    turn.id,
+                    { current: turn.id },
                     isContinuation,
                 ).catch(() => undefined);
             }
@@ -588,6 +596,18 @@ export function useTrueFoundryAgentMessages({
                 isContinuation && continuationTurnId != null
                     ? continuationTurnId
                     : generateId();
+            // First turns must send previousTurnId: null.
+            const isFirstTurnInSession =
+                "userMessage" in options &&
+                options.previousTurnId === undefined &&
+                snapshotRef.current.turns.length === 0 &&
+                snapshotRef.current.pendingUser == null &&
+                snapshotRef.current.activeStream == null;
+
+            // Mutable ref so runStream always reads the latest ID. For new
+            // user-message turns the local `generateId()` value is replaced
+            // with the gateway-assigned ID once the first SSE event arrives.
+            const turnIdRef = { current: turnId };
 
             if ("inputs" in options) {
                 applyUserToolResponsesToFold(
@@ -676,14 +696,33 @@ export function useTrueFoundryAgentMessages({
                             userMessage: options.userMessage,
                             ...(options.previousTurnId !== undefined
                                 ? { previousTurnId: options.previousTurnId }
-                                : {}),
+                                : isFirstTurnInSession
+                                  ? { previousTurnId: null }
+                                  : {}),
                             ...streamHeaders,
                         },
                         signal,
                         groupRootBaseline,
+                        // Rename the optimistic local ID to the gateway turn ID
+                        // so that edit/retry can resolve the turn via the gateway.
+                        (gatewayTurnId) => {
+                            const oldId = turnIdRef.current;
+                            if (gatewayTurnId === oldId) return;
+                            turnIdRef.current = gatewayTurnId;
+                            // Rename in the ref immediately so any synchronous read
+                            // (e.g. commitActiveStream) sees the correct ID.
+                            const renamePendingUser = (prev: SessionSnapshot): SessionSnapshot => {
+                                if (prev.pendingUser?.turnId !== oldId) return prev;
+                                return replaceSessionSnapshot(prev, {
+                                    pendingUser: { ...prev.pendingUser, turnId: gatewayTurnId },
+                                });
+                            };
+                            snapshotRef.current = renamePendingUser(snapshotRef.current);
+                            setSnapshot(renamePendingUser);
+                        },
                     );
                 },
-                turnId,
+                turnIdRef,
                 isContinuation,
             );
         },
@@ -784,7 +823,7 @@ export function useTrueFoundryAgentMessages({
                     undefined,
                     snapshotRef.current.groupRootBaseline,
                 ),
-            turn.id,
+            { current: turn.id },
             true,
         );
     }, [runStream]);