@truefoundry/assistant-ui-runtime 0.1.0-rc.1 → 0.1.1

package/dist/index.js

@@ -276,7 +276,10 @@ function parseAskUserQuestionArgs(argsText) {
 
 // src/createSubAgent.ts
 function isCreateSubAgentToolCall(toolCall) {
-  return toolCall.toolInfo.type === "truefoundry-system" && toolCall.toolInfo.name === "create_sub_agent";
+  if (toolCall.toolInfo?.type === "truefoundry-system" && toolCall.toolInfo?.name === "create_sub_agent") {
+    return true;
+  }
+  return toolCall.function?.name === "create_sub_agent";
 }
 
 // src/modelMessageImageContent.ts
@@ -735,10 +738,13 @@ function attachSubAgentMessages(state, parentThreadId, parts) {
         });
       }
     }
-    if (messages.length === 0) {
+    if (messages.length === 0 && subAgents.length === 0) {
       return part;
     }
     const artifact = { subAgents };
+    if (messages.length === 0) {
+      return { ...part, artifact };
+    }
     return { ...part, messages, artifact };
   });
 }
@@ -1144,6 +1150,27 @@ function deriveSandboxId(messages) {
   return void 0;
 }
 
+// src/extractTurnUserText.ts
+function extractTurnUserText(input) {
+  const parts = [];
+  for (const item of input ?? []) {
+    if (item.type !== "user.message") {
+      continue;
+    }
+    const { content } = item;
+    if (typeof content === "string") {
+      parts.push(content);
+      continue;
+    }
+    for (const part of content) {
+      if (part.type === "text") {
+        parts.push(part.text);
+      }
+    }
+  }
+  return parts.join("\n").trim();
+}
+
 // src/mcpAuth.ts
 var MCP_AUTH_RESUME_RUN_CUSTOM_KEY = "resumeMcpAuth";
 function buildMcpAuthTextParts(servers) {
@@ -1230,27 +1257,6 @@ function appendMcpAuthToTurnContent(content, turn) {
   );
 }
 
-// src/extractTurnUserText.ts
-function extractTurnUserText(input) {
-  const parts = [];
-  for (const item of input ?? []) {
-    if (item.type !== "user.message") {
-      continue;
-    }
-    const { content } = item;
-    if (typeof content === "string") {
-      parts.push(content);
-      continue;
-    }
-    for (const part of content) {
-      if (part.type === "text") {
-        parts.push(part.text);
-      }
-    }
-  }
-  return parts.join("\n").trim();
-}
-
 // src/sessionSnapshot.ts
 function emptyRequiredActionsOverlay() {
   return {
@@ -1275,6 +1281,18 @@ function turnToSessionRecord(turn) {
     input: turn.input
   };
 }
+function sessionEventsToSessionRecord(turnId, createdEvent, doneEvent, rootModelMessageIds, sandboxId) {
+  const userText = extractTurnUserText(createdEvent.input);
+  return {
+    id: turnId,
+    ...userText ? { userText } : {},
+    createdAt: createdEvent.createdAt,
+    state: doneEvent.state,
+    input: createdEvent.input,
+    rootModelMessageIds,
+    ...sandboxId != null ? { sandboxId } : {}
+  };
+}
 function replaceSessionSnapshot(snapshot, patch) {
   return {
     ...snapshot,
@@ -1285,6 +1303,7 @@ function replaceSessionSnapshot(snapshot, patch) {
 
 // src/convertTurnMessages.ts
 var TURN_EVENTS_PAGE_SIZE = 25;
+var SESSION_EVENTS_PAGE_SIZE = 100;
 function assistantStatusFromTurnState(state) {
   switch (state.status) {
     case "done":
@@ -1691,6 +1710,92 @@ function projectSessionMessages(snapshot, options) {
   );
 }
 var DEFAULT_LIST_EVENTS_CONCURRENCY = 5;
+async function fetchAllSessionEvents(session, options) {
+  const items = [];
+  for await (const item of await session.listEvents({
+    limit: SESSION_EVENTS_PAGE_SIZE,
+    ...options?.lastTurnId != null ? { lastTurnId: options.lastTurnId } : {}
+  })) {
+    items.push(item);
+  }
+  items.reverse();
+  return items;
+}
+function ingestSessionEventsIntoSnapshot(snapshot, items, onTurnComplete) {
+  let currentTurnId = null;
+  let currentCreatedEvent = null;
+  let currentContentEvents = [];
+  let beforeCount = 0;
+  for (const item of items) {
+    const { turnId, event } = item;
+    if (event.type === "turn.created") {
+      currentTurnId = turnId;
+      currentCreatedEvent = event;
+      currentContentEvents = [];
+      beforeCount = snapshot.fold.threads.get(ROOT_THREAD_ID)?.modelMessageIds.length ?? 0;
+    } else if (event.type === "turn.done") {
+      if (currentTurnId == null || currentCreatedEvent == null) {
+        continue;
+      }
+      const afterBucket = snapshot.fold.threads.get(ROOT_THREAD_ID);
+      const rootModelMessageIds = (afterBucket?.modelMessageIds ?? []).slice(
+        beforeCount
+      );
+      const sandboxEvent = currentContentEvents.find(
+        (ev) => ev.type === "sandbox.created"
+      );
+      applyUserToolResponsesToFold(
+        snapshot.fold,
+        currentCreatedEvent.input ?? []
+      );
+      snapshot.turns.push(
+        sessionEventsToSessionRecord(
+          currentTurnId,
+          currentCreatedEvent,
+          event,
+          rootModelMessageIds,
+          sandboxEvent?.sandboxId
+        )
+      );
+      onTurnComplete?.(replaceSessionSnapshot(snapshot, {}));
+      currentTurnId = null;
+      currentCreatedEvent = null;
+      currentContentEvents = [];
+    } else {
+      if (currentTurnId != null) {
+        ingestTurnEvent(snapshot.fold, event);
+        currentContentEvents.push(event);
+      }
+    }
+  }
+}
+async function buildSnapshotFromSessionEvents(session, onProgress) {
+  let runningTurn;
+  for await (const turn of await session.listTurns({ limit: 1 })) {
+    if (turn.state?.status === "running") {
+      runningTurn = turn;
+    }
+  }
+  const items = await fetchAllSessionEvents(session);
+  const snapshot = createEmptySessionSnapshot();
+  ingestSessionEventsIntoSnapshot(snapshot, items, onProgress);
+  if (runningTurn == null) {
+    return snapshot;
+  }
+  const pendingUserText = extractTurnUserText(runningTurn.input);
+  return replaceSessionSnapshot(snapshot, {
+    runningTurn,
+    unstable_resume: true,
+    groupRootBaseline: computeGroupRootBaseline(snapshot.turns),
+    ...pendingUserText ? {
+      pendingUser: {
+        turnId: runningTurn.id,
+        content: extractTurnUserMessageContent(runningTurn.input),
+        createdAt: new Date(runningTurn.createdAt)
+      }
+    } : {}
+  });
+}
 function ingestTurnsIntoSnapshot(snapshot, turns, eventArrays) {
   let runningTurn;
   for (let i = 0; i < turns.length; i++) {
@@ -1719,37 +1824,32 @@ function ingestTurnsIntoSnapshot(snapshot, turns, eventArrays) {
   return runningTurn;
 }
 async function buildSnapshotFromSession(session, concurrency = DEFAULT_LIST_EVENTS_CONCURRENCY) {
-  const turns = [];
-  for await (const turn of await session.listTurns()) {
-    turns.push(turn);
+  const snapshot = await buildSnapshotFromSessionEvents(session);
+  if (snapshot.runningTurn == null) {
+    return snapshot;
   }
-  turns.reverse();
-  const eventArrays = await fetchAllTurnEventsWithConcurrency(turns, concurrency);
-  const snapshot = createEmptySessionSnapshot();
-  const runningTurn = ingestTurnsIntoSnapshot(snapshot, turns, eventArrays);
+  const turn = snapshot.runningTurn;
+  const eventArrays = await fetchAllTurnEventsWithConcurrency([turn], concurrency);
+  ingestTurnsIntoSnapshot(snapshot, [turn], eventArrays);
   return replaceSessionSnapshot(snapshot, {
-    ...runningTurn != null ? {
-      runningTurn,
-      unstable_resume: true,
-      groupRootBaseline: computeGroupRootBaseline(snapshot.turns)
-    } : {}
+    runningTurn: turn,
+    unstable_resume: true,
+    groupRootBaseline: computeGroupRootBaseline(snapshot.turns)
   });
 }
-async function buildSnapshotBeforeTurnIndex(session, turnIndex, concurrency = DEFAULT_LIST_EVENTS_CONCURRENCY, orderedTurns) {
+async function buildSnapshotBeforeTurnIndex(session, turnIndex, _concurrency = DEFAULT_LIST_EVENTS_CONCURRENCY, orderedTurns) {
   if (turnIndex <= 0) {
     return createEmptySessionSnapshot();
   }
   const turns = orderedTurns ?? await listSessionTurnsOrdered(session);
   const turnsToInclude = turns.slice(0, turnIndex);
-  if (turnsToInclude.length === 0) {
+  const lastTurnId = turnsToInclude.at(-1)?.id;
+  if (lastTurnId == null) {
     return createEmptySessionSnapshot();
   }
-  const eventArrays = await fetchAllTurnEventsWithConcurrency(
-    turnsToInclude,
-    concurrency
-  );
+  const items = await fetchAllSessionEvents(session, { lastTurnId });
   const snapshot = createEmptySessionSnapshot();
-  ingestTurnsIntoSnapshot(snapshot, turnsToInclude, eventArrays);
+  ingestSessionEventsIntoSnapshot(snapshot, items);
   return snapshot;
 }
 async function resolveGatewayBranchPreviousTurnId(session, turnIndex, orderedTurns) {
@@ -2024,7 +2124,7 @@ function repositoryItemsFromMessages(messages) {
   return items;
 }
 
-// src/draftSessionBridge.ts
+// src/private/draftSessionBridge.ts
 function createDraftSessionBridge(gateway) {
   async function getDraft(draftSessionId) {
     const response = await gateway.agents.private.draftSessions.get(draftSessionId);
@@ -2041,7 +2141,7 @@ function createDraftSessionBridge(gateway) {
   };
 }
 
-// src/agentSpec.ts
+// src/private/agentSpec.ts
 function mergeAgentSpec(base, update) {
   const { model: modelUpdate, ...rest } = update;
   const next = {
@@ -2069,7 +2169,7 @@ function sessionListStartTimestamp() {
   return start.toISOString();
 }
 
-// src/truefoundryDraftThreadListAdapter.ts
+// src/private/truefoundryDraftThreadListAdapter.ts
 var THREAD_LIST_PAGE_SIZE = 20;
 function createTrueFoundryDraftThreadListAdapter(options) {
   const { gateway, defaultAgentSpec, getAgentSpec } = options;
@@ -2137,7 +2237,7 @@ var EMPTY_DRAFT_EXTRAS = {
   }
 };
 
-// ../../node_modules/.pnpm/truefoundry-gateway-sdk@0.2.0/node_modules/truefoundry-gateway-sdk/dist/esm/core/fetcher/HttpResponsePromise.mjs
+// ../../node_modules/.pnpm/truefoundry-gateway-sdk@0.3.1/node_modules/truefoundry-gateway-sdk/dist/esm/core/fetcher/HttpResponsePromise.mjs
 var __awaiter = function(thisArg, _arguments, P, generator) {
   function adopt(value) {
     return value instanceof P ? value : new P(function(resolve) {
@@ -2252,7 +2352,7 @@ var HttpResponsePromise = class _HttpResponsePromise extends Promise {
   }
 };
 
-// ../../node_modules/.pnpm/truefoundry-gateway-sdk@0.2.0/node_modules/truefoundry-gateway-sdk/dist/esm/core/pagination/Page.mjs
+// ../../node_modules/.pnpm/truefoundry-gateway-sdk@0.3.1/node_modules/truefoundry-gateway-sdk/dist/esm/core/pagination/Page.mjs
 var __awaiter2 = function(thisArg, _arguments, P, generator) {
   function adopt(value) {
     return value instanceof P ? value : new P(function(resolve) {
@@ -2407,7 +2507,7 @@ var Page = class {
   }
 };
 
-// ../../node_modules/.pnpm/truefoundry-gateway-sdk@0.2.0/node_modules/truefoundry-gateway-sdk/dist/esm/agents/TurnStreamData.mjs
+// ../../node_modules/.pnpm/truefoundry-gateway-sdk@0.3.1/node_modules/truefoundry-gateway-sdk/dist/esm/agents/TurnStreamData.mjs
 function parseSequenceNumber(id) {
   if (!id) {
     throw new Error("Missing SSE sequence number id.");
@@ -2419,7 +2519,7 @@ function parseSequenceNumber(id) {
   return parsed;
 }
 
-// ../../node_modules/.pnpm/truefoundry-gateway-sdk@0.2.0/node_modules/truefoundry-gateway-sdk/dist/esm/agents/Turn.mjs
+// ../../node_modules/.pnpm/truefoundry-gateway-sdk@0.3.1/node_modules/truefoundry-gateway-sdk/dist/esm/agents/Turn.mjs
 var __awaiter3 = function(thisArg, _arguments, P, generator) {
   function adopt(value) {
     return value instanceof P ? value : new P(function(resolve) {
@@ -2657,7 +2757,7 @@ _Turn_client = /* @__PURE__ */ new WeakMap(), _Turn_state = /* @__PURE__ */ new
     __classPrivateFieldSet(this, _Turn_state, event.state, "f");
 };
 
-// ../../node_modules/.pnpm/truefoundry-gateway-sdk@0.2.0/node_modules/truefoundry-gateway-sdk/dist/esm/agents/PreparedTurn.mjs
+// ../../node_modules/.pnpm/truefoundry-gateway-sdk@0.3.1/node_modules/truefoundry-gateway-sdk/dist/esm/agents/PreparedTurn.mjs
 var __awaiter4 = function(thisArg, _arguments, P, generator) {
   function adopt(value) {
     return value instanceof P ? value : new P(function(resolve) {
@@ -3012,7 +3112,7 @@ var PreparedTurn = class {
 };
 _PreparedTurn_client = /* @__PURE__ */ new WeakMap(), _PreparedTurn_input = /* @__PURE__ */ new WeakMap(), _PreparedTurn_previousTurnIdInput = /* @__PURE__ */ new WeakMap(), _PreparedTurn_start = /* @__PURE__ */ new WeakMap(), _PreparedTurn_turn = /* @__PURE__ */ new WeakMap();
 
-// ../../node_modules/.pnpm/truefoundry-gateway-sdk@0.2.0/node_modules/truefoundry-gateway-sdk/dist/esm/agents/AgentSession.mjs
+// ../../node_modules/.pnpm/truefoundry-gateway-sdk@0.3.1/node_modules/truefoundry-gateway-sdk/dist/esm/agents/AgentSession.mjs
 var __awaiter5 = function(thisArg, _arguments, P, generator) {
   function adopt(value) {
     return value instanceof P ? value : new P(function(resolve) {
@@ -3122,10 +3222,22 @@ var AgentSession = class {
       yield __classPrivateFieldGet3(this, _AgentSession_client, "f").agents.sessions.cancel(this.id, {}, requestOptions);
     });
   }
+  /**
+   * Paginated session events across turns (newest first); subscribe to a running turn for live events.
+   *
+   * @param opts.pageToken - Token from the previous response nextPageToken.
+   * @param opts.lastTurnId - Newest turn in the listing window (initial load only). Omit to use the session last turn.
+   * @param opts.limit - Page size. Default 100.
+   * @param requestOptions - Overrides client timeout, retries, abortSignal, headers, queryParams.
+   * @returns {Promise<core.Page<TrueFoundryGatewayApi.SessionEventItem, TrueFoundryGatewayApi.ListSessionEventsResponse>>} Paginated session events.
+   */
+  listEvents(opts, requestOptions) {
+    return __classPrivateFieldGet3(this, _AgentSession_client, "f").agents.sessions.listEvents(this.id, opts, requestOptions);
+  }
 };
 _AgentSession_client = /* @__PURE__ */ new WeakMap();
 
-// src/bindDraftAgentSession.ts
+// src/private/bindDraftAgentSession.ts
 var inflightByDraftId = /* @__PURE__ */ new Map();
 function getGatewayFromSessionClient(client) {
   const internal = client;
@@ -3260,7 +3372,7 @@ function resolveTrueFoundryAgentRuntimeOptions(options) {
   };
 }
 
-// src/useDraftAgentSpec.ts
+// src/private/useDraftAgentSpec.ts
 import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 var SPEC_SYNC_DEBOUNCE_MS = 400;
 function useDraftAgentSpec({
@@ -3408,11 +3520,13 @@ import { useCallback as useCallback2, useEffect as useEffect2, useMemo as useMem
 
 // src/loadSessionSnapshot.ts
 var inflightBySessionId2 = /* @__PURE__ */ new Map();
-function loadSessionSnapshot(client, sessionId, concurrency, sessionOptions) {
+function loadSessionSnapshot(client, sessionId, sessionOptions, onProgress) {
   const cacheKey = sessionOptions?.draftGateway != null ? `draft:${sessionId}` : sessionId;
   let inflight = inflightBySessionId2.get(cacheKey);
   if (inflight == null) {
-    inflight = getSession(client, sessionId, sessionOptions).then((session) => buildSnapshotFromSession(session, concurrency)).finally(() => {
+    inflight = getSession(client, sessionId, sessionOptions).then(
+      (session) => buildSnapshotFromSessionEvents(session, onProgress)
+    ).finally(() => {
       if (inflightBySessionId2.get(cacheKey) === inflight) {
         inflightBySessionId2.delete(cacheKey);
       }
@@ -3654,6 +3768,7 @@ function resolveTurnInput(snapshot, turnId) {
 function useTrueFoundryAgentMessages({
   client,
   sessionId,
+  isMain,
   listEventsConcurrency,
   onError,
   initializeSession,
@@ -3667,13 +3782,21 @@ function useTrueFoundryAgentMessages({
   const [snapshot, setSnapshot] = useState2(createEmptySessionSnapshot);
   const [isRunning, setIsRunning] = useState2(false);
   const [isLoading, setIsLoading] = useState2(false);
+  const [loadRetryTrigger, setLoadRetryTrigger] = useState2(0);
   const snapshotRef = useRef2(snapshot);
   snapshotRef.current = snapshot;
+  const onErrorRef = useRef2(onError);
+  onErrorRef.current = onError;
+  const resolveConversationSessionIdRef = useRef2(resolveConversationSessionId);
+  resolveConversationSessionIdRef.current = resolveConversationSessionId;
+  const initializeSessionRef = useRef2(initializeSession);
+  initializeSessionRef.current = initializeSession;
   const createdAtByMessageIdRef = useRef2(/* @__PURE__ */ new Map());
   const abortControllerRef = useRef2(null);
   const activeRunRef = useRef2(null);
   const runningTurnRef = useRef2(void 0);
   const loadGenerationRef = useRef2(0);
+  const streamGenerationRef = useRef2(0);
   const lazilyCreatedSessionIdRef = useRef2(void 0);
   const projectOptions = useMemo2(
     () => ({
@@ -3693,41 +3816,62 @@ function useTrueFoundryAgentMessages({
     () => projectSessionMessages(snapshot, projectOptions),
     [snapshot, projectOptions]
   );
-  const applyStreamUpdate = useCallback2(
-    (update, turnId, isContinuation) => {
-      setSnapshot(
-        (prev) => replaceSessionSnapshot(prev, {
-          activeStream: {
-            turnId,
-            update,
-            isContinuation
-          }
-        })
-      );
-    },
-    []
-  );
   const runStream = useCallback2(
     (createStream, turnId, isContinuation) => {
+      const streamGeneration = ++streamGenerationRef.current;
       abortControllerRef.current?.abort();
       const abortController = new AbortController();
       abortControllerRef.current = abortController;
       setIsRunning(true);
       const run = (async () => {
+        let pendingStreamUpdate = null;
+        let streamUpdateRaf = null;
+        const flushPendingStreamUpdate = () => {
+          streamUpdateRaf = null;
+          const pending = pendingStreamUpdate;
+          pendingStreamUpdate = null;
+          if (pending == null || streamGeneration !== streamGenerationRef.current) {
+            return;
+          }
+          const { update, turnId: pendingTurnId, isContinuation: pendingIsContinuation } = pending;
+          setSnapshot(
+            (prev) => replaceSessionSnapshot(prev, {
+              activeStream: {
+                turnId: pendingTurnId,
+                update,
+                isContinuation: pendingIsContinuation
+              }
+            })
+          );
+        };
+        const applyStreamUpdate = (update) => {
+          pendingStreamUpdate = { update, turnId, isContinuation };
+          if (streamUpdateRaf == null) {
+            streamUpdateRaf = requestAnimationFrame(flushPendingStreamUpdate);
+          }
+        };
         try {
           for await (const update of createStream(abortController.signal)) {
             if (abortController.signal.aborted) {
               return;
             }
-            applyStreamUpdate(update, turnId, isContinuation);
+            applyStreamUpdate(update);
           }
         } catch (error) {
           if (error instanceof Error && error.name === "AbortError") {
             return;
           }
-          onError?.(error);
+          onErrorRef.current?.(error);
           throw error;
         } finally {
+          if (streamUpdateRaf != null) {
+            cancelAnimationFrame(streamUpdateRaf);
+            streamUpdateRaf = null;
+          }
+          if (streamGeneration !== streamGenerationRef.current) {
+            return;
+          }
+          flushPendingStreamUpdate();
           if (abortControllerRef.current === abortController) {
             abortControllerRef.current = null;
           }
@@ -3758,7 +3902,7 @@ function useTrueFoundryAgentMessages({
       });
       return run;
     },
-    [applyStreamUpdate, onError]
+    [onError]
   );
   const load = useCallback2(async () => {
     if (sessionId == null) {
@@ -3766,22 +3910,33 @@ function useTrueFoundryAgentMessages({
       setSnapshot(createEmptySessionSnapshot());
       return;
     }
+    if (isMain === false) return;
+    if (lazilyCreatedSessionIdRef.current != null && sessionId !== lazilyCreatedSessionIdRef.current) {
+      lazilyCreatedSessionIdRef.current = void 0;
+    }
     if (sessionId === lazilyCreatedSessionIdRef.current) {
       return;
     }
     const generation = ++loadGenerationRef.current;
+    ++streamGenerationRef.current;
     abortControllerRef.current?.abort();
+    createdAtByMessageIdRef.current = /* @__PURE__ */ new Map();
+    setSnapshot(createEmptySessionSnapshot());
     setIsLoading(true);
     try {
       const conversationSessionId = await resolveActiveSessionId(
         sessionId,
-        resolveConversationSessionId
+        resolveConversationSessionIdRef.current
       );
       const loadedSnapshot = await loadSessionSnapshot(
         client,
         conversationSessionId,
-        listEventsConcurrency,
-        sessionOptions
+        sessionOptions,
+        (snap) => {
+          if (generation === loadGenerationRef.current) {
+            setSnapshot(snap);
+          }
+        }
       );
       if (generation !== loadGenerationRef.current) {
         return;
@@ -3789,30 +3944,33 @@ function useTrueFoundryAgentMessages({
       createdAtByMessageIdRef.current = /* @__PURE__ */ new Map();
       setSnapshot(loadedSnapshot);
       runningTurnRef.current = loadedSnapshot.runningTurn;
+      setIsLoading(false);
       if (loadedSnapshot.runningTurn != null) {
         const turn = loadedSnapshot.runningTurn;
         const isContinuation = !extractTurnUserText(turn.input);
-        await runStream(
+        void runStream(
           (signal) => resumeTurnStream(
             turn,
-            snapshotRef.current.fold,
+            loadedSnapshot.fold,
             signal,
             void 0,
-            snapshotRef.current.groupRootBaseline
+            loadedSnapshot.groupRootBaseline
           ),
           turn.id,
           isContinuation
-        );
+        ).catch(() => void 0);
       }
     } catch (error) {
-      onError?.(error);
+      if (generation === loadGenerationRef.current) {
+        onErrorRef.current?.(error);
+      }
       throw error;
     } finally {
       if (generation === loadGenerationRef.current) {
         setIsLoading(false);
       }
     }
-  }, [client, listEventsConcurrency, onError, resolveConversationSessionId, runStream, sessionId, sessionOptions]);
+  }, [client, runStream, sessionId, sessionOptions, loadRetryTrigger, isMain]);
   useEffect2(() => {
     void load().catch(() => void 0);
   }, [load]);
@@ -3820,16 +3978,16 @@ function useTrueFoundryAgentMessages({
     async (options) => {
       let activeSessionId = sessionId;
       if (activeSessionId == null) {
-        if (initializeSession == null) {
+        if (initializeSessionRef.current == null) {
           throw new Error("Cannot send a turn without an active session.");
         }
-        const { remoteId } = await initializeSession();
+        const { remoteId } = await initializeSessionRef.current();
         activeSessionId = remoteId;
         lazilyCreatedSessionIdRef.current = remoteId;
       }
       const conversationSessionId = await resolveActiveSessionId(
         activeSessionId,
-        resolveConversationSessionId
+        resolveConversationSessionIdRef.current
       );
       const session = await getSession(client, conversationSessionId, sessionOptions);
       const isContinuation = "inputs" in options || "resumeMcpAuth" in options && options.resumeMcpAuth === true;
@@ -3841,29 +3999,48 @@ function useTrueFoundryAgentMessages({
           options.inputs
         );
       }
-      setSnapshot(
-        (prev) => commitActiveStream(
-          prev,
-          "inputs" in options ? options.inputs : void 0
-        )
-      );
+      const branchBase = "userMessage" in options ? options.branchFromSnapshot : void 0;
       let groupRootBaseline;
-      if ("userMessage" in options) {
-        const rootBucket = snapshotRef.current.fold.threads.get(ROOT_THREAD_ID);
+      if (branchBase != null && "userMessage" in options) {
+        const rootBucket = branchBase.fold.threads.get(ROOT_THREAD_ID);
         groupRootBaseline = [...rootBucket?.modelMessageIds ?? []];
+        const nextSnapshot = replaceSessionSnapshot(branchBase, {
+          pendingUser: {
+            turnId,
+            content: options.userMessage,
+            createdAt: /* @__PURE__ */ new Date()
+          },
+          activeStream: void 0,
+          groupRootBaseline
+        });
+        snapshotRef.current = nextSnapshot;
+        setSnapshot(nextSnapshot);
+      } else {
         setSnapshot(
-          (prev) => replaceSessionSnapshot(prev, {
-            pendingUser: {
-              turnId,
-              content: options.userMessage,
-              createdAt: /* @__PURE__ */ new Date()
-            },
-            activeStream: void 0,
-            groupRootBaseline
-          })
+          (prev) => commitActiveStream(
+            prev,
+            "inputs" in options ? options.inputs : void 0
+          )
         );
-      } else {
-        groupRootBaseline = snapshotRef.current.groupRootBaseline ?? computeGroupRootBaseline(snapshotRef.current.turns);
+        if ("userMessage" in options) {
+          const rootBucket = snapshotRef.current.fold.threads.get(ROOT_THREAD_ID);
+          groupRootBaseline = [...rootBucket?.modelMessageIds ?? []];
+          setSnapshot((prev) => {
+            const next = replaceSessionSnapshot(prev, {
+              pendingUser: {
+                turnId,
+                content: options.userMessage,
+                createdAt: /* @__PURE__ */ new Date()
+              },
+              activeStream: void 0,
+              groupRootBaseline
+            });
+            snapshotRef.current = next;
+            return next;
+          });
+        } else {
+          groupRootBaseline = snapshotRef.current.groupRootBaseline ?? computeGroupRootBaseline(snapshotRef.current.turns);
+        }
       }
       await runStream(
         (signal) => {
@@ -3900,7 +4077,7 @@ function useTrueFoundryAgentMessages({
         isContinuation
       );
     },
-    [client, initializeSession, resolveConversationSessionId, runStream, sessionId, sessionOptions]
+    [client, runStream, sessionId, sessionOptions]
   );
   const cancel = useCallback2(async () => {
     if (sessionId == null) {
@@ -3909,12 +4086,12 @@ function useTrueFoundryAgentMessages({
     }
     const conversationSessionId = await resolveActiveSessionId(
       sessionId,
-      resolveConversationSessionId
+      resolveConversationSessionIdRef.current
     );
     const session = await getSession(client, conversationSessionId, sessionOptions);
     await session.cancel().catch(() => void 0);
     await activeRunRef.current?.catch(() => void 0);
-  }, [client, resolveConversationSessionId, sessionId, sessionOptions]);
+  }, [client, sessionId, sessionOptions]);
   const isRunningRef = useRef2(isRunning);
   isRunningRef.current = isRunning;
   const trySendCollectedRequiredActions = useCallback2(
@@ -3929,10 +4106,10 @@ function useTrueFoundryAgentMessages({
       }
       const inputs = collectRequiredActionInputs(paused);
       if (inputs.length > 0) {
-        void sendTurn({ inputs }).catch((error) => onError?.(error));
+        void sendTurn({ inputs }).catch((error) => onErrorRef.current?.(error));
       }
     },
-    [onError, projectOptions, sendTurn]
+    [projectOptions, sendTurn]
   );
   const respondToToolApproval = useCallback2(
     (response) => {
@@ -3998,7 +4175,7 @@ function useTrueFoundryAgentMessages({
       await cancel();
       const conversationSessionId = await resolveActiveSessionId(
         activeSessionId,
-        resolveConversationSessionId
+        resolveConversationSessionIdRef.current
       );
       const session = await getSession(client, conversationSessionId, sessionOptions);
       const previousTurnId = await resolveGatewayBranchPreviousTurnId(
@@ -4011,17 +4188,18 @@ function useTrueFoundryAgentMessages({
         listEventsConcurrency
       );
       createdAtByMessageIdRef.current = /* @__PURE__ */ new Map();
+      snapshotRef.current = rewound;
       setSnapshot(rewound);
       await sendTurn({
         userMessage,
-        previousTurnId
+        previousTurnId,
+        branchFromSnapshot: rewound
       });
     },
     [
       cancel,
       client,
       listEventsConcurrency,
-      resolveConversationSessionId,
       sendTurn,
       sessionId,
       sessionOptions
@@ -4054,11 +4232,14 @@ function useTrueFoundryAgentMessages({
     },
     [branchFromTurn]
   );
+  const retryLoad = useCallback2(() => {
+    setLoadRetryTrigger((n) => n + 1);
+  }, []);
   return {
     messages,
     isRunning,
     isLoading,
-    load,
+    retryLoad,
     sendTurn,
     cancel,
     respondToToolApproval,
@@ -4088,6 +4269,9 @@ function useTrueFoundryAgentRuntimeImpl(options, pendingAgentSpecRef) {
     (state) => agent.mode === "draft" ? state.threadListItem.remoteId ?? void 0 : void 0
   );
   const sessionId = useAuiState((state) => state.threadListItem.remoteId ?? void 0);
+  const isMain = useAuiState(
+    (state) => state.threads.mainThreadId === state.threadListItem.id
+  );
   const draftSpec = useDraftAgentSpec({
     draftSessionId,
     draftBridge: draftBridgeRef.current,
@@ -4112,10 +4296,12 @@ function useTrueFoundryAgentRuntimeImpl(options, pendingAgentSpecRef) {
     respondToToolResponse,
     resumeRun,
     editFromTurn,
-    resetFromTurn
+    resetFromTurn,
+    retryLoad
   } = useTrueFoundryAgentMessages({
     client,
     sessionId,
+    isMain,
     listEventsConcurrency,
     onError,
     initializeSession,
@@ -4187,6 +4373,7 @@ function useTrueFoundryAgentRuntimeImpl(options, pendingAgentSpecRef) {
       resetFromTurn: (turnId) => resetFromTurn(turnId).catch((error) => {
         onError?.(error);
       }),
+      reload: retryLoad,
       draft: draftExtras
     }),
     unstable_enableToolInvocations: true,
@@ -4235,32 +4422,32 @@ function useTrueFoundryAgentRuntimeImpl(options, pendingAgentSpecRef) {
   });
 }
 function useTrueFoundryAgentRuntime(options) {
-  const resolved = useMemo3(
-    () => resolveTrueFoundryAgentRuntimeOptions(options),
-    [options]
-  );
+  const resolved = resolveTrueFoundryAgentRuntimeOptions(options);
   const { client, agent, gateway } = resolved;
   const pendingAgentSpecRef = useRef3(
     agent.mode === "draft" ? agent.defaultAgentSpec : void 0
   );
+  const agentMode = agent.mode;
+  const namedAgentName = agent.mode === "named" ? agent.agentName : void 0;
   const threadListAdapter = useMemo3(() => {
-    if (agent.mode === "draft") {
+    if (agentMode === "draft") {
       if (gateway == null) {
         throw new Error(
           "Draft agent mode requires a `gateway` TrueFoundryGateway client."
         );
       }
+      const draftAgent = agent;
       return createTrueFoundryDraftThreadListAdapter({
         gateway,
-        defaultAgentSpec: agent.defaultAgentSpec,
-        getAgentSpec: () => pendingAgentSpecRef.current ?? agent.defaultAgentSpec
+        defaultAgentSpec: draftAgent.defaultAgentSpec,
+        getAgentSpec: () => pendingAgentSpecRef.current ?? draftAgent.defaultAgentSpec
       });
     }
     return createTrueFoundryThreadListAdapter({
       client,
-      agentName: agent.agentName
+      agentName: namedAgentName
     });
-  }, [agent, client, gateway]);
+  }, [agentMode, namedAgentName, client, gateway]);
   return useRemoteThreadListRuntime({
     allowNesting: true,
     adapter: threadListAdapter,
@@ -4331,6 +4518,10 @@ var useTrueFoundryCancel = () => {
   const aui = useAui2();
   return () => trueFoundryExtras.get(aui).cancel();
 };
+var useTrueFoundryReload = () => {
+  const aui = useAui2();
+  return () => trueFoundryExtras.get(aui).reload();
+};
 var useTrueFoundryResetFromTurn = () => {
   const aui = useAui2();
   return (turnId) => trueFoundryExtras.get(aui).resetFromTurn(turnId);
@@ -4429,6 +4620,7 @@ export {
   useTrueFoundryCancel,
   useTrueFoundryDownloadSandboxFile,
   useTrueFoundryMcpAuth,
+  useTrueFoundryReload,
   useTrueFoundryResetFromTurn,
   useTrueFoundryRespondToToolApproval,
   useTrueFoundryRespondToToolResponse,