@assistant-ui/react 0.7.85 → 0.7.87

package/src/runtimes/core/BaseThreadRuntimeCore.tsx

@@ -167,13 +167,14 @@ export abstract class BaseThreadRuntimeCore implements ThreadRuntimeCore {
     }
   }
 
+  // TODO import()/export() on external store doesn't make much sense
   public export() {
     return this.repository.export();
   }
 
   public import(data: ExportedMessageRepository) {
     this.ensureInitialized();
-
+    this.repository.clear();
     this.repository.import(data);
     this._notifySubscribers();
   }

package/src/runtimes/external-store/useExternalStoreRuntime.tsx

@@ -4,7 +4,6 @@ import { useEffect, useMemo, useState } from "react";
 import { ExternalStoreRuntimeCore } from "./ExternalStoreRuntimeCore";
 import { ExternalStoreAdapter } from "./ExternalStoreAdapter";
 import { AssistantRuntimeImpl } from "../../api/AssistantRuntime";
-import { ThreadRuntimeImpl } from "../../api/ThreadRuntime";
 import { useRuntimeAdapters } from "../adapters/RuntimeAdapterProvider";
 
 export const useExternalStoreRuntime = <T,>(store: ExternalStoreAdapter<T>) => {
@@ -21,8 +20,5 @@ export const useExternalStoreRuntime = <T,>(store: ExternalStoreAdapter<T>) => {
     return runtime.registerModelContextProvider(modelContext);
   }, [modelContext, runtime]);
 
-  return useMemo(
-    () => AssistantRuntimeImpl.create(runtime, ThreadRuntimeImpl),
-    [runtime],
-  );
+  return useMemo(() => new AssistantRuntimeImpl(runtime), [runtime]);
 };

package/src/runtimes/index.ts

@@ -5,3 +5,5 @@ export * from "./edge";
 export * from "./external-store";
 export * from "./local";
 export * from "./remote-thread-list";
+
+export { ExportedMessageRepository } from "./utils/MessageRepository";

package/src/runtimes/local/LocalRuntimeCore.tsx

@@ -1,28 +1,9 @@
-import type { CoreMessage } from "../../types/AssistantTypes";
 import { BaseAssistantRuntimeCore } from "../core/BaseAssistantRuntimeCore";
 import { LocalThreadRuntimeCore } from "./LocalThreadRuntimeCore";
 import { LocalRuntimeOptionsBase } from "./LocalRuntimeOptions";
 import { LocalThreadListRuntimeCore } from "./LocalThreadListRuntimeCore";
 import { ExportedMessageRepository } from "../utils/MessageRepository";
 import { ThreadMessageLike } from "../external-store";
-import { fromThreadMessageLike } from "../external-store/ThreadMessageLike";
-import { generateId } from "../../internal";
-import { getAutoStatus } from "../external-store/auto-status";
-
-const getExportFromInitialMessages = (
-  initialMessages: readonly ThreadMessageLike[],
-): ExportedMessageRepository => {
-  const messages = initialMessages.map((i, idx) => {
-    const isLast = idx === initialMessages.length - 1;
-    return fromThreadMessageLike(i, generateId(), getAutoStatus(isLast, false));
-  });
-  return {
-    messages: messages.map((m, idx) => ({
-      parentId: messages[idx - 1]?.id ?? null,
-      message: m,
-    })),
-  };
-};
 
 export class LocalRuntimeCore extends BaseAssistantRuntimeCore {
   public readonly threads;
@@ -45,20 +26,7 @@ export class LocalRuntimeCore extends BaseAssistantRuntimeCore {
     if (initialMessages) {
       this.threads
         .getMainThreadRuntimeCore()
-        .import(getExportFromInitialMessages(initialMessages));
+        .import(ExportedMessageRepository.fromArray(initialMessages));
     }
   }
-
-  public reset({
-    initialMessages,
-  }: {
-    initialMessages?: readonly CoreMessage[] | undefined;
-  } = {}) {
-    this.threads.switchToNewThread();
-    if (!initialMessages) return;
-
-    this.threads
-      .getMainThreadRuntimeCore()
-      .import(getExportFromInitialMessages(initialMessages));
-  }
 }

package/src/runtimes/local/useLocalRuntime.tsx

@@ -4,37 +4,10 @@ import { useEffect, useMemo, useState } from "react";
 import type { ChatModelAdapter } from "./ChatModelAdapter";
 import { LocalRuntimeCore } from "./LocalRuntimeCore";
 import { LocalRuntimeOptions } from "./LocalRuntimeOptions";
-import {
-  AssistantRuntime,
-  AssistantRuntimeImpl,
-} from "../../api/AssistantRuntime";
-import { ThreadRuntimeImpl } from "../../internal";
 import { useRuntimeAdapters } from "../adapters/RuntimeAdapterProvider";
 import { useRemoteThreadListRuntime } from "../remote-thread-list/useRemoteThreadListRuntime";
 import { useCloudThreadListAdapter } from "../remote-thread-list/adapter/cloud";
-
-export type LocalRuntime = AssistantRuntime & {
-  reset: (options?: Parameters<LocalRuntimeCore["reset"]>[0]) => void;
-};
-
-class LocalRuntimeImpl extends AssistantRuntimeImpl implements LocalRuntime {
-  private constructor(private core: LocalRuntimeCore) {
-    super(core, ThreadRuntimeImpl);
-  }
-
-  public override __internal_bindMethods() {
-    super.__internal_bindMethods();
-    this.reset = this.reset.bind(this);
-  }
-
-  public reset(options?: Parameters<LocalRuntimeCore["reset"]>[0]) {
-    this.core.reset(options);
-  }
-
-  public static override create(_core: LocalRuntimeCore): LocalRuntime {
-    return new LocalRuntimeImpl(_core);
-  }
-}
+import { AssistantRuntimeImpl } from "../../internal";
 
 const useLocalThreadRuntime = (
   adapter: ChatModelAdapter,
@@ -65,7 +38,7 @@ const useLocalThreadRuntime = (
     return runtime.registerModelContextProvider(modelContext);
   }, [modelContext, runtime]);
 
-  return useMemo(() => LocalRuntimeImpl.create(runtime), [runtime]);
+  return useMemo(() => new AssistantRuntimeImpl(runtime), [runtime]);
 };
 
 export const useLocalRuntime = (

package/src/runtimes/remote-thread-list/useRemoteThreadListRuntime.tsx

@@ -34,5 +34,5 @@ export const useRemoteThreadListRuntime = (
     runtime.threads.__internal_setOptions(options);
     runtime.threads.__internal_load();
   }, [runtime, options]);
-  return useMemo(() => AssistantRuntimeImpl.create(runtime), [runtime]);
+  return useMemo(() => new AssistantRuntimeImpl(runtime), [runtime]);
 };

package/src/runtimes/utils/MessageRepository.tsx

@@ -1,6 +1,9 @@
 import type { CoreMessage, ThreadMessage } from "../../types/AssistantTypes";
-import { generateOptimisticId } from "../../utils/idUtils";
+import { generateId, generateOptimisticId } from "../../utils/idUtils";
 import { fromCoreMessage } from "../edge/converters/fromCoreMessage";
+import { ThreadMessageLike } from "../external-store";
+import { getAutoStatus } from "../external-store/auto-status";
+import { fromThreadMessageLike } from "../external-store/ThreadMessageLike";
 
 type RepositoryParent = {
   children: string[];
@@ -18,13 +21,30 @@ export type ExportedMessageRepositoryItem = {
   parentId: string | null;
 };
 
-export interface ExportedMessageRepository {
+export type ExportedMessageRepository = {
   headId?: string | null;
   messages: Array<{
     message: ThreadMessage;
     parentId: string | null;
   }>;
-}
+};
+
+export const ExportedMessageRepository = {
+  fromArray: (
+    messages: readonly ThreadMessageLike[],
+  ): ExportedMessageRepository => {
+    const conv = messages.map((m) =>
+      fromThreadMessageLike(m, generateId(), getAutoStatus(false, false)),
+    );
+
+    return {
+      messages: conv.map((m, idx) => ({
+        parentId: idx > 0 ? conv[idx - 1]!.id : null,
+        message: m,
+      })),
+    };
+  },
+};
 
 const findHead = (
   message: RepositoryMessage | RepositoryParent,
@@ -287,6 +307,16 @@ export class MessageRepository {
     this._messages.dirty();
   }
 
+  clear(): void {
+    this.messages.clear();
+    this.head = null;
+    this.root = {
+      children: [],
+      next: null,
+    };
+    this._messages.dirty();
+  }
+
   export(): ExportedMessageRepository {
     const exportItems: ExportedMessageRepository["messages"] = [];
 

package/src/tailwindcss/index.ts

@@ -1,9 +1,9 @@
 import plugin from "tailwindcss/plugin.js";
-import baseComponentsCSS from "../../dist/styles/tailwindcss/base-components.css.json";
-import threadCSS from "../../dist/styles/tailwindcss/thread.css.json";
-import modalCSS from "../../dist/styles/tailwindcss/modal.css.json";
-import defaultThemeCSS from "../../dist/styles/themes/default.css.json";
-import shadcnExtrasCSS from "../../dist/styles/themes/shadcn-extras.css.json";
+import baseComponentsCSS from "../../dist/styles/tailwindcss/base-components.css.json" with { type: "json" };
+import threadCSS from "../../dist/styles/tailwindcss/thread.css.json" with { type: "json" };
+import modalCSS from "../../dist/styles/tailwindcss/modal.css.json" with { type: "json" };
+import defaultThemeCSS from "../../dist/styles/themes/default.css.json" with { type: "json" };
+import shadcnExtrasCSS from "../../dist/styles/themes/shadcn-extras.css.json" with { type: "json" };
 
 type AssistantTailwindPluginColors = {
   border: string;
@@ -77,10 +77,7 @@ const auiPlugin = plugin.withOptions<AssisstantTailwindPluginOptions>(
         addComponents(modalCSS);
       }
     },
-  ({
-    colors = {},
-    shadcn = false,
-  } = {}) => {
+  ({ colors = {}, shadcn = false } = {}) => {
     const prefix = !shadcn ? "--aui-" : "--";
     return {
       safelist: [{ pattern: /^aui-/ }],

package/src/utils/json/fix-json.ts

@@ -40,7 +40,7 @@ type State =
 // Please note that invalid JSON is not considered/covered, because it
 // is assumed that the resulting JSON will be processed by a standard
 // JSON parser that will detect any invalid JSON.
-export function fixJson(input: string): string {
+export function fixJson(input: string): [string, number] {
   const stack: State[] = ["ROOT"];
   let lastValidIndex = -1;
   let literalStart: number | null = null;
@@ -372,6 +372,7 @@ export function fixJson(input: string): string {
   }
 
   let result = input.slice(0, lastValidIndex + 1);
+  let partialCount = 0;
 
   for (let i = stack.length - 1; i >= 0; i--) {
     const state = stack[i];
@@ -379,6 +380,7 @@ export function fixJson(input: string): string {
     switch (state) {
       case "INSIDE_STRING": {
         result += '"';
+        partialCount++;
         break;
       }
 
@@ -389,6 +391,7 @@ export function fixJson(input: string): string {
       case "INSIDE_OBJECT_BEFORE_VALUE":
       case "INSIDE_OBJECT_AFTER_VALUE": {
         result += "}";
+        partialCount++;
         break;
       }
 
@@ -396,6 +399,7 @@ export function fixJson(input: string): string {
       case "INSIDE_ARRAY_AFTER_COMMA":
       case "INSIDE_ARRAY_AFTER_VALUE": {
         result += "]";
+        partialCount++;
         break;
       }
 
@@ -413,5 +417,5 @@ export function fixJson(input: string): string {
     }
   }
 
-  return result;
+  return [result, partialCount];
 }

package/src/utils/json/parse-partial-json.ts

@@ -1,14 +1,74 @@
 import sjson from "secure-json-parse";
 import { fixJson } from "./fix-json";
+import {
+  ContentPartStatus,
+  ToolCallContentPartStatus,
+} from "../../types/AssistantTypes";
+import { useContentPart } from "../../context";
 
+const PARTIAL_JSON_COUNT_SYMBOL = Symbol("partial-json-count");
 export const parsePartialJson = (json: string) => {
   try {
     return sjson.parse(json);
   } catch {
     try {
-      return sjson.parse(fixJson(json));
+      const [fixedJson, partialCount] = fixJson(json);
+      const res = sjson.parse(fixedJson);
+      res[PARTIAL_JSON_COUNT_SYMBOL] = partialCount;
+      return res;
     } catch {
       return undefined;
     }
   }
 };
+
+const COMPLETE_STATUS = Object.freeze({ type: "complete" });
+
+const getFieldStatus = (
+  lastState: ContentPartStatus,
+  args: unknown,
+  fieldPath: string[],
+  partialCount: number,
+): ContentPartStatus => {
+  if (fieldPath.length === 0) return lastState;
+  if (typeof args !== "object" || args === null) return COMPLETE_STATUS;
+
+  const path = fieldPath.at(-1)!;
+
+  // If the expected property does not exist, mark as incomplete
+  if (!Object.prototype.hasOwnProperty.call(args, path)) {
+    return lastState;
+  }
+
+  const argsKeys = Object.keys(args);
+  const isLast = argsKeys[argsKeys.length - 1] === path;
+  if (!isLast) return COMPLETE_STATUS;
+
+  return getFieldStatus(
+    lastState,
+    (args as Record<string, unknown>)[path],
+    fieldPath.slice(0, -1),
+    partialCount - 1,
+  );
+};
+
+const getToolArgsFieldStatus = (
+  status: ToolCallContentPartStatus,
+  args: Record<string, unknown>,
+  fieldPath: string[],
+): ContentPartStatus => {
+  const partialCount = (args as any)[PARTIAL_JSON_COUNT_SYMBOL] ?? 0;
+  if (partialCount === 0) return COMPLETE_STATUS;
+
+  const lastState: ContentPartStatus =
+    status.type !== "requires-action" ? status : COMPLETE_STATUS;
+
+  return getFieldStatus(lastState, args, fieldPath, partialCount);
+};
+
+export const useToolArgsFieldStatus = (fieldPath: string[]) => {
+  return useContentPart((p) => {
+    if (p.type !== "tool-call") throw new Error("not a tool call");
+    return getToolArgsFieldStatus(p.status, p.args, fieldPath);
+  });
+};