@truefoundry/assistant-ui-runtime 0.1.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,
@@ -69,6 +70,11 @@ 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 =
@@ -289,6 +295,7 @@ export function useTrueFoundryAgentMessages({
     initializeSession,
     resolveConversationSessionId,
     draftGateway,
+    getTurnHeaders,
 }: UseTrueFoundryAgentMessagesOptions) {
     const sessionOptions = useMemo<GetSessionOptions | undefined>(
         () =>
@@ -299,10 +306,12 @@ 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;
@@ -310,6 +319,8 @@ export function useTrueFoundryAgentMessages({
     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);
@@ -477,9 +488,11 @@ export function useTrueFoundryAgentMessages({
         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(
@@ -564,6 +577,9 @@ export function useTrueFoundryAgentMessages({
                 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);
@@ -639,7 +655,7 @@ export function useTrueFoundryAgentMessages({
                         return streamTurnContent(
                             session,
                             snapshotRef.current.fold,
-                            { inputs: options.inputs },
+                            { inputs: options.inputs, ...streamHeaders },
                             signal,
                             groupRootBaseline,
                         );
@@ -648,7 +664,7 @@ export function useTrueFoundryAgentMessages({
                         return streamTurnContent(
                             session,
                             snapshotRef.current.fold,
-                            { resumeMcpAuth: true },
+                            { resumeMcpAuth: true, ...streamHeaders },
                             signal,
                             groupRootBaseline,
                         );
@@ -661,6 +677,7 @@ export function useTrueFoundryAgentMessages({
                             ...(options.previousTurnId !== undefined
                                 ? { previousTurnId: options.previousTurnId }
                                 : {}),
+                            ...streamHeaders,
                         },
                         signal,
                         groupRootBaseline,
@@ -782,8 +799,6 @@ export function useTrueFoundryAgentMessages({
             const committed = commitActiveStream(snapshotRef.current);
             setSnapshot(committed);
 
-            const turnIndex = findTurnIndex(committed, turnId);
-
             await cancel();
 
             const conversationSessionId = await resolveActiveSessionId(
@@ -791,13 +806,13 @@ export function useTrueFoundryAgentMessages({
                 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();
@@ -855,10 +870,70 @@ export function useTrueFoundryAgentMessages({
         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,
+        isLoadingOlderHistory,
+        hasOlderHistory,
+        loadOlderHistory,
         retryLoad,
         sendTurn,
         cancel,

package/src/useTrueFoundryAgentRuntime.ts

@@ -26,7 +26,10 @@ import {
     extractEditedText,
     parseTurnIdFromMessageId,
 } from "./convertTurnMessages.js";
-import { createDraftSessionBridge } from "./private/draftSessionBridge.js";
+import {
+    createDraftSessionBridge,
+    DRAFT_SESSION_LAST_UPDATED_AT_HEADER,
+} from "./private/draftSessionBridge.js";
 import { MCP_AUTH_RESUME_RUN_CUSTOM_KEY } from "./mcpAuth.js";
 import { createTrueFoundryDraftThreadListAdapter } from "./private/truefoundryDraftThreadListAdapter.js";
 import { trueFoundryExtras } from "./truefoundryExtras.js";
@@ -76,6 +79,20 @@ function useTrueFoundryAgentRuntimeImpl(
         onError,
     });
 
+    const takeTurnHeaderTimestampRef = useRef(draftSpec.takeTurnHeaderTimestamp);
+    takeTurnHeaderTimestampRef.current = draftSpec.takeTurnHeaderTimestamp;
+
+    const getTurnHeaders = useCallback(async () => {
+        if (agent.mode !== "draft") {
+            return undefined;
+        }
+        const updatedAt = await takeTurnHeaderTimestampRef.current();
+        if (updatedAt == null) {
+            return undefined;
+        }
+        return { [DRAFT_SESSION_LAST_UPDATED_AT_HEADER]: updatedAt };
+    }, [agent.mode]);
+
     const aui = useAui();
     const initializeSession = useCallback(
         () => aui.threadListItem().initialize(),
@@ -90,6 +107,9 @@ function useTrueFoundryAgentRuntimeImpl(
         messages,
         isRunning,
         isLoading,
+        isLoadingOlderHistory,
+        hasOlderHistory,
+        loadOlderHistory,
         sendTurn,
         cancel,
         respondToToolApproval,
@@ -106,6 +126,7 @@ function useTrueFoundryAgentRuntimeImpl(
         onError,
         initializeSession,
         draftGateway: agent.mode === "draft" ? gateway : undefined,
+        getTurnHeaders: agent.mode === "draft" ? getTurnHeaders : undefined,
     });
 
     if (agent.mode === "draft" && draftSpec.agentSpec != null) {
@@ -181,6 +202,9 @@ function useTrueFoundryAgentRuntimeImpl(
                     onError?.(error);
                 }),
             reload: retryLoad,
+            hasOlderHistory,
+            isLoadingOlderHistory,
+            loadOlderHistory,
             draft: draftExtras,
         }),
         unstable_enableToolInvocations: true,