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

package/src/useTrueFoundryAgentMessages.ts

@@ -11,16 +11,17 @@ import type {
     TurnInputItem,
     TurnStateDone,
 } from "truefoundry-gateway-sdk/agents";
-import type { TrueFoundryGateway } from "truefoundry-gateway-sdk";
+import type { PrivateAgentSessionClient } from "truefoundry-gateway-sdk/agents/private";
 
 import { ROOT_THREAD_ID } from "./constants.js";
 import {
     buildEditedUserMessageContent,
-    buildSnapshotBeforeTurnIndex,
+    buildSnapshotBeforeTurn,
     computeGroupRootBaseline,
     extractTurnUserMessageContent,
+    prependOlderSessionHistory,
     projectSessionMessages,
-    resolveGatewayBranchPreviousTurnId,
+    resolveGatewayBranchPreviousTurnIdForTurn,
     rootModelMessageIdsSinceBaseline,
     userMessageContentToText,
     type UserMessageContent,
@@ -67,8 +68,13 @@ export type UseTrueFoundryAgentMessagesOptions = {
     }>;
     /** Maps a thread `remoteId` to the gateway session id used for turns. */
     resolveConversationSessionId?: (remoteId: string) => Promise<string>;
-    /** When set, turns bind to `/agents/sessions/{draftSessionId}/turns` after draft validation. */
-    draftGateway?: TrueFoundryGateway;
+    /** When set, turns bind via PrivateAgentSessionClient.getDraftSession. */
+    privateClient?: PrivateAgentSessionClient;
+    /**
+     * 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 =
@@ -288,21 +294,24 @@ export function useTrueFoundryAgentMessages({
     onError,
     initializeSession,
     resolveConversationSessionId,
-    draftGateway,
+    privateClient,
+    getTurnHeaders,
 }: UseTrueFoundryAgentMessagesOptions) {
     const sessionOptions = useMemo<GetSessionOptions | undefined>(
         () =>
-            draftGateway != null ? { draftGateway } : undefined,
-        [draftGateway],
+            privateClient != null ? { privateClient } : undefined,
+        [privateClient],
     );
 
     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,11 @@ 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 { getGatewayFromPrivateClient } from "./private/getGatewayFromPrivateClient.js";
 import { MCP_AUTH_RESUME_RUN_CUSTOM_KEY } from "./mcpAuth.js";
 import { createTrueFoundryDraftThreadListAdapter } from "./private/truefoundryDraftThreadListAdapter.js";
 import { trueFoundryExtras } from "./truefoundryExtras.js";
@@ -43,7 +47,7 @@ function useTrueFoundryAgentRuntimeImpl(
     const {
         client,
         agent,
-        gateway,
+        privateClient,
         adapters,
         onError,
         listEventsConcurrency,
@@ -51,8 +55,8 @@ function useTrueFoundryAgentRuntimeImpl(
     } = options;
 
     const draftBridgeRef = useRef(
-        agent.mode === "draft" && gateway != null
-            ? createDraftSessionBridge(gateway)
+        agent.mode === "draft" && privateClient != null
+            ? createDraftSessionBridge(privateClient)
             : null,
     );
 
@@ -76,6 +80,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 +108,9 @@ function useTrueFoundryAgentRuntimeImpl(
         messages,
         isRunning,
         isLoading,
+        isLoadingOlderHistory,
+        hasOlderHistory,
+        loadOlderHistory,
         sendTurn,
         cancel,
         respondToToolApproval,
@@ -105,7 +126,8 @@ function useTrueFoundryAgentRuntimeImpl(
         listEventsConcurrency,
         onError,
         initializeSession,
-        draftGateway: agent.mode === "draft" ? gateway : undefined,
+        privateClient: agent.mode === "draft" ? privateClient : undefined,
+        getTurnHeaders: agent.mode === "draft" ? getTurnHeaders : undefined,
     });
 
     if (agent.mode === "draft" && draftSpec.agentSpec != null) {
@@ -130,9 +152,9 @@ function useTrueFoundryAgentRuntimeImpl(
 
     const downloadSandboxFile = useCallback(
         async (path: string) => {
-            if (gateway == null) {
+            if (privateClient == null) {
                 const error = new Error(
-                    "Downloading a sandbox file requires a `gateway` TrueFoundryGateway client.",
+                    "Downloading a sandbox file requires a `privateClient` PrivateAgentSessionClient.",
                 );
                 onError?.(error);
                 throw error;
@@ -142,10 +164,12 @@ function useTrueFoundryAgentRuntimeImpl(
                 onError?.(error);
                 throw error;
             }
+            // Same gateway.agents.downloadSandboxFile path as before, via the wrapped client.
+            const gateway = getGatewayFromPrivateClient(privateClient);
             const response = await gateway.agents.downloadSandboxFile(sandboxId, { path });
             return await response.blob();
         },
-        [gateway, sandboxId, onError],
+        [privateClient, sandboxId, onError],
     );
 
     const draftExtras = useMemo(() => {
@@ -181,6 +205,9 @@ function useTrueFoundryAgentRuntimeImpl(
                     onError?.(error);
                 }),
             reload: retryLoad,
+            hasOlderHistory,
+            isLoadingOlderHistory,
+            loadOlderHistory,
             draft: draftExtras,
         }),
         unstable_enableToolInvocations: true,
@@ -240,7 +267,7 @@ export function useTrueFoundryAgentRuntime(options: UseTrueFoundryAgentRuntimeOp
     // 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 { client, agent, privateClient } = resolved;
 
     const pendingAgentSpecRef = useRef<AgentSpec | undefined>(
         agent.mode === "draft" ? agent.defaultAgentSpec : undefined,
@@ -252,14 +279,14 @@ export function useTrueFoundryAgentRuntime(options: UseTrueFoundryAgentRuntimeOp
     const namedAgentName = agent.mode === "named" ? agent.agentName : undefined;
     const threadListAdapter = useMemo(() => {
         if (agentMode === "draft") {
-            if (gateway == null) {
+            if (privateClient == null) {
                 throw new Error(
-                    "Draft agent mode requires a `gateway` TrueFoundryGateway client.",
+                    "Draft agent mode requires a `privateClient` PrivateAgentSessionClient.",
                 );
             }
             const draftAgent = agent as Extract<typeof agent, { mode: "draft" }>;
             return createTrueFoundryDraftThreadListAdapter({
-                gateway,
+                privateClient,
                 defaultAgentSpec: draftAgent.defaultAgentSpec,
                 getAgentSpec: () => pendingAgentSpecRef.current ?? draftAgent.defaultAgentSpec,
             });
@@ -269,7 +296,7 @@ export function useTrueFoundryAgentRuntime(options: UseTrueFoundryAgentRuntimeOp
             agentName: namedAgentName!,
         });
         // eslint-disable-next-line react-hooks/exhaustive-deps
-    }, [agentMode, namedAgentName, client, gateway]);
+    }, [agentMode, namedAgentName, client, privateClient]);
 
     return useRemoteThreadListRuntime({
         allowNesting: true,

package/src/agentSessionModule.d.ts

@@ -1,37 +0,0 @@
-declare module "truefoundry-gateway-sdk/dist/esm/agents/AgentSession.mjs" {
-    import type { TrueFoundryGateway } from "truefoundry-gateway-sdk";
-    import type * as TrueFoundryGatewayApi from "truefoundry-gateway-sdk";
-    import type { PreparedTurn } from "truefoundry-gateway-sdk/agents";
-    import type * as core from "truefoundry-gateway-sdk/core";
-
-    export class AgentSession implements TrueFoundryGatewayApi.Session {
-        readonly id: string;
-        readonly agentName: string;
-        readonly title?: string;
-        readonly createdBySubject: TrueFoundryGatewayApi.Subject;
-        readonly createdAt: string;
-        readonly updatedAt: string;
-        constructor(session: TrueFoundryGatewayApi.Session, client: TrueFoundryGateway);
-        prepareTurn(opts?: {
-            input?: TrueFoundryGatewayApi.TurnInputItem[];
-            previousTurnId?: TrueFoundryGatewayApi.PreviousTurnIdInput;
-        }): PreparedTurn;
-        listTurns(
-            opts?: TrueFoundryGatewayApi.agents.SessionsListTurnsRequest,
-            requestOptions?: unknown,
-        ): Promise<core.Page<unknown, TrueFoundryGatewayApi.ListTurnsResponse>>;
-        /**
-         * Paginated session events across turns (newest first); subscribe to a
-         * running turn for live events.
-         *
-         * @param opts.lastTurnId - Newest turn in the listing window (initial load only).
-         *   Omit to use the session last turn. Running-turn events are excluded.
-         */
-        listEvents(
-            opts?: TrueFoundryGatewayApi.agents.SessionsListEventsRequest,
-            requestOptions?: unknown,
-        ): Promise<core.Page<TrueFoundryGatewayApi.SessionEventItem, TrueFoundryGatewayApi.ListSessionEventsResponse>>;
-        getTurn(opts: { turnId: string }, requestOptions?: unknown): Promise<unknown>;
-        cancel(requestOptions?: unknown): Promise<void>;
-    }
-}