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

package/src/useTrueFoundryAgentMessages.ts

@@ -57,6 +57,8 @@ import type { TurnStreamUpdate } from "./turnStreamUpdate.js";
 export type UseTrueFoundryAgentMessagesOptions = {
     client: AgentSessionClient;
     sessionId: string | undefined;
+    /** When true the thread is the currently selected (main) thread. */
+    isMain?: boolean | undefined;
     listEventsConcurrency?: number | undefined;
     onError?: ((error: unknown) => void) | undefined;
     initializeSession?: () => Promise<{
@@ -70,7 +72,16 @@ export type UseTrueFoundryAgentMessagesOptions = {
 };
 
 export type SendTurnOptions =
-    | { userMessage: UserMessageContent; previousTurnId?: string | null }
+    | {
+          userMessage: UserMessageContent;
+          previousTurnId?: string | null;
+          /**
+           * When branching (edit/reset), the already-rewound history to send from.
+           * Applied atomically with `pendingUser` so a stale React snapshot cannot
+           * keep pre-branch turns while the new user message is appended.
+           */
+          branchFromSnapshot?: SessionSnapshot;
+      }
     | { inputs: RequiredActionInput[] }
     | { resumeMcpAuth: true };
 
@@ -272,6 +283,7 @@ function resolveTurnInput(
 export function useTrueFoundryAgentMessages({
     client,
     sessionId,
+    isMain,
     listEventsConcurrency,
     onError,
     initializeSession,
@@ -287,15 +299,24 @@ export function useTrueFoundryAgentMessages({
     const [snapshot, setSnapshot] = useState<SessionSnapshot>(createEmptySessionSnapshot);
     const [isRunning, setIsRunning] = useState(false);
     const [isLoading, setIsLoading] = useState(false);
+    const [loadRetryTrigger, setLoadRetryTrigger] = useState(0);
 
     const snapshotRef = useRef(snapshot);
     snapshotRef.current = snapshot;
 
+    const onErrorRef = useRef(onError);
+    onErrorRef.current = onError;
+    const resolveConversationSessionIdRef = useRef(resolveConversationSessionId);
+    resolveConversationSessionIdRef.current = resolveConversationSessionId;
+    const initializeSessionRef = useRef(initializeSession);
+    initializeSessionRef.current = initializeSession;
+
     const createdAtByMessageIdRef = useRef(new Map<string, Date>());
     const abortControllerRef = useRef<AbortController | null>(null);
     const activeRunRef = useRef<Promise<void> | null>(null);
     const runningTurnRef = useRef<Turn | undefined>(undefined);
     const loadGenerationRef = useRef(0);
+    const streamGenerationRef = useRef(0);
     const lazilyCreatedSessionIdRef = useRef<string | undefined>(undefined);
 
     const projectOptions = useMemo(
@@ -318,47 +339,81 @@ export function useTrueFoundryAgentMessages({
         [snapshot, projectOptions],
     );
 
-    const applyStreamUpdate = useCallback(
-        (update: TurnStreamUpdate, turnId: string, isContinuation: boolean) => {
-            setSnapshot((prev) =>
-                replaceSessionSnapshot(prev, {
-                    activeStream: {
-                        turnId,
-                        update,
-                        isContinuation,
-                    },
-                }),
-            );
-        },
-        [],
-    );
-
     const runStream = useCallback(
         (
             createStream: (signal: AbortSignal) => AsyncGenerator<TurnStreamUpdate>,
             turnId: string,
             isContinuation: boolean,
         ): Promise<void> => {
+            const streamGeneration = ++streamGenerationRef.current;
             abortControllerRef.current?.abort();
             const abortController = new AbortController();
             abortControllerRef.current = abortController;
             setIsRunning(true);
 
             const run = (async () => {
+                // Sub-agent turns can emit 100+ stream events per frame. Coalesce to one
+                // setSnapshot per animation frame so assistant-ui does not remount the whole
+                // 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;
+
+                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: TurnStreamUpdate) => {
+                    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;
                     }
@@ -392,7 +447,7 @@ export function useTrueFoundryAgentMessages({
                 });
             return run;
         },
-        [applyStreamUpdate, onError],
+        [onError],
     );
 
     const load = useCallback(async () => {
@@ -401,24 +456,45 @@ export function useTrueFoundryAgentMessages({
             setSnapshot(createEmptySessionSnapshot());
             return;
         }
+
+        // Thread components are never unmounted on navigation — isMain going
+        // false→true is the only reliable signal that the user has clicked on
+        // this thread.  Skip the load when this thread is not the active one
+        // so that isMain being a dep triggers a fresh load on every selection.
+        if (isMain === false) return;
+
+        // When we are loading a *different* session the user has navigated away
+        // from the lazily-created one — clear the guard so navigating back to it
+        // later triggers a proper reload instead of silently skipping.
+        if (lazilyCreatedSessionIdRef.current != null && sessionId !== lazilyCreatedSessionIdRef.current) {
+            lazilyCreatedSessionIdRef.current = undefined;
+        }
+
         if (sessionId === lazilyCreatedSessionIdRef.current) {
             return;
         }
 
         const generation = ++loadGenerationRef.current;
+        ++streamGenerationRef.current;
         abortControllerRef.current?.abort();
+        createdAtByMessageIdRef.current = new Map();
+        setSnapshot(createEmptySessionSnapshot());
         setIsLoading(true);
 
         try {
             const conversationSessionId = await resolveActiveSessionId(
                 sessionId,
-                resolveConversationSessionId,
+                resolveConversationSessionIdRef.current,
             );
             const loadedSnapshot = await loadSessionSnapshot(
                 client,
                 conversationSessionId,
-                listEventsConcurrency,
                 sessionOptions,
+                (snap) => {
+                    if (generation === loadGenerationRef.current) {
+                        setSnapshot(snap);
+                    }
+                },
             );
             if (generation !== loadGenerationRef.current) {
                 return;
@@ -428,32 +504,44 @@ export function useTrueFoundryAgentMessages({
             setSnapshot(loadedSnapshot);
             runningTurnRef.current = loadedSnapshot.runningTurn;
 
+            // History (and any seeded pendingUser for the in-flight turn) is
+            // ready — clear loading before resuming. Awaiting the subscribe
+            // stream here previously kept isLoading true for the entire
+            // backend run, so reconnect UIs stayed on shimmers even though
+            // subscribe was already live.
+            setIsLoading(false);
+
             if (loadedSnapshot.runningTurn != null) {
                 const turn = loadedSnapshot.runningTurn;
                 const isContinuation = !extractTurnUserText(turn.input);
                 // TODO: pass afterSequenceNumber once stream ingestion tracks sequence numbers.
-                await runStream(
+                // Use loadedSnapshot directly — snapshotRef.current still points at
+                // the empty snapshot cleared above until the setSnapshot(loadedSnapshot)
+                // call re-renders.
+                void runStream(
                     (signal) =>
                         resumeTurnStream(
                             turn,
-                            snapshotRef.current.fold,
+                            loadedSnapshot.fold,
                             signal,
                             undefined,
-                            snapshotRef.current.groupRootBaseline,
+                            loadedSnapshot.groupRootBaseline,
                         ),
                     turn.id,
                     isContinuation,
-                );
+                ).catch(() => undefined);
             }
         } 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]);
 
     useEffect(() => {
         void load().catch(() => undefined);
@@ -463,17 +551,17 @@ export function useTrueFoundryAgentMessages({
         async (options: SendTurnOptions) => {
             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 =
@@ -492,34 +580,57 @@ export function useTrueFoundryAgentMessages({
                 );
             }
 
-            setSnapshot((prev) =>
-                commitActiveStream(
-                    prev,
-                    "inputs" in options ? options.inputs : undefined,
-                ),
-            );
+            const branchBase =
+                "userMessage" in options ? options.branchFromSnapshot : undefined;
 
             let groupRootBaseline: readonly string[] | undefined;
 
-            if ("userMessage" in options) {
-                const rootBucket =
-                    snapshotRef.current.fold.threads.get(ROOT_THREAD_ID);
+            if (branchBase != null && "userMessage" in options) {
+                // Atomic apply: never merge pendingUser onto a stale React `prev`
+                // that still holds pre-branch turns (edit would show old + new).
+                const rootBucket = branchBase.fold.threads.get(ROOT_THREAD_ID);
                 groupRootBaseline = [...(rootBucket?.modelMessageIds ?? [])];
+                const nextSnapshot = replaceSessionSnapshot(branchBase, {
+                    pendingUser: {
+                        turnId,
+                        content: options.userMessage,
+                        createdAt: new Date(),
+                    },
+                    activeStream: undefined,
+                    groupRootBaseline,
+                });
+                snapshotRef.current = nextSnapshot;
+                setSnapshot(nextSnapshot);
+            } else {
                 setSnapshot((prev) =>
-                    replaceSessionSnapshot(prev, {
-                        pendingUser: {
-                            turnId,
-                            content: options.userMessage,
-                            createdAt: new Date(),
-                        },
-                        activeStream: undefined,
-                        groupRootBaseline,
-                    }),
+                    commitActiveStream(
+                        prev,
+                        "inputs" in options ? options.inputs : undefined,
+                    ),
                 );
-            } 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: new Date(),
+                            },
+                            activeStream: undefined,
+                            groupRootBaseline,
+                        });
+                        snapshotRef.current = next;
+                        return next;
+                    });
+                } else {
+                    groupRootBaseline =
+                        snapshotRef.current.groupRootBaseline ??
+                        computeGroupRootBaseline(snapshotRef.current.turns);
+                }
             }
 
             await runStream(
@@ -559,7 +670,7 @@ export function useTrueFoundryAgentMessages({
                 isContinuation,
             );
         },
-        [client, initializeSession, resolveConversationSessionId, runStream, sessionId, sessionOptions],
+        [client, runStream, sessionId, sessionOptions],
     );
 
     const cancel = useCallback(async () => {
@@ -569,7 +680,7 @@ export function useTrueFoundryAgentMessages({
         }
         const conversationSessionId = await resolveActiveSessionId(
             sessionId,
-            resolveConversationSessionId,
+            resolveConversationSessionIdRef.current,
         );
         const session = await getSession(client, conversationSessionId, sessionOptions);
         // Request cancellation but keep consuming the stream. After cancel(),
@@ -581,7 +692,7 @@ export function useTrueFoundryAgentMessages({
         // reconcile is needed here — the cancelled turn is terminal and local
         // state reconciles against the event log on the next session load.
         await activeRunRef.current?.catch(() => undefined);
-    }, [client, resolveConversationSessionId, sessionId, sessionOptions]);
+    }, [client, sessionId, sessionOptions]);
 
     const isRunningRef = useRef(isRunning);
     isRunningRef.current = isRunning;
@@ -598,10 +709,10 @@ export 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 = useCallback(
@@ -677,7 +788,7 @@ export function useTrueFoundryAgentMessages({
 
             const conversationSessionId = await resolveActiveSessionId(
                 activeSessionId,
-                resolveConversationSessionId,
+                resolveConversationSessionIdRef.current,
             );
             const session = await getSession(client, conversationSessionId, sessionOptions);
             const previousTurnId = await resolveGatewayBranchPreviousTurnId(
@@ -690,18 +801,21 @@ export function useTrueFoundryAgentMessages({
                 listEventsConcurrency,
             );
             createdAtByMessageIdRef.current = new Map();
+            // Keep the ref aligned before awaiting sendTurn so any intermediate
+            // reads (and the atomic pendingUser apply) see the rewound history.
+            snapshotRef.current = rewound;
             setSnapshot(rewound);
 
             await sendTurn({
                 userMessage,
                 previousTurnId,
+                branchFromSnapshot: rewound,
             });
         },
         [
             cancel,
             client,
             listEventsConcurrency,
-            resolveConversationSessionId,
             sendTurn,
             sessionId,
             sessionOptions,
@@ -737,11 +851,15 @@ export function useTrueFoundryAgentMessages({
         [branchFromTurn],
     );
 
+    const retryLoad = useCallback(() => {
+        setLoadRetryTrigger((n) => n + 1);
+    }, []);
+
     return {
         messages,
         isRunning,
         isLoading,
-        load,
+        retryLoad,
         sendTurn,
         cancel,
         respondToToolApproval,

package/src/useTrueFoundryAgentRuntime.ts

@@ -14,7 +14,7 @@ import { useAui, useAuiState } from "@assistant-ui/store";
 import type { MutableRefObject } from "react";
 import { useCallback, useMemo, useRef, useState } from "react";
 
-import type { AgentSpec } from "./agentSpec.js";
+import type { AgentSpec } from "./private/agentSpec.js";
 import {
     collectPendingApprovals,
     collectPendingToolResponses,
@@ -26,14 +26,14 @@ import {
     extractEditedText,
     parseTurnIdFromMessageId,
 } from "./convertTurnMessages.js";
-import { createDraftSessionBridge } from "./draftSessionBridge.js";
+import { createDraftSessionBridge } from "./private/draftSessionBridge.js";
 import { MCP_AUTH_RESUME_RUN_CUSTOM_KEY } from "./mcpAuth.js";
-import { createTrueFoundryDraftThreadListAdapter } from "./truefoundryDraftThreadListAdapter.js";
+import { createTrueFoundryDraftThreadListAdapter } from "./private/truefoundryDraftThreadListAdapter.js";
 import { trueFoundryExtras } from "./truefoundryExtras.js";
 import { createTrueFoundryThreadListAdapter } from "./truefoundryThreadListAdapter.js";
 import type { UseTrueFoundryAgentRuntimeOptions } from "./types.js";
 import { resolveTrueFoundryAgentRuntimeOptions } from "./types.js";
-import { useDraftAgentSpec } from "./useDraftAgentSpec.js";
+import { useDraftAgentSpec } from "./private/useDraftAgentSpec.js";
 import { useTrueFoundryAgentMessages } from "./useTrueFoundryAgentMessages.js";
 
 function useTrueFoundryAgentRuntimeImpl(
@@ -63,6 +63,9 @@ function useTrueFoundryAgentRuntimeImpl(
                 : undefined,
     );
     const sessionId = useAuiState((state) => state.threadListItem.remoteId ?? undefined);
+    const isMain = useAuiState(
+        (state) => state.threads.mainThreadId === state.threadListItem.id,
+    );
 
     const draftSpec = useDraftAgentSpec({
         draftSessionId,
@@ -94,9 +97,11 @@ function useTrueFoundryAgentRuntimeImpl(
         resumeRun,
         editFromTurn,
         resetFromTurn,
+        retryLoad,
     } = useTrueFoundryAgentMessages({
         client,
         sessionId,
+        isMain,
         listEventsConcurrency,
         onError,
         initializeSession,
@@ -175,6 +180,7 @@ function useTrueFoundryAgentRuntimeImpl(
                 resetFromTurn(turnId).catch((error) => {
                     onError?.(error);
                 }),
+            reload: retryLoad,
             draft: draftExtras,
         }),
         unstable_enableToolInvocations: true,
@@ -228,34 +234,42 @@ function useTrueFoundryAgentRuntimeImpl(
 }
 
 export function useTrueFoundryAgentRuntime(options: UseTrueFoundryAgentRuntimeOptions) {
-    const resolved = useMemo(
-        () => resolveTrueFoundryAgentRuntimeOptions(options),
-        [options],
-    );
+    // resolveTrueFoundryAgentRuntimeOptions is a cheap pure function; memoizing on the
+    // whole `options` object (which callers typically pass as an inline literal) means
+    // `resolved` — and everything derived from it — would be recreated on every render,
+    // causing `threadListAdapter` to be a new object each render.  We compute `resolved`
+    // unconditionally and stabilize individual downstream values with primitives/refs.
+    const resolved = resolveTrueFoundryAgentRuntimeOptions(options);
     const { client, agent, gateway } = resolved;
 
     const pendingAgentSpecRef = useRef<AgentSpec | undefined>(
         agent.mode === "draft" ? agent.defaultAgentSpec : undefined,
     );
 
+    // Use stable primitive deps so that threadListAdapter is not recreated every render
+    // when options is an inline object literal (new reference on every parent render).
+    const agentMode = agent.mode;
+    const namedAgentName = agent.mode === "named" ? agent.agentName : undefined;
     const threadListAdapter = useMemo(() => {
-        if (agent.mode === "draft") {
+        if (agentMode === "draft") {
             if (gateway == null) {
                 throw new Error(
                     "Draft agent mode requires a `gateway` TrueFoundryGateway client.",
                 );
             }
+            const draftAgent = agent as Extract<typeof agent, { mode: "draft" }>;
             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]);
+        // eslint-disable-next-line react-hooks/exhaustive-deps
+    }, [agentMode, namedAgentName, client, gateway]);
 
     return useRemoteThreadListRuntime({
         allowNesting: true,