@assistant-ui/react 0.11.35 → 0.11.36

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/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;
   });