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

package/dist/index.js

@@ -2830,18 +2830,25 @@ function bindAbort(session, abortSignal) {
   abortSignal.addEventListener("abort", onAbort, { once: true });
   return onAbort;
 }
-async function* streamTurnContent(session, foldState, options, abortSignal, groupRootBaseline) {
-  const previousTurnId = options.previousTurnId === null ? void 0 : options.previousTurnId ?? "auto";
+async function* streamTurnContent(session, foldState, options, abortSignal, groupRootBaseline, onTurnIdAvailable) {
+  const previousTurnId = options.previousTurnId === null ? null : options.previousTurnId ?? "auto";
   const turn = session.prepareTurn({
     input: buildTurnInput(options),
-    ...previousTurnId != null ? { previousTurnId } : {}
+    ...previousTurnId !== void 0 ? { previousTurnId } : {}
   });
   const onAbort = bindAbort(session, abortSignal);
   if (abortSignal.aborted) {
     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 },
         {
@@ -2851,8 +2858,13 @@ async function* streamTurnContent(session, foldState, options, abortSignal, grou
       ),
       foldState,
       groupRootBaseline
-    );
+    )) {
+      notifyTurnIdIfAvailable();
+      yield update;
+    }
+    notifyTurnIdIfAvailable();
   } catch (error) {
+    notifyTurnIdIfAvailable();
     if (error instanceof Error && error.name === "AbortError") {
       return;
     }
@@ -3067,7 +3079,7 @@ function useTrueFoundryAgentMessages({
     [snapshot, projectOptions]
   );
   const runStream = useCallback2(
-    (createStream, turnId, isContinuation) => {
+    (createStream, turnIdRef, isContinuation) => {
       const streamGeneration = ++streamGenerationRef.current;
       abortControllerRef.current?.abort();
       const abortController = new AbortController();
@@ -3083,11 +3095,14 @@ function useTrueFoundryAgentMessages({
           if (pending == null || streamGeneration !== streamGenerationRef.current) {
             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
               }
@@ -3095,7 +3110,7 @@ function useTrueFoundryAgentMessages({
           );
         };
         const applyStreamUpdate = (update) => {
-          pendingStreamUpdate = { update, turnId, isContinuation };
+          pendingStreamUpdate = { update, isContinuation };
           if (streamUpdateRaf == null) {
             streamUpdateRaf = requestAnimationFrame(flushPendingStreamUpdate);
           }
@@ -3208,7 +3223,7 @@ function useTrueFoundryAgentMessages({
             void 0,
             loadedSnapshot.groupRootBaseline
           ),
-          turn.id,
+          { current: turn.id },
           isContinuation
         ).catch(() => void 0);
       }
@@ -3247,6 +3262,8 @@ function useTrueFoundryAgentMessages({
       const isContinuation = "inputs" in options || "resumeMcpAuth" in options && options.resumeMcpAuth === true;
       const continuationTurnId = snapshotRef.current.activeStream?.turnId;
       const turnId = isContinuation && continuationTurnId != null ? continuationTurnId : generateId();
+      const isFirstTurnInSession = "userMessage" in options && options.previousTurnId === void 0 && snapshotRef.current.turns.length === 0 && snapshotRef.current.pendingUser == null && snapshotRef.current.activeStream == null;
+      const turnIdRef = { current: turnId };
       if ("inputs" in options) {
         applyUserToolResponsesToFold(
           snapshotRef.current.fold,
@@ -3321,14 +3338,29 @@ function useTrueFoundryAgentMessages({
             snapshotRef.current.fold,
             {
               userMessage: options.userMessage,
-              ...options.previousTurnId !== void 0 ? { previousTurnId: options.previousTurnId } : {},
+              ...options.previousTurnId !== void 0 ? { previousTurnId: options.previousTurnId } : isFirstTurnInSession ? { previousTurnId: null } : {},
               ...streamHeaders
             },
             signal,
-            groupRootBaseline
+            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;
+              const renamePendingUser = (prev) => {
+                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
       );
     },
@@ -3414,7 +3446,7 @@ function useTrueFoundryAgentMessages({
         void 0,
         snapshotRef.current.groupRootBaseline
       ),
-      turn.id,
+      { current: turn.id },
       true
     );
   }, [runStream]);