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

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,17 @@ import {
     extractEditedText,
     parseTurnIdFromMessageId,
 } from "./convertTurnMessages.js";
-import { createDraftSessionBridge } from "./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 "./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 +66,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,
@@ -73,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(),
@@ -87,6 +107,9 @@ function useTrueFoundryAgentRuntimeImpl(
         messages,
         isRunning,
         isLoading,
+        isLoadingOlderHistory,
+        hasOlderHistory,
+        loadOlderHistory,
         sendTurn,
         cancel,
         respondToToolApproval,
@@ -94,13 +117,16 @@ function useTrueFoundryAgentRuntimeImpl(
         resumeRun,
         editFromTurn,
         resetFromTurn,
+        retryLoad,
     } = useTrueFoundryAgentMessages({
         client,
         sessionId,
+        isMain,
         listEventsConcurrency,
         onError,
         initializeSession,
         draftGateway: agent.mode === "draft" ? gateway : undefined,
+        getTurnHeaders: agent.mode === "draft" ? getTurnHeaders : undefined,
     });
 
     if (agent.mode === "draft" && draftSpec.agentSpec != null) {
@@ -175,6 +201,10 @@ function useTrueFoundryAgentRuntimeImpl(
                 resetFromTurn(turnId).catch((error) => {
                     onError?.(error);
                 }),
+            reload: retryLoad,
+            hasOlderHistory,
+            isLoadingOlderHistory,
+            loadOlderHistory,
             draft: draftExtras,
         }),
         unstable_enableToolInvocations: true,
@@ -228,34 +258,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,