@truefoundry/assistant-ui-runtime 0.1.6-rc.0 → 0.1.7

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 "./private/agentSpec.js";
+import type { AgentSpec } from "./server/types.js";
 import {
     collectPendingApprovals,
     collectPendingToolResponses,
@@ -25,18 +25,19 @@ import {
     buildUserMessageContent,
     extractEditedText,
     parseTurnIdFromMessageId,
+    userMessageContentToText,
 } from "./convertTurnMessages.js";
 import {
     createDraftSessionBridge,
     DRAFT_SESSION_LAST_UPDATED_AT_HEADER,
-} from "./private/draftSessionBridge.js";
+} from "./draft/draftSessionBridge.js";
 import { MCP_AUTH_RESUME_RUN_CUSTOM_KEY } from "./mcpAuth.js";
-import { createTrueFoundryDraftThreadListAdapter } from "./private/truefoundryDraftThreadListAdapter.js";
+import { createTrueFoundryDraftThreadListAdapter } from "./draft/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 "./private/useDraftAgentSpec.js";
+import { useDraftAgentSpec } from "./draft/useDraftAgentSpec.js";
 import { useTrueFoundryAgentMessages } from "./useTrueFoundryAgentMessages.js";
 
 function useTrueFoundryAgentRuntimeImpl(
@@ -66,6 +67,10 @@ function useTrueFoundryAgentRuntimeImpl(
     const isMain = useAuiState(
         (state) => state.threads.mainThreadId === state.threadListItem.id,
     );
+    // On a hard refresh, the URL session runtime mounts before assistant-ui
+    // promotes it to the main thread. Allow that one session to hydrate early.
+    const isInitialSession =
+        sessionId != null && sessionId === options.initialSessionId;
 
     const draftSpec = useDraftAgentSpec({
         draftSessionId,
@@ -119,6 +124,7 @@ function useTrueFoundryAgentRuntimeImpl(
         server,
         sessionId,
         isMain,
+        isInitialSession,
         listEventsConcurrency,
         onError,
         initializeSession,
@@ -148,20 +154,16 @@ function useTrueFoundryAgentRuntimeImpl(
     const downloadSandboxFile = useCallback(
         async (path: string) => {
             if (server.downloadSandboxFile == null) {
-                const error = new Error(
+                throw new Error(
                     "Downloading a sandbox file requires AgentChatServer.downloadSandboxFile.",
                 );
-                onError?.(error);
-                throw error;
             }
             if (sandboxId == null) {
-                const error = new Error("No sandbox is available yet for this session.");
-                onError?.(error);
-                throw error;
+                throw new Error("No sandbox is available yet for this session.");
             }
             return await server.downloadSandboxFile(sandboxId, { path });
         },
-        [server, sandboxId, onError],
+        [server, sandboxId],
     );
 
     const draftExtras = useMemo(() => {
@@ -171,6 +173,7 @@ function useTrueFoundryAgentRuntimeImpl(
         return {
             agentSpec: draftSpec.agentSpec,
             draftSessionId: draftSpec.draftSessionId,
+            isSpecLoading: draftSpec.isSpecLoading,
             isSpecSyncing: draftSpec.isSpecSyncing,
             specError: draftSpec.specError,
             updateAgentSpec: draftSpec.updateAgentSpec,
@@ -192,10 +195,9 @@ function useTrueFoundryAgentRuntimeImpl(
             resumeMcpAuth,
             downloadSandboxFile,
             cancel,
+            // resetFromTurn/branchFromTurn/sendTurn already report via onError.
             resetFromTurn: (turnId: string) =>
-                resetFromTurn(turnId).catch((error) => {
-                    onError?.(error);
-                }),
+                resetFromTurn(turnId).catch(() => undefined),
             reload: retryLoad,
             hasOlderHistory,
             isLoadingOlderHistory,
@@ -224,7 +226,19 @@ function useTrueFoundryAgentRuntimeImpl(
                 return;
             }
 
-            await sendTurn({ userMessage: buildUserMessageContent(message) });
+            const userMessage = buildUserMessageContent(message);
+            await sendTurn({
+                userMessage,
+                // The composer clears before onNew runs. Restore its text only when
+                // the turn failed before turn.created registered it in the backend.
+                onPreTurnFailure: () => {
+                    const text = userMessageContentToText(userMessage);
+                    const composer = aui.thread().composer();
+                    if (text && !composer.getState().text.trim()) {
+                        composer.setText(text);
+                    }
+                },
+            });
         },
         onCancel: async () => {
             await cancel();
@@ -242,12 +256,8 @@ function useTrueFoundryAgentRuntimeImpl(
             }
             const turnId = parseTurnIdFromMessageId(sourceId);
             const editedText = extractEditedText(message);
-            try {
-                await editFromTurn(turnId, editedText);
-            } catch (error) {
-                onError?.(error);
-                throw error;
-            }
+            // editFromTurn/branchFromTurn/sendTurn already report via onError.
+            await editFromTurn(turnId, editedText);
         },
     });
 }