@assistant-ui/react 0.11.34 → 0.11.36

package/src/legacy-runtime/runtime-cores/assistant-transport/useAssistantTransportRuntime.tsx

@@ -13,6 +13,7 @@ import { useState, useRef, useMemo } from "react";
 import {
   AssistantMessageAccumulator,
   DataStreamDecoder,
+  AssistantTransportDecoder,
   unstable_createInitialMessage as createInitialMessage,
 } from "assistant-stream";
 import {
@@ -30,12 +31,14 @@ import { ToolExecutionStatus, useToolInvocations } from "./useToolInvocations";
 import { toAISDKTools, getEnabledTools, createRequestHeaders } from "./utils";
 import { useRemoteThreadListRuntime } from "../remote-thread-list/useRemoteThreadListRuntime";
 import { InMemoryThreadListAdapter } from "../remote-thread-list/adapter/in-memory";
-import { useAssistantApi } from "../../../context/react";
+import { useAssistantApi, useAssistantState } from "../../../context/react";
+import { UserExternalState } from "../../../augmentations";
 
 const symbolAssistantTransportExtras = Symbol("assistant-transport-extras");
 type AssistantTransportExtras = {
   [symbolAssistantTransportExtras]: true;
   sendCommand: (command: AssistantTransportCommand) => void;
+  state: UserExternalState;
 };
 
 const asAssistantTransportExtras = (
@@ -63,6 +66,18 @@ export const useAssistantTransportSendCommand = () => {
   };
 };
 
+export function useAssistantTransportState(): UserExternalState;
+export function useAssistantTransportState<T>(
+  selector: (state: UserExternalState) => T,
+): T;
+export function useAssistantTransportState<T>(
+  selector: (state: UserExternalState) => T = (t) => t as T,
+): T | UserExternalState {
+  return useAssistantState(({ thread }) =>
+    selector(asAssistantTransportExtras(thread.extras).state),
+  );
+}
+
 const useAssistantTransportThreadRuntime = <T,>(
   options: AssistantTransportOptions<T>,
 ): AssistantRuntime => {
@@ -114,21 +129,26 @@ const useAssistantTransportThreadRuntime = <T,>(
         throw new Error("Response body is null");
       }
 
+      // Select decoder based on protocol option
+      const protocol = options.protocol ?? "data-stream";
+      const decoder =
+        protocol === "assistant-transport"
+          ? new AssistantTransportDecoder()
+          : new DataStreamDecoder();
+
       let err: string | undefined;
-      const stream = response.body
-        .pipeThrough(new DataStreamDecoder())
-        .pipeThrough(
-          new AssistantMessageAccumulator({
-            initialMessage: createInitialMessage({
-              unstable_state:
-                (agentStateRef.current as ReadonlyJSONValue) ?? null,
-            }),
-            throttle: isResume,
-            onError: (error) => {
-              err = error;
-            },
+      const stream = response.body.pipeThrough(decoder).pipeThrough(
+        new AssistantMessageAccumulator({
+          initialMessage: createInitialMessage({
+            unstable_state:
+              (agentStateRef.current as ReadonlyJSONValue) ?? null,
           }),
-        );
+          throttle: isResume,
+          onError: (error) => {
+            err = error;
+          },
+        }),
+      );
 
       let markedDelivered = false;
 
@@ -206,6 +226,7 @@ const useAssistantTransportThreadRuntime = <T,>(
       sendCommand: (command: AssistantTransportCommand) => {
         commandQueue.enqueue(command);
       },
+      state: agentStateRef.current as UserExternalState,
     } satisfies AssistantTransportExtras,
     onNew: async (message: AppendMessage): Promise<void> => {
       if (message.role !== "user")

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

@@ -7,6 +7,7 @@ export type AssistantRuntimeCore = {
   readonly threads: ThreadListRuntimeCore;
 
   registerModelContextProvider: (provider: ModelContextProvider) => Unsubscribe;
+  getModelContextProvider: () => ModelContextProvider;
 
   /**
    * EXPERIMENTAL: A component that is rendered inside the AssistantRuntimeProvider.

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

@@ -13,4 +13,8 @@ export abstract class BaseAssistantRuntimeCore implements AssistantRuntimeCore {
   ): Unsubscribe {
     return this._contextProvider.registerModelContextProvider(provider);
   }
+
+  public getModelContextProvider(): ModelContextProvider {
+    return this._contextProvider;
+  }
 }

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

@@ -212,8 +212,10 @@ export class ExternalStoreThreadListRuntimeCore
     onDelete(threadId);
   }
 
-  public initialize(): never {
-    throw new Error("Method not implemented.");
+  public initialize(
+    threadId: string,
+  ): Promise<{ remoteId: string; externalId: string | undefined }> {
+    return Promise.resolve({ remoteId: threadId, externalId: undefined });
   }
 
   public generateTitle(): never {

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

@@ -162,7 +162,13 @@ export class ExternalStoreThreadRuntimeCore
             if (!store.convertMessage) return m;
 
             const isLast = idx === store.messages!.length - 1;
-            const autoStatus = getAutoStatus(isLast, isRunning, false, false);
+            const autoStatus = getAutoStatus(
+              isLast,
+              isRunning,
+              false,
+              false,
+              undefined,
+            );
 
             if (
               cache &&

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

@@ -1,34 +1,67 @@
+import { ReadonlyJSONValue } from "assistant-stream/utils";
 import { MessageStatus } from "../../../types";
 
-const AUTO_STATUS_RUNNING = Object.freeze({ type: "running" });
-const AUTO_STATUS_COMPLETE = Object.freeze({
-  type: "complete",
-  reason: "unknown",
-});
+const symbolAutoStatus = Symbol("autoStatus");
 
-const AUTO_STATUS_PENDING = Object.freeze({
-  type: "requires-action",
-  reason: "tool-calls",
-});
+const AUTO_STATUS_RUNNING = Object.freeze(
+  Object.assign({ type: "running" as const }, { [symbolAutoStatus]: true }),
+);
+const AUTO_STATUS_COMPLETE = Object.freeze(
+  Object.assign(
+    {
+      type: "complete" as const,
+      reason: "unknown" as const,
+    },
+    { [symbolAutoStatus]: true },
+  ),
+);
 
-const AUTO_STATUS_INTERRUPT = Object.freeze({
-  type: "requires-action",
-  reason: "interrupt",
-});
+const AUTO_STATUS_PENDING = Object.freeze(
+  Object.assign(
+    {
+      type: "requires-action" as const,
+      reason: "tool-calls" as const,
+    },
+    { [symbolAutoStatus]: true },
+  ),
+);
+
+const AUTO_STATUS_INTERRUPT = Object.freeze(
+  Object.assign(
+    {
+      type: "requires-action" as const,
+      reason: "interrupt" as const,
+    },
+    { [symbolAutoStatus]: true },
+  ),
+);
 
 export const isAutoStatus = (status: MessageStatus) =>
-  status === AUTO_STATUS_RUNNING || status === AUTO_STATUS_COMPLETE;
+  (status as any)[symbolAutoStatus] === true;
 
 export const getAutoStatus = (
   isLast: boolean,
   isRunning: boolean,
   hasInterruptedToolCalls: boolean,
   hasPendingToolCalls: boolean,
-) =>
-  isLast && isRunning
+  error?: ReadonlyJSONValue,
+): MessageStatus => {
+  if (isLast && error) {
+    return Object.assign(
+      {
+        type: "incomplete" as const,
+        reason: "error" as const,
+        error: error,
+      },
+      { [symbolAutoStatus]: true },
+    );
+  }
+
+  return isLast && isRunning
     ? AUTO_STATUS_RUNNING
     : hasInterruptedToolCalls
       ? AUTO_STATUS_INTERRUPT
       : hasPendingToolCalls
         ? AUTO_STATUS_PENDING
         : AUTO_STATUS_COMPLETE;
+};

package/src/legacy-runtime/runtime-cores/external-store/external-message-converter.tsx

@@ -10,6 +10,7 @@ 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";
 
 export namespace useExternalMessageConverter {
   export type Message =
@@ -30,6 +31,7 @@ export namespace useExternalMessageConverter {
 
   export type Metadata = {
     readonly toolStatuses?: Record<string, ToolExecutionStatus>;
+    readonly error?: ReadonlyJSONValue;
   };
 
   export type Callback<T> = (
@@ -263,6 +265,7 @@ export const convertExternalMessages = <T extends WeakKey>(
       isRunning,
       hasSuspendedToolCalls,
       hasPendingToolCalls,
+      isLast ? metadata.error : undefined,
     );
     const newMessage = fromThreadMessageLike(
       joined,
@@ -348,6 +351,7 @@ export const useExternalMessageConverter = <T extends WeakKey>({
           isRunning,
           hasSuspendedToolCalls,
           hasPendingToolCalls,
+          isLast ? state.metadata.error : undefined,
         );
 
         if (

package/src/legacy-runtime/runtime-cores/local/LocalThreadListRuntimeCore.tsx

@@ -104,8 +104,10 @@ export class LocalThreadListRuntimeCore
     throw new Error("Method not implemented.");
   }
 
-  public initialize(): never {
-    throw new Error("Method not implemented.");
+  public initialize(
+    threadId: string,
+  ): Promise<{ remoteId: string; externalId: string | undefined }> {
+    return Promise.resolve({ remoteId: threadId, externalId: undefined });
   }
 
   public generateTitle(): never {

package/src/legacy-runtime/runtime-cores/utils/MessageRepository.tsx

@@ -67,7 +67,7 @@ export const ExportedMessageRepository = {
       fromThreadMessageLike(
         m,
         generateId(),
-        getAutoStatus(false, false, false, false),
+        getAutoStatus(false, false, false, false, undefined),
       ),
     );
 

package/src/model-context/index.ts

@@ -19,5 +19,14 @@ export { tool } from "./tool";
 
 export { makeAssistantVisible } from "./makeAssistantVisible";
 
+export {
+  Toolkit,
+  type ToolDefinition,
+  type ToolkitFallback,
+  type ToolkitLayout,
+} from "./toolbox";
+
+export { Tools } from "../client/Tools";
+
 export * from "./registry";
 export * from "./frame";

package/src/model-context/makeAssistantVisible.tsx

@@ -76,7 +76,7 @@ export const makeAssistantVisible = <T extends ComponentType<any>>(
 
       const { clickable, editable } = config ?? {};
       useEffect(() => {
-        return api.registerModelContextProvider({
+        return api.modelContext().register({
           getModelContext: () => {
             return {
               tools: {

package/src/model-context/toolbox.tsx

@@ -0,0 +1,37 @@
+"use client";
+
+import type { ComponentType, ReactNode } from "react";
+import type { Tool } from "assistant-stream";
+import type { ToolCallMessagePartComponent } from "../types/MessagePartComponentTypes";
+
+export type ToolDefinition<
+  TArgs extends Record<string, unknown>,
+  TResult,
+> = Tool<TArgs, TResult> & {
+  render?: ToolCallMessagePartComponent<TArgs, TResult> | undefined;
+};
+
+export const FallbackSymbol = Symbol("Toolkit.Fallback");
+export const LayoutSymbol = Symbol("Toolkit.Layout");
+
+export type ToolkitFallback = {
+  render: ToolCallMessagePartComponent<unknown, unknown>;
+};
+
+export type ToolkitLayout = {
+  render: ComponentType<{ children: ReactNode }>;
+};
+
+export type Toolkit = Record<string, ToolDefinition<any, any>> & {
+  [FallbackSymbol]?: ToolkitFallback;
+  [LayoutSymbol]?: ToolkitLayout;
+};
+
+export const Toolkit = {
+  Fallback: FallbackSymbol,
+  Layout: LayoutSymbol,
+} as const;
+
+export type ToolsConfig = {
+  toolkit: Toolkit;
+};

package/src/model-context/useAssistantInstructions.tsx

@@ -27,7 +27,7 @@ export const useAssistantInstructions = (
     const config = {
       system: instruction,
     };
-    return api.registerModelContextProvider({
+    return api.modelContext().register({
       getModelContext: () => config,
     });
   }, [api, instruction, disabled]);

package/src/model-context/useAssistantTool.tsx

@@ -33,7 +33,7 @@ export const useAssistantTool = <
         [toolName]: rest,
       },
     };
-    return api.registerModelContextProvider({
+    return api.modelContext().register({
       getModelContext: () => context,
     });
   }, [api, tool]);

package/src/primitives/message/MessageParts.tsx

@@ -196,8 +196,9 @@ const ToolUIDisplay = ({
   Fallback: ToolCallMessagePartComponent | undefined;
 } & ToolCallMessagePartProps) => {
   const Render = useAssistantState(({ toolUIs }) => {
-    const Render = toolUIs[props.toolName] ?? Fallback;
-    if (Array.isArray(Render)) return Render[0];
+    const Render =
+      toolUIs.tools[props.toolName] ?? toolUIs.fallback ?? Fallback;
+    if (Array.isArray(Render)) return Render[0] ?? Fallback;
     return Render;
   });
   if (!Render) return null;

package/src/primitives/message/MessagePartsGrouped.tsx

@@ -221,8 +221,9 @@ const ToolUIDisplay = ({
   Fallback: ToolCallMessagePartComponent | undefined;
 } & ToolCallMessagePartProps) => {
   const Render = useAssistantState(({ toolUIs }) => {
-    const Render = toolUIs[props.toolName] ?? Fallback;
-    if (Array.isArray(Render)) return Render[0];
+    const Render =
+      toolUIs.tools[props.toolName] ?? toolUIs.fallback ?? Fallback;
+    if (Array.isArray(Render)) return Render[0] ?? Fallback;
     return Render;
   });
   if (!Render) return null;

package/src/tests/setup.ts

@@ -1,11 +1,5 @@
 // This file contains setup code for tests
 import { vi } from "vitest";
 
-// Set up global mocks if needed
-// Using a fixed date to avoid recursive calls
-const OriginalDate = global.Date;
-const fixedDate = new OriginalDate("2023-01-01");
-global.Date = vi.fn(() => fixedDate) as any;
-global.Date.now = vi.fn(() => fixedDate.getTime());
-
-// Add any other global setup needed for tests
+// Use Vitest's built-in to mock Date
+vi.setSystemTime(new Date("2023-01-01"));

package/src/utils/tap-store/derived-scopes.ts

@@ -0,0 +1,185 @@
+import { resource, tapEffect } from "@assistant-ui/tap";
+import { tapMemo, tapRef, tapResource, tapResources } from "@assistant-ui/tap";
+import { createAssistantApiField } from "../../context/react/AssistantApiContext";
+import type {
+  AssistantApi,
+  AssistantApiField,
+} from "../../context/react/AssistantApiContext";
+import type {
+  AssistantEvent,
+  AssistantEventCallback,
+  AssistantEventSelector,
+} from "../../types/EventTypes";
+import type { ResourceElement, Unsubscribe } from "@assistant-ui/tap";
+
+/**
+ * Extract the API return type from an AssistantApiField
+ */
+type ExtractApiType<T> =
+  T extends AssistantApiField<infer TApi, any> ? TApi : never;
+
+/**
+ * Extract the metadata type from an AssistantApiField
+ *
+ * Used in DerivedScopesInput to validate that each field's source/query types match
+ * the expected types from AssistantApi.
+ */
+type ExtractMeta<T> =
+  T extends AssistantApiField<any, infer TMeta> ? TMeta : never;
+
+/**
+ * Get only the field names from AssistantApi (exclude method names)
+ */
+type AssistantApiFieldNames = {
+  [K in keyof AssistantApi]: AssistantApi[K] extends { source: any; query: any }
+    ? K
+    : never;
+}[keyof AssistantApi];
+
+/**
+ * Configuration for a derived scope field - infers types from the actual values provided
+ */
+export type DerivedScopeConfig<TSource extends string | null, TQuery, TApi> = {
+  source: TSource;
+  query: TQuery;
+  get: () => TApi;
+};
+
+/**
+ * Type for the special `on` callback function
+ */
+export type OnCallbackFn = <TEvent extends AssistantEvent>(
+  selector: AssistantEventSelector<TEvent>,
+  callback: AssistantEventCallback<TEvent>,
+) => Unsubscribe;
+
+/**
+ * Type for the special `subscribe` callback function
+ */
+export type SubscribeCallbackFn = (listener: () => void) => Unsubscribe;
+
+/**
+ * Type for the special `flushSync` callback function
+ */
+export type FlushSyncCallbackFn = () => void;
+
+/**
+ * Type for special non-field functions in AssistantApi
+ */
+export type SpecialCallbacks = {
+  on?: OnCallbackFn;
+  subscribe?: SubscribeCallbackFn;
+  flushSync?: FlushSyncCallbackFn;
+};
+
+/**
+ * Type for the scopes parameter - allows both DerivedScope elements and special callbacks.
+ * Field names are restricted to valid AssistantApi field names.
+ * TypeScript validates that the source/query/get types match the expected field type.
+ */
+export type DerivedScopesInput = {
+  [K in AssistantApiFieldNames]?: ResourceElement<
+    AssistantApiField<
+      ExtractApiType<AssistantApi[K]>,
+      {
+        source: ExtractMeta<AssistantApi[K]>["source"];
+        query: ExtractMeta<AssistantApi[K]>["query"];
+      }
+    >
+  >;
+} & SpecialCallbacks;
+
+/**
+ * DerivedScope resource - memoizes an AssistantApiField based on source and query.
+ * The get callback always calls the most recent version (useEffectEvent pattern).
+ * TypeScript infers TSource, TQuery, and TApi from the config object.
+ * Validation happens at the DerivedScopesInput level.
+ */
+export const DerivedScope = resource(
+  <TSource extends string | null, TQuery, TApi>(
+    config: DerivedScopeConfig<TSource, TQuery, TApi>,
+  ): AssistantApiField<
+    TApi,
+    {
+      source: TSource;
+      query: TQuery;
+    }
+  > => {
+    const getRef = tapRef(config.get);
+    tapEffect(() => {
+      getRef.current = config.get;
+    });
+
+    return tapMemo(() => {
+      return createAssistantApiField({
+        source: config.source,
+        query: config.query,
+        get: () => getRef.current(),
+      });
+    }, [config.source, JSON.stringify(config.query)]);
+  },
+);
+
+/**
+ * Helper resource to wrap each scope field - stable resource identity for proper memoization.
+ * Creating this outside the map ensures tapResources can properly track and memoize each field.
+ */
+const ScopeFieldWithNameResource = resource(
+  (config: {
+    fieldName: string;
+    scopeElement: ReturnType<typeof DerivedScope>;
+  }) => {
+    const field = tapResource(config.scopeElement);
+    return tapMemo(
+      () => [config.fieldName, field] as const,
+      [config.fieldName, field],
+    );
+  },
+);
+
+/**
+ * DerivedScopes resource - takes an object of DerivedScope resource elements and special callbacks,
+ * and returns a Partial<AssistantApi> with all the derived fields.
+ */
+export const DerivedScopes = resource(
+  (scopes: DerivedScopesInput): Partial<AssistantApi> => {
+    const { on, subscribe, flushSync, ...scopeFields } = scopes;
+    const callbacksRef = tapRef({ on, subscribe, flushSync });
+    tapEffect(() => {
+      callbacksRef.current = { on, subscribe, flushSync };
+    });
+
+    const results = tapResources(
+      Object.entries(scopeFields).map(([fieldName, scopeElement]) =>
+        ScopeFieldWithNameResource(
+          {
+            fieldName,
+            scopeElement: scopeElement as ReturnType<typeof DerivedScope>,
+          },
+          { key: fieldName },
+        ),
+      ),
+    );
+
+    return tapMemo(() => {
+      const result = Object.fromEntries(results) as Partial<AssistantApi>;
+
+      const {
+        on: onCb,
+        subscribe: subCb,
+        flushSync: flushCb,
+      } = callbacksRef.current;
+
+      if (onCb) {
+        result.on = <TEvent extends AssistantEvent>(
+          selector: AssistantEventSelector<TEvent>,
+          callback: AssistantEventCallback<TEvent>,
+        ) => onCb(selector, callback);
+      }
+      if (subCb) result.subscribe = (listener) => subCb(listener);
+      if (flushCb) result.flushSync = () => flushCb();
+
+      return result;
+    }, [...results]);
+  },
+);

package/src/utils/tap-store/index.ts

@@ -1,2 +1,11 @@
 export { type Store, asStore } from "./store";
 export { tapApi, type ApiObject } from "./tap-api";
+export {
+  DerivedScope,
+  type DerivedScopeConfig,
+  type OnCallbackFn,
+  type SubscribeCallbackFn,
+  type FlushSyncCallbackFn,
+  type SpecialCallbacks,
+  type DerivedScopesInput,
+} from "./derived-scopes";

package/src/utils/tap-store/tap-api.ts

@@ -40,7 +40,7 @@ export const tapApi = <TApi extends ApiObject & { getState: () => any }>(
     key?: string | undefined;
   },
 ) => {
-  const ref = tapRef(() => api);
+  const ref = tapRef(api);
   tapEffect(() => {
     ref.current = api;
   });