@assistant-ui/react 0.4.8 → 0.5.2

package/dist/index.mjs

@@ -307,13 +307,17 @@ var ThreadProvider = ({
     useCallback2(
       (thread) => {
         const onThreadUpdate = () => {
-          context.useThread.setState(
-            Object.freeze({
-              isRunning: context.useThreadRuntime.getState().isRunning
-            }),
-            true
-          );
-          context.useThreadMessages.setState(thread.messages, true);
+          if (thread.isRunning !== context.useThread.getState().isRunning) {
+            context.useThread.setState(
+              Object.freeze({
+                isRunning: thread.isRunning
+              }),
+              true
+            );
+          }
+          if (thread.messages !== context.useThreadMessages.getState()) {
+            context.useThreadMessages.setState(thread.messages, true);
+          }
         };
         onThreadUpdate();
         return thread.subscribe(onThreadUpdate);
@@ -1169,12 +1173,35 @@ import { memo as memo2 } from "react";
 import { useEffect as useEffect6, useState as useState4 } from "react";
 import { create as create10 } from "zustand";
 import { jsx as jsx13 } from "react/jsx-runtime";
-var DONE_STATUS = { type: "done" };
+var COMPLETE_STATUS = {
+  type: "complete"
+};
+var toContentPartStatus = (message, partIndex, part) => {
+  if (message.role !== "assistant") return COMPLETE_STATUS;
+  const isLastPart = partIndex === Math.max(0, message.content.length - 1);
+  if (part.type !== "tool-call") {
+    if ("reason" in message.status && message.status.reason === "tool-calls" && isLastPart)
+      throw new Error(
+        "Encountered unexpected requires-action status. This is likely an internal bug in assistant-ui."
+      );
+    return isLastPart ? message.status : COMPLETE_STATUS;
+  }
+  if (!!part.result) {
+    return COMPLETE_STATUS;
+  }
+  return message.status;
+};
+var EMPTY_CONTENT = Object.freeze({ type: "text", text: "" });
 var syncContentPart = ({ message }, useContentPart, partIndex) => {
-  const part = message.content[partIndex];
-  if (!part) return;
-  const messageStatus = message.role === "assistant" ? message.status : DONE_STATUS;
-  const status = partIndex === message.content.length - 1 ? messageStatus : DONE_STATUS;
+  let part = message.content[partIndex];
+  if (!part) {
+    if (message.content.length === 0 && partIndex === 0) {
+      part = EMPTY_CONTENT;
+    } else {
+      return;
+    }
+  }
+  const status = toContentPartStatus(message, partIndex, part);
   const currentState = useContentPart.getState();
   if (currentState.part === part && currentState.status === status) return;
   useContentPart.setState(
@@ -1299,7 +1326,7 @@ var ContentPartPrimitiveText = forwardRef7(({ smooth = true, ...rest }, forwarde
     part: { text }
   } = useContentPartText();
   const smoothText = useSmooth(text, smooth);
-  return /* @__PURE__ */ jsx14(Primitive4.span, { "data-status": status, ...rest, ref: forwardedRef, children: smoothText });
+  return /* @__PURE__ */ jsx14(Primitive4.span, { "data-status": status.type, ...rest, ref: forwardedRef, children: smoothText });
 });
 ContentPartPrimitiveText.displayName = "ContentPartPrimitive.Text";
 
@@ -1327,7 +1354,7 @@ ContentPartPrimitiveDisplay.displayName = "ContentPartPrimitive.Display";
 // src/primitives/contentPart/ContentPartInProgress.tsx
 var ContentPartPrimitiveInProgress = ({ children }) => {
   const { useContentPart } = useContentPartContext();
-  const isInProgress = useContentPart((c) => c.status.type === "in_progress");
+  const isInProgress = useContentPart((c) => c.status.type === "running");
   return isInProgress ? children : null;
 };
 ContentPartPrimitiveInProgress.displayName = "ContentPartPrimitive.InProgress";
@@ -1366,10 +1393,16 @@ var MessageContentPartComponent = ({
   const type = part.type;
   switch (type) {
     case "text":
+      if (status.type === "requires-action")
+        throw new Error("Encountered unexpected requires-action status");
       return /* @__PURE__ */ jsx16(Text2, { part, status });
     case "image":
+      if (status.type === "requires-action")
+        throw new Error("Encountered unexpected requires-action status");
       return /* @__PURE__ */ jsx16(Image2, { part, status });
     case "ui":
+      if (status.type === "requires-action")
+        throw new Error("Encountered unexpected requires-action status");
       return /* @__PURE__ */ jsx16(UI, { part, status });
     case "tool-call": {
       const Tool = by_name[part.toolName] || Fallback2;
@@ -1399,7 +1432,7 @@ var MessagePrimitiveContent = ({
   components
 }) => {
   const { useMessage } = useMessageContext();
-  const contentLength = useMessage((s) => s.message.content.length);
+  const contentLength = useMessage((s) => s.message.content.length) || 1;
   return new Array(contentLength).fill(null).map((_, idx) => {
     const partIndex = idx;
     return /* @__PURE__ */ jsx16(
@@ -1703,7 +1736,7 @@ var useThreadViewportAutoScroll = ({
     const div = divRef.current;
     if (!div) return;
     const isAtBottom = useViewport.getState().isAtBottom;
-    const newIsAtBottom = div.scrollHeight - div.scrollTop <= div.clientHeight;
+    const newIsAtBottom = div.scrollHeight - div.scrollTop <= div.clientHeight + 1;
     if (!newIsAtBottom && lastScrollTop.current < div.scrollTop) {
     } else {
       if (newIsAtBottom) {
@@ -1939,6 +1972,40 @@ var ThreadPrimitiveSuggestion = createActionButton(
 // src/runtimes/local/useLocalRuntime.tsx
 import { useInsertionEffect as useInsertionEffect3, useState as useState8 } from "react";
 
+// src/runtimes/core/BaseAssistantRuntime.tsx
+var BaseAssistantRuntime = class {
+  constructor(_thread) {
+    this._thread = _thread;
+    this._thread = _thread;
+  }
+  get thread() {
+    return this._thread;
+  }
+  set thread(thread) {
+    this._thread = thread;
+    this.subscriptionHandler();
+  }
+  _subscriptions = /* @__PURE__ */ new Set();
+  subscribe(callback) {
+    this._subscriptions.add(callback);
+    return () => this._subscriptions.delete(callback);
+  }
+  subscriptionHandler = () => {
+    for (const callback of this._subscriptions) callback();
+  };
+};
+
+// src/internal.ts
+var internal_exports = {};
+__export(internal_exports, {
+  BaseAssistantRuntime: () => BaseAssistantRuntime,
+  MessageRepository: () => MessageRepository,
+  ProxyConfigProvider: () => ProxyConfigProvider,
+  TooltipIconButton: () => TooltipIconButton,
+  generateId: () => generateId,
+  useSmooth: () => useSmooth
+});
+
 // src/utils/idUtils.tsx
 import { customAlphabet } from "nanoid/non-secure";
 var generateId = customAlphabet(
@@ -1948,6 +2015,54 @@ var generateId = customAlphabet(
 var optimisticPrefix = "__optimistic__";
 var generateOptimisticId = () => `${optimisticPrefix}${generateId()}`;
 
+// src/runtimes/edge/converters/fromCoreMessage.ts
+var fromCoreMessages = (message) => {
+  return message.map((message2) => fromCoreMessage(message2));
+};
+var fromCoreMessage = (message, {
+  id = generateId(),
+  status = { type: "complete", reason: "unknown" }
+} = {}) => {
+  const commonProps = {
+    id,
+    createdAt: /* @__PURE__ */ new Date()
+  };
+  const role = message.role;
+  switch (role) {
+    case "assistant":
+      return {
+        ...commonProps,
+        role,
+        content: message.content.map((part) => {
+          if (part.type === "tool-call") {
+            return {
+              ...part,
+              argsText: JSON.stringify(part.args)
+            };
+          }
+          return part;
+        }),
+        status
+      };
+    case "user":
+      return {
+        ...commonProps,
+        role,
+        content: message.content
+      };
+    case "system":
+      return {
+        ...commonProps,
+        role,
+        content: message.content
+      };
+    default: {
+      const unsupportedRole = role;
+      throw new Error(`Unknown message role: ${unsupportedRole}`);
+    }
+  }
+};
+
 // src/runtimes/utils/MessageRepository.tsx
 var findHead = (message) => {
   if (message.next) return findHead(message.next);
@@ -2038,12 +2153,13 @@ var MessageRepository = class {
     do {
       optimisticId = generateOptimisticId();
     } while (this.messages.has(optimisticId));
-    this.addOrUpdateMessage(parentId, {
-      ...message,
-      id: optimisticId,
-      createdAt: /* @__PURE__ */ new Date(),
-      ...message.role === "assistant" ? { status: "in_progress" } : void 0
-    });
+    this.addOrUpdateMessage(
+      parentId,
+      fromCoreMessage(message, {
+        id: optimisticId,
+        status: { type: "running" }
+      })
+    );
     return optimisticId;
   }
   deleteMessage(messageId, replacementId) {
@@ -2110,40 +2226,6 @@ var MessageRepository = class {
   }
 };
 
-// src/runtimes/core/BaseAssistantRuntime.tsx
-var BaseAssistantRuntime = class {
-  constructor(_thread) {
-    this._thread = _thread;
-    this._thread = _thread;
-  }
-  get thread() {
-    return this._thread;
-  }
-  set thread(thread) {
-    this._thread = thread;
-    this.subscriptionHandler();
-  }
-  _subscriptions = /* @__PURE__ */ new Set();
-  subscribe(callback) {
-    this._subscriptions.add(callback);
-    return () => this._subscriptions.delete(callback);
-  }
-  subscriptionHandler = () => {
-    for (const callback of this._subscriptions) callback();
-  };
-};
-
-// src/internal.ts
-var internal_exports = {};
-__export(internal_exports, {
-  BaseAssistantRuntime: () => BaseAssistantRuntime,
-  MessageRepository: () => MessageRepository,
-  ProxyConfigProvider: () => ProxyConfigProvider,
-  TooltipIconButton: () => TooltipIconButton,
-  generateId: () => generateId,
-  useSmooth: () => useSmooth
-});
-
 // src/ui/base/tooltip-icon-button.tsx
 import { forwardRef as forwardRef17 } from "react";
 
@@ -2460,35 +2542,43 @@ var fromLanguageModelMessages = (lm, { mergeRoundtrips }) => {
   return messages;
 };
 
-// src/runtimes/edge/converters/fromCoreMessage.ts
-var fromCoreMessages = (message) => {
-  return message.map((message2) => {
-    return {
-      id: generateId(),
-      createdAt: /* @__PURE__ */ new Date(),
-      ...message2.role === "assistant" ? {
-        status: { type: "done" }
-      } : void 0,
-      ...message2
-    };
-  });
-};
-
 // src/runtimes/edge/converters/toCoreMessages.ts
 var toCoreMessages = (message) => {
-  return message.map((message2) => {
-    return {
-      role: message2.role,
-      content: message2.content.map((part) => {
-        if (part.type === "ui") throw new Error("UI parts are not supported");
-        if (part.type === "tool-call") {
-          const { argsText, ...rest } = part;
-          return rest;
-        }
-        return part;
-      })
-    };
-  });
+  return message.map(toCoreMessage);
+};
+var toCoreMessage = (message) => {
+  const role = message.role;
+  switch (role) {
+    case "assistant":
+      return {
+        role,
+        content: message.content.map((part) => {
+          if (part.type === "ui") throw new Error("UI parts are not supported");
+          if (part.type === "tool-call") {
+            const { argsText, ...rest } = part;
+            return rest;
+          }
+          return part;
+        })
+      };
+    case "user":
+      return {
+        role,
+        content: message.content.map((part) => {
+          if (part.type === "ui") throw new Error("UI parts are not supported");
+          return part;
+        })
+      };
+    case "system":
+      return {
+        role,
+        content: message.content
+      };
+    default: {
+      const unsupportedRole = role;
+      throw new Error(`Unknown message role: ${unsupportedRole}`);
+    }
+  }
 };
 
 // src/runtimes/edge/converters/fromLanguageModelTools.ts
@@ -2958,9 +3048,10 @@ var parsePartialJson = (json) => {
 };
 
 // src/runtimes/edge/streams/runResultStream.ts
-function runResultStream(initialContent) {
+function runResultStream() {
   let message = {
-    content: initialContent
+    content: [],
+    status: { type: "running" }
   };
   const currentToolCall = { toolCallId: "", argsText: "" };
   return new TransformStream({
@@ -3008,7 +3099,13 @@ function runResultStream(initialContent) {
           break;
         }
         case "error": {
-          throw chunk.error;
+          if (chunk.error instanceof Error && chunk.error.name === "AbortError") {
+            message = appendOrUpdateCancel(message);
+            controller.enqueue(message);
+            break;
+          } else {
+            throw chunk.error;
+          }
         }
         default: {
           const unhandledType = chunkType;
@@ -3082,11 +3179,42 @@ var appendOrUpdateToolResult = (message, toolCallId, toolName, result) => {
 };
 var appendOrUpdateFinish = (message, chunk) => {
   const { type, ...rest } = chunk;
+  return {
+    ...message,
+    status: getStatus(chunk),
+    roundtrips: [
+      ...message.roundtrips ?? [],
+      {
+        logprobs: rest.logprobs,
+        usage: rest.usage
+      }
+    ]
+  };
+};
+var getStatus = (chunk) => {
+  if (chunk.finishReason === "tool-calls") {
+    return {
+      type: "requires-action",
+      reason: "tool-calls"
+    };
+  } else if (chunk.finishReason === "stop" || chunk.finishReason === "unknown") {
+    return {
+      type: "complete",
+      reason: chunk.finishReason
+    };
+  } else {
+    return {
+      type: "incomplete",
+      reason: chunk.finishReason
+    };
+  }
+};
+var appendOrUpdateCancel = (message) => {
   return {
     ...message,
     status: {
-      type: "done",
-      ...rest
+      type: "incomplete",
+      reason: "cancelled"
     }
   };
 };
@@ -3186,7 +3314,7 @@ var EdgeChatAdapter = class {
   constructor(options) {
     this.options = options;
   }
-  async roundtrip(initialContent, { messages, abortSignal, config, onUpdate }) {
+  async run({ messages, abortSignal, config, onUpdate }) {
     const result = await fetch(this.options.api, {
       method: "POST",
       headers: {
@@ -3206,46 +3334,14 @@ var EdgeChatAdapter = class {
         `Edge runtime returned status ${result.status}: ${await result.text()}`
       );
     }
-    const stream = result.body.pipeThrough(new TextDecoderStream()).pipeThrough(chunkByLineStream()).pipeThrough(assistantDecoderStream()).pipeThrough(toolResultStream(config.tools)).pipeThrough(runResultStream(initialContent));
-    let message;
+    const stream = result.body.pipeThrough(new TextDecoderStream()).pipeThrough(chunkByLineStream()).pipeThrough(assistantDecoderStream()).pipeThrough(toolResultStream(config.tools)).pipeThrough(runResultStream());
     let update;
     for await (update of asAsyncIterable(stream)) {
-      message = onUpdate(update);
+      onUpdate(update);
     }
     if (update === void 0)
       throw new Error("No data received from Edge Runtime");
-    return [message, update];
-  }
-  async run({ messages, abortSignal, config, onUpdate }) {
-    let roundtripAllowance = this.options.maxToolRoundtrips ?? 1;
-    let usage = {
-      promptTokens: 0,
-      completionTokens: 0
-    };
-    let result;
-    let assistantMessage;
-    do {
-      [assistantMessage, result] = await this.roundtrip(result?.content ?? [], {
-        messages: assistantMessage ? [...messages, assistantMessage] : messages,
-        abortSignal,
-        config,
-        onUpdate
-      });
-      if (result.status?.type === "done") {
-        usage.promptTokens += result.status.usage?.promptTokens ?? 0;
-        usage.completionTokens += result.status.usage?.completionTokens ?? 0;
-      }
-    } while (result.status?.type === "done" && result.status.finishReason === "tool-calls" && result.content.every((c) => c.type !== "tool-call" || !!c.result) && roundtripAllowance-- > 0);
-    if (result.status?.type === "done" && usage.promptTokens > 0) {
-      result = {
-        ...result,
-        status: {
-          ...result.status,
-          usage
-        }
-      };
-    }
-    return result;
+    return update;
   }
 };
 
@@ -3258,30 +3354,10 @@ var useEdgeRuntime = ({
   return useLocalRuntime(adapter, { initialMessages });
 };
 
-// src/runtimes/local/LocalRuntime.tsx
-var LocalRuntime = class extends BaseAssistantRuntime {
-  _proxyConfigProvider;
-  constructor(adapter, options) {
-    const proxyConfigProvider = new ProxyConfigProvider();
-    super(new LocalThreadRuntime(proxyConfigProvider, adapter, options));
-    this._proxyConfigProvider = proxyConfigProvider;
-  }
-  set adapter(adapter) {
-    this.thread.adapter = adapter;
-  }
-  registerModelConfigProvider(provider) {
-    return this._proxyConfigProvider.registerModelConfigProvider(provider);
-  }
-  switchToThread(threadId) {
-    if (threadId) {
-      throw new Error("LocalRuntime does not yet support switching threads");
-    }
-    return this.thread = new LocalThreadRuntime(
-      this._proxyConfigProvider,
-      this.thread.adapter
-    );
-  }
-};
+// src/runtimes/local/shouldContinue.tsx
+var shouldContinue = (result) => result.status?.type === "requires-action" && result.status.reason === "tool-calls" && result.content.every((c) => c.type !== "tool-call" || !!c.result);
+
+// src/runtimes/local/LocalThreadRuntime.tsx
 var CAPABILITIES = Object.freeze({
   edit: true,
   reload: true,
@@ -3292,6 +3368,7 @@ var LocalThreadRuntime = class {
   constructor(configProvider, adapter, options) {
     this.configProvider = configProvider;
     this.adapter = adapter;
+    this.options = options;
     if (options?.initialMessages) {
       let parentId = null;
       const messages = fromCoreMessages(options.initialMessages);
@@ -3335,27 +3412,51 @@ var LocalThreadRuntime = class {
   }
   async startRun(parentId) {
     this.repository.resetHead(parentId);
-    const messages = this.repository.getMessages();
+    const id = generateId();
     let message = {
-      id: generateId(),
+      id,
       role: "assistant",
-      status: { type: "in_progress" },
-      content: [{ type: "text", text: "" }],
+      status: { type: "running" },
+      content: [],
       createdAt: /* @__PURE__ */ new Date()
     };
+    do {
+      message = await this.performRoundtrip(parentId, message);
+    } while (shouldContinue(message));
+  }
+  async performRoundtrip(parentId, message) {
+    const messages = this.repository.getMessages();
     this.abortController?.abort();
     this.abortController = new AbortController();
-    this.repository.addOrUpdateMessage(parentId, { ...message });
-    this.notifySubscribers();
+    const initialContent = message.content;
+    const initialRoundtrips = message.roundtrips;
     const updateMessage = (m) => {
       message = {
         ...message,
-        ...m
+        ...m.content ? { content: [...initialContent, ...m.content ?? []] } : void 0,
+        status: m.status ?? message.status,
+        ...m.roundtrips?.length ? { roundtrips: [...initialRoundtrips ?? [], ...m.roundtrips] } : void 0
       };
       this.repository.addOrUpdateMessage(parentId, message);
       this.notifySubscribers();
-      return message;
     };
+    const maxToolRoundtrips = this.options?.maxToolRoundtrips ?? 1;
+    const toolRoundtrips = message.roundtrips?.length ?? 0;
+    if (toolRoundtrips > maxToolRoundtrips) {
+      updateMessage({
+        status: {
+          type: "incomplete",
+          reason: "tool-calls"
+        }
+      });
+      return message;
+    } else {
+      updateMessage({
+        status: {
+          type: "running"
+        }
+      });
+    }
     try {
       const result = await this.adapter.run({
         messages,
@@ -3363,21 +3464,29 @@ var LocalThreadRuntime = class {
         config: this.configProvider.getModelConfig(),
         onUpdate: updateMessage
       });
-      if (result.status?.type === "in_progress")
+      if (result.status?.type === "running")
         throw new Error(
-          "Unexpected in_progress status returned from ChatModelAdapter"
+          "Unexpected running status returned from ChatModelAdapter"
         );
       this.abortController = null;
-      updateMessage({ status: { type: "done" }, ...result });
-      this.repository.addOrUpdateMessage(parentId, { ...message });
-    } catch (e) {
-      const isAbortError = e instanceof Error && e.name === "AbortError";
-      this.abortController = null;
       updateMessage({
-        status: isAbortError ? { type: "cancelled" } : { type: "error", error: e }
+        status: { type: "complete", reason: "unknown" },
+        ...result
       });
-      if (!isAbortError) throw e;
+    } catch (e) {
+      this.abortController = null;
+      if (e instanceof Error && e.name === "AbortError") {
+        updateMessage({
+          status: { type: "incomplete", reason: "cancelled" }
+        });
+      } else {
+        updateMessage({
+          status: { type: "incomplete", reason: "error", error: e }
+        });
+        throw e;
+      }
     }
+    return message;
   }
   cancelRun() {
     if (!this.abortController) return;
@@ -3392,14 +3501,16 @@ var LocalThreadRuntime = class {
     return () => this._subscriptions.delete(callback);
   }
   addToolResult({ messageId, toolCallId, result }) {
-    const { parentId, message } = this.repository.getMessage(messageId);
+    let { parentId, message } = this.repository.getMessage(messageId);
     if (message.role !== "assistant")
       throw new Error("Tried to add tool result to non-assistant message");
+    let added = false;
     let found = false;
     const newContent = message.content.map((c) => {
       if (c.type !== "tool-call") return c;
       if (c.toolCallId !== toolCallId) return c;
       found = true;
+      if (!c.result) added = true;
       return {
         ...c,
         result
@@ -3407,10 +3518,39 @@ var LocalThreadRuntime = class {
     });
     if (!found)
       throw new Error("Tried to add tool result to non-existing tool call");
-    this.repository.addOrUpdateMessage(parentId, {
+    message = {
       ...message,
       content: newContent
-    });
+    };
+    this.repository.addOrUpdateMessage(parentId, message);
+    if (added && shouldContinue(message)) {
+      this.performRoundtrip(parentId, message);
+    }
+  }
+};
+
+// src/runtimes/local/LocalRuntime.tsx
+var LocalRuntime = class extends BaseAssistantRuntime {
+  _proxyConfigProvider;
+  constructor(adapter, options) {
+    const proxyConfigProvider = new ProxyConfigProvider();
+    super(new LocalThreadRuntime(proxyConfigProvider, adapter, options));
+    this._proxyConfigProvider = proxyConfigProvider;
+  }
+  set adapter(adapter) {
+    this.thread.adapter = adapter;
+  }
+  registerModelConfigProvider(provider) {
+    return this._proxyConfigProvider.registerModelConfigProvider(provider);
+  }
+  switchToThread(threadId) {
+    if (threadId) {
+      throw new Error("LocalRuntime does not yet support switching threads");
+    }
+    return this.thread = new LocalThreadRuntime(
+      this._proxyConfigProvider,
+      this.thread.adapter
+    );
   }
 };
 
@@ -3605,7 +3745,7 @@ var Text = ({ status }) => {
     {
       className: classNames2(
         "aui-text",
-        status.type === "in_progress" && "aui-text-in-progress"
+        status.type === "running" && "aui-text-in-progress"
       ),
       children: /* @__PURE__ */ jsx33(contentPart_exports.Text, {})
     }
@@ -4131,11 +4271,13 @@ export {
   thread_welcome_default as ThreadWelcome,
   user_action_bar_default as UserActionBar,
   user_message_default as UserMessage,
+  fromCoreMessage,
   fromCoreMessages,
   fromLanguageModelMessages,
   fromLanguageModelTools,
   makeAssistantTool,
   makeAssistantToolUI,
+  toCoreMessage,
   toCoreMessages,
   toLanguageModelMessages,
   toLanguageModelTools,