@assistant-ui/react 0.12.10 → 0.12.12

package/src/legacy-runtime/runtime-cores/assistant-transport/useToolInvocations.ts

@@ -41,23 +41,21 @@ export type ToolExecutionStatus =
   | { type: "executing" }
   | { type: "interrupt"; payload: { type: "human"; payload: unknown } };
 
+type ToolState = {
+  argsText: string;
+  hasResult: boolean;
+  argsComplete: boolean;
+  streamToolCallId: string;
+  controller: ToolCallStreamController;
+};
+
 export function useToolInvocations({
   state,
   getTools,
   onResult,
   setToolStatuses,
 }: UseToolInvocationsParams) {
-  const lastToolStates = useRef<
-    Record<
-      string,
-      {
-        argsText: string;
-        hasResult: boolean;
-        argsComplete: boolean;
-        controller: ToolCallStreamController;
-      }
-    >
-  >({});
+  const lastToolStates = useRef<Record<string, ToolState>>({});
 
   const humanInputRef = useRef<
     Map<
@@ -72,26 +70,76 @@ export function useToolInvocations({
   const acRef = useRef<AbortController>(new AbortController());
   const executingCountRef = useRef(0);
   const settledResolversRef = useRef<Array<() => void>>([]);
+  const toolCallIdAliasesRef = useRef<Map<string, string>>(new Map());
+  const ignoredResultToolCallIdsRef = useRef<Set<string>>(new Set());
+  const rewriteCounterRef = useRef(0);
+
+  const getLogicalToolCallId = (toolCallId: string) => {
+    return toolCallIdAliasesRef.current.get(toolCallId) ?? toolCallId;
+  };
+
+  const shouldIgnoreAndCleanupResult = (toolCallId: string) => {
+    if (!ignoredResultToolCallIdsRef.current.has(toolCallId)) return false;
+    ignoredResultToolCallIdsRef.current.delete(toolCallId);
+    toolCallIdAliasesRef.current.delete(toolCallId);
+    return true;
+  };
+
+  const getWrappedTools = () => {
+    const tools = getTools();
+    if (!tools) return undefined;
+
+    return Object.fromEntries(
+      Object.entries(tools).map(([name, tool]) => {
+        const execute = tool.execute;
+        const streamCall = tool.streamCall;
+
+        const wrappedTool = {
+          ...tool,
+          ...(execute !== undefined && {
+            execute: (
+              ...[args, context]: Parameters<NonNullable<typeof execute>>
+            ) =>
+              execute(args, {
+                ...context,
+                toolCallId: getLogicalToolCallId(context.toolCallId),
+              }),
+          }),
+          ...(streamCall !== undefined && {
+            streamCall: (
+              ...[reader, context]: Parameters<NonNullable<typeof streamCall>>
+            ) =>
+              streamCall(reader, {
+                ...context,
+                toolCallId: getLogicalToolCallId(context.toolCallId),
+              }),
+          }),
+        } as Tool;
+        return [name, wrappedTool];
+      }),
+    ) as Record<string, Tool>;
+  };
 
   const [controller] = useState(() => {
     const [stream, controller] = createAssistantStreamController();
     const transform = unstable_toolResultStream(
-      getTools,
+      getWrappedTools,
       () => acRef.current?.signal ?? new AbortController().signal,
       (toolCallId: string, payload: unknown) => {
+        const logicalToolCallId = getLogicalToolCallId(toolCallId);
         return new Promise<unknown>((resolve, reject) => {
           // Reject previous human input request if it exists
-          const previous = humanInputRef.current.get(toolCallId);
+          const previous = humanInputRef.current.get(logicalToolCallId);
           if (previous) {
             previous.reject(
               new Error("Human input request was superseded by a new request"),
             );
           }
 
-          humanInputRef.current.set(toolCallId, { resolve, reject });
+          humanInputRef.current.set(logicalToolCallId, { resolve, reject });
           setToolStatuses((prev) => ({
             ...prev,
-            [toolCallId]: {
+            [logicalToolCallId]: {
               type: "interrupt",
               payload: { type: "human", payload },
             },
@@ -100,17 +148,25 @@ export function useToolInvocations({
       },
       {
         onExecutionStart: (toolCallId: string) => {
+          if (ignoredResultToolCallIdsRef.current.has(toolCallId)) {
+            return;
+          }
+          const logicalToolCallId = getLogicalToolCallId(toolCallId);
           executingCountRef.current++;
           setToolStatuses((prev) => ({
             ...prev,
-            [toolCallId]: { type: "executing" },
+            [logicalToolCallId]: { type: "executing" },
           }));
         },
         onExecutionEnd: (toolCallId: string) => {
+          if (ignoredResultToolCallIdsRef.current.has(toolCallId)) {
+            return;
+          }
+          const logicalToolCallId = getLogicalToolCallId(toolCallId);
           executingCountRef.current--;
           setToolStatuses((prev) => {
             const next = { ...prev };
-            delete next[toolCallId];
+            delete next[logicalToolCallId];
             return next;
           });
           // Resolve any waiting abort promises when all tools have settled
@@ -128,13 +184,22 @@ export function useToolInvocations({
         new WritableStream({
           write(chunk) {
             if (chunk.type === "result") {
-              // the tool call result was already set by the backend
-              if (lastToolStates.current[chunk.meta.toolCallId]?.hasResult)
+              if (shouldIgnoreAndCleanupResult(chunk.meta.toolCallId)) {
                 return;
+              }
+
+              const logicalToolCallId = getLogicalToolCallId(
+                chunk.meta.toolCallId,
+              );
+              if (logicalToolCallId !== chunk.meta.toolCallId) {
+                toolCallIdAliasesRef.current.delete(chunk.meta.toolCallId);
+              }
+              // the tool call result was already set by the backend
+              if (lastToolStates.current[logicalToolCallId]?.hasResult) return;
 
               onResult({
                 type: "add-tool-result",
-                toolCallId: chunk.meta.toolCallId,
+                toolCallId: logicalToolCallId,
                 toolName: chunk.meta.toolName,
                 result: chunk.result,
                 isError: chunk.isError,
@@ -152,6 +217,33 @@ export function useToolInvocations({
   const isInitialState = useRef(true);
 
   useEffect(() => {
+    const createToolState = ({
+      controller,
+      streamToolCallId,
+    }: {
+      controller: ToolCallStreamController;
+      streamToolCallId: string;
+    }): ToolState => ({
+      argsText: "",
+      hasResult: false,
+      argsComplete: false,
+      streamToolCallId,
+      controller,
+    });
+
+    const setToolState = (toolCallId: string, state: ToolState) => {
+      lastToolStates.current[toolCallId] = state;
+      return state;
+    };
+
+    const patchToolState = (
+      toolCallId: string,
+      state: ToolState,
+      patch: Partial<ToolState>,
+    ) => {
+      return setToolState(toolCallId, { ...state, ...patch });
+    };
+
     const processMessages = (
       messages: readonly (typeof state.messages)[number][],
     ) => {
@@ -166,17 +258,21 @@ export function useToolInvocations({
               }
               let lastState = lastToolStates.current[content.toolCallId];
               if (!lastState) {
+                toolCallIdAliasesRef.current.set(
+                  content.toolCallId,
+                  content.toolCallId,
+                );
                 const toolCallController = controller.addToolCallPart({
                   toolName: content.toolName,
                   toolCallId: content.toolCallId,
                 });
-                lastState = {
-                  argsText: "",
-                  hasResult: false,
-                  argsComplete: false,
-                  controller: toolCallController,
-                };
-                lastToolStates.current[content.toolCallId] = lastState;
+                lastState = setToolState(
+                  content.toolCallId,
+                  createToolState({
+                    controller: toolCallController,
+                    streamToolCallId: content.toolCallId,
+                  }),
+                );
               }
 
               if (content.argsText !== lastState.argsText) {
@@ -200,17 +296,50 @@ export function useToolInvocations({
                       isArgsTextComplete(content.argsText)
                     ) {
                       lastState.controller.argsText.close();
-                      lastToolStates.current[content.toolCallId] = {
+                      patchToolState(content.toolCallId, lastState, {
                         argsText: content.argsText,
-                        hasResult: lastState.hasResult,
                         argsComplete: true,
-                        controller: lastState.controller,
-                      };
+                      });
                       return; // Continue to next content part
                     }
-                    throw new Error(
-                      `Tool call argsText can only be appended, not updated: ${content.argsText} does not start with ${lastState.argsText}`,
+                    if (process.env.NODE_ENV !== "production") {
+                      console.warn(
+                        "argsText rewrote previous snapshot, restarting tool args stream:",
+                        {
+                          previous: lastState.argsText,
+                          next: content.argsText,
+                          toolCallId: content.toolCallId,
+                        },
+                      );
+                    }
+
+                    ignoredResultToolCallIdsRef.current.add(
+                      lastState.streamToolCallId,
+                    );
+                    lastState.controller.argsText.close();
+
+                    const streamToolCallId = `${content.toolCallId}:rewrite:${rewriteCounterRef.current++}`;
+                    toolCallIdAliasesRef.current.set(
+                      streamToolCallId,
+                      content.toolCallId,
                     );
+                    const toolCallController = controller.addToolCallPart({
+                      toolName: content.toolName,
+                      toolCallId: streamToolCallId,
+                    });
+                    if (process.env.NODE_ENV !== "production") {
+                      console.warn("started replacement stream tool call", {
+                        toolCallId: content.toolCallId,
+                        streamToolCallId,
+                      });
+                    }
+                    lastState = setToolState(content.toolCallId, {
+                      ...createToolState({
+                        controller: toolCallController,
+                        streamToolCallId,
+                      }),
+                      hasResult: lastState.hasResult,
+                    });
                   }
 
                   const argsTextDelta = content.argsText.slice(
@@ -223,22 +352,18 @@ export function useToolInvocations({
                     lastState.controller.argsText.close();
                   }
 
-                  lastToolStates.current[content.toolCallId] = {
+                  patchToolState(content.toolCallId, lastState, {
                     argsText: content.argsText,
-                    hasResult: lastState.hasResult,
                     argsComplete: shouldClose,
-                    controller: lastState.controller,
-                  };
+                  });
                 }
               }
 
               if (content.result !== undefined && !lastState.hasResult) {
-                lastToolStates.current[content.toolCallId] = {
+                patchToolState(content.toolCallId, lastState, {
                   hasResult: true,
                   argsComplete: true,
-                  argsText: lastState.argsText,
-                  controller: lastState.controller,
-                };
+                });
 
                 lastState.controller.setResponse(
                   new ToolResponse({
@@ -287,8 +412,12 @@ export function useToolInvocations({
 
   return {
     reset: () => {
-      void abort();
       isInitialState.current = true;
+      void abort().finally(() => {
+        toolCallIdAliasesRef.current.clear();
+        ignoredResultToolCallIdsRef.current.clear();
+        rewriteCounterRef.current = 0;
+      });
     },
     abort,
     resume: (toolCallId: string, payload: unknown) => {

package/src/legacy-runtime/runtime-cores/assistant-transport/utils.ts

@@ -1,37 +1,5 @@
-import { Tool, toToolsJSONSchema, type ToolJSONSchema } from "assistant-stream";
-
-/**
- * @deprecated Use `toToolsJSONSchema` from `assistant-stream` instead.
- */
-export function toAISDKTools(
-  tools: Record<string, Tool>,
-): Record<string, ToolJSONSchema> {
-  return toToolsJSONSchema(tools, { filter: () => true });
-}
-
-/**
- * @deprecated Use `toToolsJSONSchema` from `assistant-stream` instead, which includes filtering by default.
- */
-export function getEnabledTools(
-  tools: Record<string, Tool>,
-): Record<string, Tool> {
-  return Object.fromEntries(
-    Object.entries(tools).filter(
-      ([, tool]) => !tool.disabled && tool.type !== "backend",
-    ),
-  );
-}
-
-export async function createRequestHeaders(
-  headersValue:
-    | Record<string, string>
-    | Headers
-    | (() => Promise<Record<string, string> | Headers>),
-): Promise<Headers> {
-  const resolvedHeaders =
-    typeof headersValue === "function" ? await headersValue() : headersValue;
-
-  const headers = new Headers(resolvedHeaders);
-  headers.set("Content-Type", "application/json");
-  return headers;
-}
+export {
+  toAISDKTools,
+  getEnabledTools,
+  createRequestHeaders,
+} from "@assistant-ui/core";

package/src/legacy-runtime/runtime-cores/composer/BaseComposerRuntimeCore.ts

@@ -0,0 +1 @@
+export { BaseComposerRuntimeCore } from "@assistant-ui/core/internal";

package/src/legacy-runtime/runtime-cores/composer/DefaultEditComposerRuntimeCore.ts

@@ -0,0 +1 @@
+export { DefaultEditComposerRuntimeCore } from "@assistant-ui/core/internal";

package/src/legacy-runtime/runtime-cores/composer/DefaultThreadComposerRuntimeCore.ts

@@ -0,0 +1 @@
+export { DefaultThreadComposerRuntimeCore } from "@assistant-ui/core/internal";

package/src/legacy-runtime/runtime-cores/core/AssistantRuntimeCore.ts

@@ -0,0 +1 @@
+export type { AssistantRuntimeCore } from "@assistant-ui/core";

package/src/legacy-runtime/runtime-cores/core/BaseAssistantRuntimeCore.ts

@@ -0,0 +1 @@
+export { BaseAssistantRuntimeCore } from "@assistant-ui/core/internal";

package/src/legacy-runtime/runtime-cores/core/BaseThreadRuntimeCore.ts

@@ -0,0 +1 @@
+export { BaseThreadRuntimeCore } from "@assistant-ui/core/internal";

package/src/legacy-runtime/runtime-cores/core/ComposerRuntimeCore.ts

@@ -0,0 +1,6 @@
+export type {
+  ComposerRuntimeEventType,
+  DictationState,
+  ComposerRuntimeCore,
+  ThreadComposerRuntimeCore,
+} from "@assistant-ui/core";

package/src/legacy-runtime/runtime-cores/core/ThreadListRuntimeCore.ts

@@ -0,0 +1,4 @@
+export type {
+  ThreadListItemCoreState,
+  ThreadListRuntimeCore,
+} from "@assistant-ui/core";

package/src/legacy-runtime/runtime-cores/core/ThreadRuntimeCore.ts

@@ -0,0 +1,13 @@
+export type {
+  RuntimeCapabilities,
+  AddToolResultOptions,
+  ResumeToolCallOptions,
+  SubmitFeedbackOptions,
+  ThreadSuggestion,
+  SpeechState,
+  SubmittedFeedback,
+  ThreadRuntimeEventType,
+  StartRunConfig,
+  ResumeRunConfig,
+  ThreadRuntimeCore,
+} from "@assistant-ui/core";

package/src/legacy-runtime/runtime-cores/core/index.ts

@@ -1,7 +1,7 @@
+// Re-export from @assistant-ui/core
 export type {
   AddToolResultOptions,
   SubmitFeedbackOptions,
   ThreadSuggestion,
-} from "./ThreadRuntimeCore";
-
-export type { DictationState } from "./ComposerRuntimeCore";
+  DictationState,
+} from "@assistant-ui/core";

package/src/legacy-runtime/runtime-cores/external-store/ExternalStoreAdapter.ts

@@ -0,0 +1,6 @@
+export type {
+  ExternalStoreAdapter,
+  ExternalStoreMessageConverter,
+  ExternalStoreThreadListAdapter,
+  ExternalStoreThreadData,
+} from "@assistant-ui/core";

package/src/legacy-runtime/runtime-cores/external-store/ExternalStoreRuntimeCore.ts

@@ -0,0 +1 @@
+export { ExternalStoreRuntimeCore } from "@assistant-ui/core/internal";

package/src/legacy-runtime/runtime-cores/external-store/ExternalStoreThreadListRuntimeCore.ts

@@ -0,0 +1,2 @@
+export { ExternalStoreThreadListRuntimeCore } from "@assistant-ui/core/internal";
+export type { ExternalStoreThreadFactory } from "@assistant-ui/core/internal";

package/src/legacy-runtime/runtime-cores/external-store/ExternalStoreThreadRuntimeCore.ts

@@ -0,0 +1,4 @@
+export {
+  ExternalStoreThreadRuntimeCore,
+  hasUpcomingMessage,
+} from "@assistant-ui/core/internal";

package/src/legacy-runtime/runtime-cores/external-store/ThreadMessageConverter.ts

@@ -1,23 +1,2 @@
-import { ThreadMessage } from "../../../types";
-
-export type ConverterCallback<TIn> = (
-  cache: ThreadMessage | undefined,
-  message: TIn,
-  idx: number,
-) => ThreadMessage;
-
-export class ThreadMessageConverter {
-  private readonly cache = new WeakMap<WeakKey, ThreadMessage>();
-
-  convertMessages<TIn extends WeakKey>(
-    messages: readonly TIn[],
-    converter: ConverterCallback<TIn>,
-  ): ThreadMessage[] {
-    return messages.map((m, idx) => {
-      const cached = this.cache.get(m);
-      const newMessage = converter(cached, m, idx);
-      this.cache.set(m, newMessage);
-      return newMessage;
-    });
-  }
-}
+export { ThreadMessageConverter } from "@assistant-ui/core/internal";
+export type { ConverterCallback } from "@assistant-ui/core/internal";

package/src/legacy-runtime/runtime-cores/external-store/ThreadMessageLike.ts

@@ -0,0 +1,2 @@
+export type { ThreadMessageLike } from "@assistant-ui/core";
+export { fromThreadMessageLike } from "@assistant-ui/core/internal";

package/src/legacy-runtime/runtime-cores/external-store/auto-status.ts

@@ -0,0 +1 @@
+export { isAutoStatus, getAutoStatus } from "@assistant-ui/core/internal";

package/src/legacy-runtime/runtime-cores/external-store/{createMessageConverter.tsx → createMessageConverter.ts}

@@ -1,7 +1,7 @@
 "use client";
+import type { ThreadMessage } from "@assistant-ui/core";
 import { ThreadState } from "../../runtime";
 import { useAui, useAuiState } from "@assistant-ui/store";
-import { ThreadMessage } from "../../../types";
 import {
   useExternalMessageConverter,
   convertExternalMessages,

package/src/legacy-runtime/runtime-cores/external-store/{external-message-converter.tsx → external-message-converter.ts}

@@ -5,14 +5,19 @@ import { ThreadMessageConverter } from "./ThreadMessageConverter";
 import {
   getExternalStoreMessages,
   symbolInnerMessage,
+  bindExternalStoreMessage,
 } from "./getExternalStoreMessage";
 import { fromThreadMessageLike, ThreadMessageLike } from "./ThreadMessageLike";
 import { getAutoStatus, isAutoStatus } from "./auto-status";
-import { ThreadMessage, ToolCallMessagePart } from "../../../types";
 import { ToolExecutionStatus } from "../assistant-transport/useToolInvocations";
 import { ReadonlyJSONValue } from "assistant-stream/utils";
-import { generateErrorMessageId } from "../../../utils/idUtils";
-import { ThreadAssistantMessage } from "../../../types/AssistantTypes";
+import { generateErrorMessageId } from "@assistant-ui/core/internal";
+import type {
+  MessageTiming,
+  ThreadAssistantMessage,
+  ThreadMessage,
+  ToolCallMessagePart,
+} from "@assistant-ui/core";
 
 export namespace useExternalMessageConverter {
   export type Message =
@@ -34,6 +39,7 @@ export namespace useExternalMessageConverter {
   export type Metadata = {
     readonly toolStatuses?: Record<string, ToolExecutionStatus>;
     readonly error?: ReadonlyJSONValue;
+    readonly messageTiming?: Record<string, MessageTiming>;
   };
 
   export type Callback<T> = (
@@ -160,6 +166,10 @@ const joinExternalMessages = (
                 };
               }
 
+              if (output.metadata.timing) {
+                assistantMessage.metadata.timing = output.metadata.timing;
+              }
+
               if (output.metadata.submittedFeedback) {
                 assistantMessage.metadata.submittedFeedback =
                   output.metadata.submittedFeedback;
@@ -211,7 +221,7 @@ const joinExternalMessages = (
   return assistantMessage;
 };
 
-const chunkExternalMessages = <T,>(
+const chunkExternalMessages = <T>(
   callbackResults: CallbackResult<T>[],
   joinStrategy?: "concat-content" | "none",
 ) => {
@@ -267,23 +277,22 @@ const chunkExternalMessages = <T,>(
 function createErrorAssistantMessage(
   error: ReadonlyJSONValue,
 ): ThreadAssistantMessage {
-  return Object.assign<ThreadAssistantMessage, { [symbolInnerMessage]: [] }>(
-    {
-      id: generateErrorMessageId(),
-      role: "assistant",
-      content: [],
-      status: { type: "incomplete", reason: "error", error },
-      createdAt: new Date(),
-      metadata: {
-        unstable_state: null,
-        unstable_annotations: [],
-        unstable_data: [],
-        custom: {},
-        steps: [],
-      },
+  const msg: ThreadAssistantMessage = {
+    id: generateErrorMessageId(),
+    role: "assistant",
+    content: [],
+    status: { type: "incomplete", reason: "error", error },
+    createdAt: new Date(),
+    metadata: {
+      unstable_state: null,
+      unstable_annotations: [],
+      unstable_data: [],
+      custom: {},
+      steps: [],
     },
-    { [symbolInnerMessage]: [] },
-  );
+  };
+  bindExternalStoreMessage(msg, []);
+  return msg;
 }
 
 export const convertExternalMessages = <T extends WeakKey>(
@@ -322,7 +331,7 @@ export const convertExternalMessages = <T extends WeakKey>(
       idx.toString(),
       autoStatus,
     );
-    (newMessage as any)[symbolInnerMessage] = message.inputs;
+    bindExternalStoreMessage(newMessage, message.inputs);
     return newMessage;
   });
 
@@ -430,14 +439,12 @@ export const useExternalMessageConverter = <T extends WeakKey>({
           idx.toString(),
           autoStatus,
         );
-        (newMessage as any)[symbolInnerMessage] = message.inputs;
+        bindExternalStoreMessage(newMessage, message.inputs);
         return newMessage;
       },
     );
 
-    (threadMessages as unknown as { [symbolInnerMessage]: T[] })[
-      symbolInnerMessage
-    ] = messages;
+    bindExternalStoreMessage(threadMessages, messages);
 
     if (state.metadata.error) {
       const lastMessage = threadMessages.at(-1);

package/src/legacy-runtime/runtime-cores/external-store/getExternalStoreMessage.ts

@@ -0,0 +1,6 @@
+export {
+  getExternalStoreMessage,
+  getExternalStoreMessages,
+  bindExternalStoreMessage,
+} from "@assistant-ui/core";
+export { symbolInnerMessage } from "@assistant-ui/core/internal";

package/src/legacy-runtime/runtime-cores/external-store/index.ts

@@ -1,15 +1,19 @@
+// Re-export from @assistant-ui/core
+export type { ThreadMessageLike } from "@assistant-ui/core";
+export {
+  getExternalStoreMessage,
+  getExternalStoreMessages,
+  bindExternalStoreMessage,
+} from "@assistant-ui/core";
 export type {
   ExternalStoreAdapter,
   ExternalStoreMessageConverter,
   ExternalStoreThreadListAdapter,
   ExternalStoreThreadData,
-} from "./ExternalStoreAdapter";
-export type { ThreadMessageLike } from "./ThreadMessageLike";
+} from "@assistant-ui/core";
+
+// React-specific (stay in react)
 export { useExternalStoreRuntime } from "./useExternalStoreRuntime";
-export {
-  getExternalStoreMessage,
-  getExternalStoreMessages,
-} from "./getExternalStoreMessage";
 export {
   useExternalMessageConverter,
   convertExternalMessages as unstable_convertExternalMessages,

package/src/legacy-runtime/runtime-cores/external-store/{useExternalStoreRuntime.tsx → useExternalStoreRuntime.ts}

@@ -9,7 +9,7 @@ import {
 } from "../../runtime/AssistantRuntime";
 import { useRuntimeAdapters } from "../adapters/RuntimeAdapterProvider";
 
-export const useExternalStoreRuntime = <T,>(
+export const useExternalStoreRuntime = <T>(
   store: ExternalStoreAdapter<T>,
 ): AssistantRuntime => {
   const [runtime] = useState(() => new ExternalStoreRuntimeCore(store));

package/src/legacy-runtime/runtime-cores/index.ts

@@ -5,7 +5,6 @@ export * from "./external-store";
 export * from "./local";
 export * from "./remote-thread-list";
 
-export {
-  ExportedMessageRepository,
-  type ExportedMessageRepositoryItem,
-} from "./utils/MessageRepository";
+// Re-export from @assistant-ui/core
+export type { ExportedMessageRepositoryItem } from "@assistant-ui/core";
+export { ExportedMessageRepository } from "@assistant-ui/core";