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

package/dist/index.js

@@ -2830,7 +2830,7 @@ function bindAbort(session, abortSignal) {
   abortSignal.addEventListener("abort", onAbort, { once: true });
   return onAbort;
 }
-async function* streamTurnContent(session, foldState, options, abortSignal, groupRootBaseline) {
+async function* streamTurnContent(session, foldState, options, abortSignal, groupRootBaseline, onTurnIdAvailable) {
   const previousTurnId = options.previousTurnId === null ? void 0 : options.previousTurnId ?? "auto";
   const turn = session.prepareTurn({
     input: buildTurnInput(options),
@@ -2841,7 +2841,8 @@ async function* streamTurnContent(session, foldState, options, abortSignal, grou
     return;
   }
   try {
-    yield* streamTurnEvents(
+    let turnIdNotified = false;
+    for await (const update of streamTurnEvents(
       turn.execute(
         { stream: true },
         {
@@ -2851,7 +2852,16 @@ async function* streamTurnContent(session, foldState, options, abortSignal, grou
       ),
       foldState,
       groupRootBaseline
-    );
+    )) {
+      if (!turnIdNotified && turn.id != null) {
+        onTurnIdAvailable?.(turn.id);
+        turnIdNotified = true;
+      }
+      yield update;
+    }
+    if (!turnIdNotified && turn.id != null) {
+      onTurnIdAvailable?.(turn.id);
+    }
   } catch (error) {
     if (error instanceof Error && error.name === "AbortError") {
       return;
@@ -3067,7 +3077,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 +3093,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 +3108,7 @@ function useTrueFoundryAgentMessages({
           );
         };
         const applyStreamUpdate = (update) => {
-          pendingStreamUpdate = { update, turnId, isContinuation };
+          pendingStreamUpdate = { update, isContinuation };
           if (streamUpdateRaf == null) {
             streamUpdateRaf = requestAnimationFrame(flushPendingStreamUpdate);
           }
@@ -3208,7 +3221,7 @@ function useTrueFoundryAgentMessages({
             void 0,
             loadedSnapshot.groupRootBaseline
           ),
-          turn.id,
+          { current: turn.id },
           isContinuation
         ).catch(() => void 0);
       }
@@ -3247,6 +3260,7 @@ 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 turnIdRef = { current: turnId };
       if ("inputs" in options) {
         applyUserToolResponsesToFold(
           snapshotRef.current.fold,
@@ -3325,10 +3339,25 @@ function useTrueFoundryAgentMessages({
               ...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 +3443,7 @@ function useTrueFoundryAgentMessages({
         void 0,
         snapshotRef.current.groupRootBaseline
       ),
-      turn.id,
+      { current: turn.id },
       true
     );
   }, [runStream]);