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

package/src/useTrueFoundryAgentMessages.ts

@@ -16,11 +16,12 @@ import type { TrueFoundryGateway } from "truefoundry-gateway-sdk";
 import { ROOT_THREAD_ID } from "./constants.js";
 import {
     buildEditedUserMessageContent,
-    buildSnapshotBeforeTurnIndex,
+    buildSnapshotBeforeTurn,
     computeGroupRootBaseline,
     extractTurnUserMessageContent,
+    prependOlderSessionHistory,
     projectSessionMessages,
-    resolveGatewayBranchPreviousTurnId,
+    resolveGatewayBranchPreviousTurnIdForTurn,
     rootModelMessageIdsSinceBaseline,
     userMessageContentToText,
     type UserMessageContent,
@@ -57,6 +58,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<{
@@ -67,10 +70,24 @@ export type UseTrueFoundryAgentMessagesOptions = {
     resolveConversationSessionId?: (remoteId: string) => Promise<string>;
     /** When set, turns bind to `/agents/sessions/{draftSessionId}/turns` after draft validation. */
     draftGateway?: TrueFoundryGateway;
+    /**
+     * Optional per-turn headers for createTurn. Invoked once per `sendTurn` after
+     * the session is resolved; return value is forwarded to `turn.execute`.
+     */
+    getTurnHeaders?: () => Promise<Record<string, string> | undefined>;
 };
 
 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,11 +289,13 @@ function resolveTurnInput(
 export function useTrueFoundryAgentMessages({
     client,
     sessionId,
+    isMain,
     listEventsConcurrency,
     onError,
     initializeSession,
     resolveConversationSessionId,
     draftGateway,
+    getTurnHeaders,
 }: UseTrueFoundryAgentMessagesOptions) {
     const sessionOptions = useMemo<GetSessionOptions | undefined>(
         () =>
@@ -287,15 +306,28 @@ export function useTrueFoundryAgentMessages({
     const [snapshot, setSnapshot] = useState<SessionSnapshot>(createEmptySessionSnapshot);
     const [isRunning, setIsRunning] = useState(false);
     const [isLoading, setIsLoading] = useState(false);
+    const [isLoadingOlderHistory, setIsLoadingOlderHistory] = useState(false);
+    const [loadRetryTrigger, setLoadRetryTrigger] = useState(0);
 
     const snapshotRef = useRef(snapshot);
     snapshotRef.current = snapshot;
+    const loadOlderInflightRef = useRef<Promise<void> | null>(null);
+
+    const onErrorRef = useRef(onError);
+    onErrorRef.current = onError;
+    const resolveConversationSessionIdRef = useRef(resolveConversationSessionId);
+    resolveConversationSessionIdRef.current = resolveConversationSessionId;
+    const initializeSessionRef = useRef(initializeSession);
+    initializeSessionRef.current = initializeSession;
+    const getTurnHeadersRef = useRef(getTurnHeaders);
+    getTurnHeadersRef.current = getTurnHeaders;
 
     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 +350,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 +458,7 @@ export function useTrueFoundryAgentMessages({
                 });
             return run;
         },
-        [applyStreamUpdate, onError],
+        [onError],
     );
 
     const load = useCallback(async () => {
@@ -401,24 +467,47 @@ 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();
+        loadOlderInflightRef.current = null;
+        createdAtByMessageIdRef.current = new Map();
+        setSnapshot(createEmptySessionSnapshot());
         setIsLoading(true);
+        setIsLoadingOlderHistory(false);
 
         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 +517,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,19 +564,22 @@ 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 turnHeaders = await getTurnHeadersRef.current?.();
+            const streamHeaders =
+                turnHeaders != null ? { headers: turnHeaders } : {};
             const isContinuation =
                 "inputs" in options ||
                 ("resumeMcpAuth" in options && options.resumeMcpAuth === true);
@@ -492,34 +596,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(
@@ -528,7 +655,7 @@ export function useTrueFoundryAgentMessages({
                         return streamTurnContent(
                             session,
                             snapshotRef.current.fold,
-                            { inputs: options.inputs },
+                            { inputs: options.inputs, ...streamHeaders },
                             signal,
                             groupRootBaseline,
                         );
@@ -537,7 +664,7 @@ export function useTrueFoundryAgentMessages({
                         return streamTurnContent(
                             session,
                             snapshotRef.current.fold,
-                            { resumeMcpAuth: true },
+                            { resumeMcpAuth: true, ...streamHeaders },
                             signal,
                             groupRootBaseline,
                         );
@@ -550,6 +677,7 @@ export function useTrueFoundryAgentMessages({
                             ...(options.previousTurnId !== undefined
                                 ? { previousTurnId: options.previousTurnId }
                                 : {}),
+                            ...streamHeaders,
                         },
                         signal,
                         groupRootBaseline,
@@ -559,7 +687,7 @@ export function useTrueFoundryAgentMessages({
                 isContinuation,
             );
         },
-        [client, initializeSession, resolveConversationSessionId, runStream, sessionId, sessionOptions],
+        [client, runStream, sessionId, sessionOptions],
     );
 
     const cancel = useCallback(async () => {
@@ -569,7 +697,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 +709,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 +726,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(
@@ -671,37 +799,38 @@ export function useTrueFoundryAgentMessages({
             const committed = commitActiveStream(snapshotRef.current);
             setSnapshot(committed);
 
-            const turnIndex = findTurnIndex(committed, turnId);
-
             await cancel();
 
             const conversationSessionId = await resolveActiveSessionId(
                 activeSessionId,
-                resolveConversationSessionId,
+                resolveConversationSessionIdRef.current,
             );
             const session = await getSession(client, conversationSessionId, sessionOptions);
-            const previousTurnId = await resolveGatewayBranchPreviousTurnId(
+            const previousTurnId = await resolveGatewayBranchPreviousTurnIdForTurn(
                 session,
-                turnIndex,
+                turnId,
             );
-            const rewound = await buildSnapshotBeforeTurnIndex(
+            const rewound = await buildSnapshotBeforeTurn(
                 session,
-                turnIndex,
+                turnId,
                 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 +866,75 @@ export function useTrueFoundryAgentMessages({
         [branchFromTurn],
     );
 
+    const retryLoad = useCallback(() => {
+        setLoadRetryTrigger((n) => n + 1);
+    }, []);
+
+    const hasOlderHistory = snapshot.historyPagination?.hasOlder === true;
+
+    const loadOlderHistory = useCallback(async () => {
+        if (sessionId == null || isMain === false) {
+            return;
+        }
+        if (loadOlderInflightRef.current != null) {
+            return loadOlderInflightRef.current;
+        }
+
+        const current = snapshotRef.current;
+        if (current.historyPagination?.hasOlder !== true) {
+            return;
+        }
+        if (current.historyPagination.olderPageToken == null) {
+            return;
+        }
+
+        const generation = loadGenerationRef.current;
+        setIsLoadingOlderHistory(true);
+
+        const run = (async () => {
+            try {
+                const conversationSessionId = await resolveActiveSessionId(
+                    sessionId,
+                    resolveConversationSessionIdRef.current,
+                );
+                const session = await getSession(
+                    client,
+                    conversationSessionId,
+                    sessionOptions,
+                );
+                const next = await prependOlderSessionHistory(
+                    session,
+                    snapshotRef.current,
+                );
+                if (generation !== loadGenerationRef.current) {
+                    return;
+                }
+                setSnapshot(next);
+            } catch (error) {
+                if (generation === loadGenerationRef.current) {
+                    onErrorRef.current?.(error);
+                }
+                throw error;
+            } finally {
+                if (generation === loadGenerationRef.current) {
+                    setIsLoadingOlderHistory(false);
+                }
+                loadOlderInflightRef.current = null;
+            }
+        })();
+
+        loadOlderInflightRef.current = run;
+        return run;
+    }, [client, isMain, sessionId, sessionOptions]);
+
     return {
         messages,
         isRunning,
         isLoading,
-        load,
+        isLoadingOlderHistory,
+        hasOlderHistory,
+        loadOlderHistory,
+        retryLoad,
         sendTurn,
         cancel,
         respondToToolApproval,