@assistant-ui/react 0.11.58 → 0.12.1

package/src/client/ThreadMessageClient.tsx

@@ -4,12 +4,10 @@ import {
   tapMemo,
   tapState,
   tapInlineResource,
+  withKey,
 } from "@assistant-ui/tap";
-import { AttachmentClientApi } from "./types/Attachment";
-import { MessageClientState, MessageClientApi } from "./types/Message";
-import { MessagePartClientState, MessagePartClientApi } from "./types/Part";
-import { tapLookupResources } from "./util-hooks/tapLookupResources";
-import { tapApi } from "../utils/tap-store";
+import { type ClientOutput, tapClientLookup } from "@assistant-ui/store";
+import { MessageState, PartState } from "../types/scopes";
 import {
   ThreadAssistantMessagePart,
   ThreadUserMessagePart,
@@ -19,16 +17,21 @@ import {
 import { NoOpComposerClient } from "./NoOpComposerClient";
 
 const ThreadMessagePartClient = resource(
-  ({ part }: { part: ThreadAssistantMessagePart | ThreadUserMessagePart }) => {
-    const state = tapMemo<MessagePartClientState>(() => {
+  ({
+    part,
+  }: {
+    part: ThreadAssistantMessagePart | ThreadUserMessagePart;
+  }): ClientOutput<"part"> => {
+    const state = tapMemo<PartState>(() => {
       return {
         ...part,
         status: { type: "complete" },
       };
     }, [part]);
 
-    return tapApi<MessagePartClientApi>(
-      {
+    return {
+      state,
+      methods: {
         getState: () => state,
         addToolResult: () => {
           throw new Error("Not supported");
@@ -37,28 +40,21 @@ const ThreadMessagePartClient = resource(
           throw new Error("Not supported");
         },
       },
-      {
-        key:
-          state.type === "tool-call"
-            ? `toolCallId-${state.toolCallId}`
-            : undefined,
-      },
-    );
+    };
   },
 );
+
 const ThreadMessageAttachmentClient = resource(
-  ({ attachment }: { attachment: Attachment }) => {
-    return tapApi<AttachmentClientApi>(
-      {
+  ({ attachment }: { attachment: Attachment }): ClientOutput<"attachment"> => {
+    return {
+      state: attachment,
+      methods: {
         getState: () => attachment,
         remove: () => {
           throw new Error("Not supported");
         },
       },
-      {
-        key: attachment.id,
-      },
-    );
+    };
   },
 );
 export type ThreadMessageClientProps = {
@@ -68,6 +64,7 @@ export type ThreadMessageClientProps = {
   branchNumber?: number;
   branchCount?: number;
 };
+
 export const ThreadMessageClient = resource(
   ({
     message,
@@ -75,35 +72,38 @@ export const ThreadMessageClient = resource(
     isLast = true,
     branchNumber = 1,
     branchCount = 1,
-  }: ThreadMessageClientProps) => {
+  }: ThreadMessageClientProps): ClientOutput<"message"> => {
     const [isCopiedState, setIsCopied] = tapState(false);
     const [isHoveringState, setIsHovering] = tapState(false);
 
-    const parts = tapLookupResources(
-      message.content.map((part, idx) => [
-        "toolCallId" in part && part.toolCallId != null
-          ? `toolCallId-${part.toolCallId}`
-          : `index-${idx}`,
-        ThreadMessagePartClient({ part }),
-      ]),
+    const parts = tapClientLookup(
+      () =>
+        message.content.map((part, idx) =>
+          withKey(
+            "toolCallId" in part && part.toolCallId != null
+              ? `toolCallId-${part.toolCallId}`
+              : `index-${idx}`,
+            ThreadMessagePartClient({ part }),
+          ),
+        ),
+      [message.content],
     );
 
-    const attachments = tapLookupResources(
-      message.attachments?.map((attachment) => [
-        attachment.id,
-        ThreadMessageAttachmentClient({ attachment }),
-      ]) ?? [],
+    const attachments = tapClientLookup(
+      () =>
+        (message.attachments ?? []).map((attachment) =>
+          withKey(attachment.id, ThreadMessageAttachmentClient({ attachment })),
+        ),
+      [message.attachments],
     );
 
-    const composerState = tapInlineResource(
-      NoOpComposerClient({ type: "edit" }),
-    );
+    const composer = tapInlineResource(NoOpComposerClient({ type: "edit" }));
 
-    const state = tapMemo<MessageClientState>(() => {
+    const state = tapMemo<MessageState>(() => {
       return {
         ...message,
         parts: parts.state,
-        composer: composerState.state,
+        composer: composer.state,
         parentId: null,
         index,
         isLast,
@@ -114,52 +114,65 @@ export const ThreadMessageClient = resource(
         isCopied: isCopiedState,
         isHovering: isHoveringState,
       };
-    }, [message, index, isCopiedState, isHoveringState, isLast]);
+    }, [
+      message,
+      index,
+      isCopiedState,
+      isHoveringState,
+      isLast,
+      parts.state,
+      composer.state,
+      branchNumber,
+      branchCount,
+    ]);
 
-    return tapApi<MessageClientApi>({
-      getState: () => state,
-      composer: composerState.api,
-      part: (selector) => {
-        if ("index" in selector) {
-          return parts.api({ index: selector.index });
-        } else {
-          return parts.api({ key: `toolCallId-${selector.toolCallId}` });
-        }
-      },
-      attachment: (selector) => {
-        if ("id" in selector) {
-          return attachments.api({ key: selector.id });
-        } else {
-          return attachments.api(selector);
-        }
-      },
-      reload: () => {
-        throw new Error("Not supported in ThreadMessageProvider");
-      },
-      speak: () => {
-        throw new Error("Not supported in ThreadMessageProvider");
-      },
-      stopSpeaking: () => {
-        throw new Error("Not supported in ThreadMessageProvider");
-      },
-      submitFeedback: () => {
-        throw new Error("Not supported in ThreadMessageProvider");
-      },
-      switchToBranch: () => {
-        throw new Error("Not supported in ThreadMessageProvider");
-      },
-      getCopyText: () => {
-        return message.content
-          .map((part) => {
-            if ("text" in part && typeof part.text === "string") {
-              return part.text;
-            }
-            return "";
-          })
-          .join("\n");
+    return {
+      state,
+      methods: {
+        getState: () => state,
+        composer: composer.methods,
+        part: (selector) => {
+          if ("index" in selector) {
+            return parts.get({ index: selector.index });
+          } else {
+            return parts.get({ key: `toolCallId-${selector.toolCallId}` });
+          }
+        },
+        attachment: (selector) => {
+          if ("id" in selector) {
+            return attachments.get({ key: selector.id });
+          } else {
+            return attachments.get(selector);
+          }
+        },
+        reload: () => {
+          throw new Error("Not supported in ThreadMessageProvider");
+        },
+        speak: () => {
+          throw new Error("Not supported in ThreadMessageProvider");
+        },
+        stopSpeaking: () => {
+          throw new Error("Not supported in ThreadMessageProvider");
+        },
+        submitFeedback: () => {
+          throw new Error("Not supported in ThreadMessageProvider");
+        },
+        switchToBranch: () => {
+          throw new Error("Not supported in ThreadMessageProvider");
+        },
+        getCopyText: () => {
+          return message.content
+            .map((part) => {
+              if ("text" in part && typeof part.text === "string") {
+                return part.text;
+              }
+              return "";
+            })
+            .join("\n");
+        },
+        setIsCopied,
+        setIsHovering,
       },
-      setIsCopied,
-      setIsHovering,
-    });
+    };
   },
 );

package/src/client/Tools.ts

@@ -1,77 +1,88 @@
-import { resource, tapState, tapEffect } from "@assistant-ui/tap";
-import { tapApi } from "../utils/tap-store";
-import { tapModelContext } from "./ModelContext";
-import { ToolsState, ToolsApi } from "./types/Tools";
+import { resource, tapState, tapEffect, tapCallback } from "@assistant-ui/tap";
+import { tapAssistantClientRef, type ClientOutput } from "@assistant-ui/store";
+import { ToolsState } from "../types/scopes";
 import type { Tool } from "assistant-stream";
 import { type Toolkit } from "../model-context/toolbox";
+import { ToolCallMessagePartComponent } from "../types";
 
-export const Tools = resource(({ toolkit }: { toolkit?: Toolkit }) => {
-  const [state, setState] = tapState<ToolsState>(() => ({
-    tools: {},
-  }));
+export const Tools = resource(
+  ({ toolkit }: { toolkit?: Toolkit }): ClientOutput<"tools"> => {
+    const [state, setState] = tapState<ToolsState>(() => ({
+      tools: {},
+    }));
 
-  const modelContext = tapModelContext();
+    const clientRef = tapAssistantClientRef();
 
-  tapEffect(() => {
-    if (!toolkit) return;
-    const unsubscribes: (() => void)[] = [];
+    const setToolUI = tapCallback(
+      (toolName: string, render: ToolCallMessagePartComponent) => {
+        setState((prev) => {
+          return {
+            ...prev,
+            tools: {
+              ...prev.tools,
+              [toolName]: [...(prev.tools[toolName] ?? []), render],
+            },
+          };
+        });
 
-    // Register tool UIs (exclude symbols)
-    for (const [toolName, tool] of Object.entries(toolkit)) {
-      if (tool.render) {
-        unsubscribes.push(setToolUI(toolName, tool.render));
-      }
-    }
-
-    // Register tools with model context (exclude symbols)
-    const toolsWithoutRender = Object.entries(toolkit).reduce(
-      (acc, [name, tool]) => {
-        const { render, ...rest } = tool;
-        acc[name] = rest;
-        return acc;
+        return () => {
+          setState((prev) => {
+            return {
+              ...prev,
+              tools: {
+                ...prev.tools,
+                [toolName]:
+                  prev.tools[toolName]?.filter((r) => r !== render) ?? [],
+              },
+            };
+          });
+        };
       },
-      {} as Record<string, Tool<any, any>>,
+      [],
     );
 
-    const modelContextProvider = {
-      getModelContext: () => ({
-        tools: toolsWithoutRender,
-      }),
-    };
+    tapEffect(() => {
+      if (!toolkit) return;
+      const unsubscribes: (() => void)[] = [];
 
-    unsubscribes.push(modelContext.register(modelContextProvider));
-
-    return () => {
-      unsubscribes.forEach((fn) => fn());
-    };
-  }, [toolkit, modelContext]);
+      // Register tool UIs (exclude symbols)
+      for (const [toolName, tool] of Object.entries(toolkit)) {
+        if (tool.render) {
+          unsubscribes.push(setToolUI(toolName, tool.render));
+        }
+      }
 
-  const setToolUI = (toolName: string, render: any) => {
-    setState((prev) => {
-      return {
-        ...prev,
-        tools: {
-          ...prev.tools,
-          [toolName]: [...(prev.tools[toolName] ?? []), render],
+      // Register tools with model context (exclude symbols)
+      const toolsWithoutRender = Object.entries(toolkit).reduce(
+        (acc, [name, tool]) => {
+          const { render, ...rest } = tool;
+          acc[name] = rest;
+          return acc;
         },
+        {} as Record<string, Tool<any, any>>,
+      );
+
+      const modelContextProvider = {
+        getModelContext: () => ({
+          tools: toolsWithoutRender,
+        }),
       };
-    });
 
-    return () => {
-      setState((prev) => {
-        return {
-          ...prev,
-          tools: {
-            ...prev.tools,
-            [toolName]: prev.tools[toolName]?.filter((r) => r !== render) ?? [],
-          },
-        };
-      });
-    };
-  };
+      unsubscribes.push(
+        clientRef.current!.modelContext().register(modelContextProvider),
+      );
 
-  return tapApi<ToolsApi>({
-    getState: () => state,
-    setToolUI,
-  });
-});
+      return () => {
+        unsubscribes.forEach((fn) => fn());
+      };
+    }, [toolkit, setToolUI, clientRef]);
+
+    return {
+      state,
+      methods: {
+        getState: () => state,
+        setToolUI,
+      },
+    };
+  },
+);

package/src/context/providers/AttachmentByIndexProvider.tsx

@@ -1,29 +1,22 @@
 "use client";
 
 import { type FC, type PropsWithChildren } from "react";
-
-import {
-  AssistantProvider,
-  useAssistantApi,
-  useExtendedAssistantApi,
-} from "../react/AssistantApiContext";
-import { DerivedScope } from "../../utils/tap-store/derived-scopes";
+import { useAui, AuiProvider, Derived } from "@assistant-ui/store";
 
 export const MessageAttachmentByIndexProvider: FC<
   PropsWithChildren<{
     index: number;
   }>
 > = ({ index, children }) => {
-  const baseApi = useAssistantApi();
-  const api = useExtendedAssistantApi({
-    attachment: DerivedScope({
+  const aui = useAui({
+    attachment: Derived({
       source: "message",
       query: { type: "index", index },
-      get: () => baseApi.message().attachment({ index }),
+      get: (aui) => aui.message().attachment({ index }),
     }),
   });
 
-  return <AssistantProvider api={api}>{children}</AssistantProvider>;
+  return <AuiProvider value={aui}>{children}</AuiProvider>;
 };
 
 export const ComposerAttachmentByIndexProvider: FC<
@@ -31,14 +24,13 @@ export const ComposerAttachmentByIndexProvider: FC<
     index: number;
   }>
 > = ({ index, children }) => {
-  const baseApi = useAssistantApi();
-  const api = useExtendedAssistantApi({
-    attachment: DerivedScope({
+  const aui = useAui({
+    attachment: Derived({
       source: "composer",
       query: { type: "index", index },
-      get: () => baseApi.composer().attachment({ index }),
+      get: (aui) => aui.composer().attachment({ index }),
     }),
   });
 
-  return <AssistantProvider api={api}>{children}</AssistantProvider>;
+  return <AuiProvider value={aui}>{children}</AuiProvider>;
 };

package/src/context/providers/MessageByIndexProvider.tsx

@@ -1,50 +1,25 @@
 "use client";
 
 import { type FC, type PropsWithChildren } from "react";
-import {
-  AssistantProvider,
-  useAssistantApi,
-  useExtendedAssistantApi,
-} from "../react/AssistantApiContext";
-import {
-  checkEventScope,
-  normalizeEventSelector,
-} from "../../types/EventTypes";
-import { DerivedScope } from "../../utils/tap-store/derived-scopes";
+import { useAui, AuiProvider, Derived } from "@assistant-ui/store";
 
 export const MessageByIndexProvider: FC<
   PropsWithChildren<{
     index: number;
   }>
 > = ({ index, children }) => {
-  const baseApi = useAssistantApi();
-  const api = useExtendedAssistantApi({
-    message: DerivedScope({
+  const aui = useAui({
+    message: Derived({
       source: "thread",
       query: { type: "index", index },
-      get: () => baseApi.thread().message({ index }),
+      get: (aui) => aui.thread().message({ index }),
     }),
-    composer: DerivedScope({
+    composer: Derived({
       source: "message",
       query: {},
-      get: () => baseApi.thread().message({ index }).composer,
+      get: (aui) => aui.thread().message({ index }).composer,
     }),
-    on(selector, callback) {
-      const getMessage = () => baseApi.thread().message({ index });
-      const { event, scope } = normalizeEventSelector(selector);
-      if (
-        !checkEventScope("composer", scope, event) &&
-        !checkEventScope("message", scope, event)
-      )
-        return baseApi.on(selector, callback);
-
-      return baseApi.on({ scope: "thread", event }, (e) => {
-        if (e.messageId === getMessage().getState().id) {
-          callback(e);
-        }
-      });
-    },
   });
 
-  return <AssistantProvider api={api}>{children}</AssistantProvider>;
+  return <AuiProvider value={aui}>{children}</AuiProvider>;
 };

package/src/context/providers/MessageProvider.tsx

@@ -1,30 +1,18 @@
 "use client";
 
 import { type FC, type PropsWithChildren } from "react";
-import {
-  AssistantProvider,
-  useExtendedAssistantApi,
-} from "../react/AssistantApiContext";
-import { useResource } from "@assistant-ui/tap/react";
-import { asStore } from "../../utils/tap-store";
+import { useAui, AuiProvider } from "@assistant-ui/store";
 import {
   ThreadMessageClientProps,
   ThreadMessageClient,
 } from "../../client/ThreadMessageClient";
-import { DerivedScope } from "../../utils/tap-store/derived-scopes";
 
 export const MessageProvider: FC<
   PropsWithChildren<ThreadMessageClientProps>
 > = ({ children, ...props }) => {
-  const store = useResource(asStore(ThreadMessageClient(props)));
-  const api = useExtendedAssistantApi({
-    message: DerivedScope({
-      source: "root",
-      query: {},
-      get: () => store.getState().api,
-    }),
-    subscribe: store.subscribe,
+  const aui = useAui({
+    message: ThreadMessageClient(props),
   });
 
-  return <AssistantProvider api={api}>{children}</AssistantProvider>;
+  return <AuiProvider value={aui}>{children}</AuiProvider>;
 };

package/src/context/providers/PartByIndexProvider.tsx

@@ -1,26 +1,20 @@
 "use client";
 
 import { type FC, type PropsWithChildren } from "react";
-import {
-  AssistantProvider,
-  useAssistantApi,
-  useExtendedAssistantApi,
-} from "../react/AssistantApiContext";
-import { DerivedScope } from "../../utils/tap-store/derived-scopes";
+import { useAui, AuiProvider, Derived } from "@assistant-ui/store";
 
 export const PartByIndexProvider: FC<
   PropsWithChildren<{
     index: number;
   }>
 > = ({ index, children }) => {
-  const baseApi = useAssistantApi();
-  const api = useExtendedAssistantApi({
-    part: DerivedScope({
+  const aui = useAui({
+    part: Derived({
       source: "message",
       query: { type: "index", index },
-      get: () => baseApi.message().part({ index }),
+      get: (aui) => aui.message().part({ index }),
     }),
   });
 
-  return <AssistantProvider api={api}>{children}</AssistantProvider>;
+  return <AuiProvider value={aui}>{children}</AuiProvider>;
 };

package/src/context/providers/TextMessagePartProvider.tsx

@@ -1,22 +1,19 @@
 "use client";
 
 import { type FC, type PropsWithChildren } from "react";
-import {
-  AssistantProvider,
-  useExtendedAssistantApi,
-} from "../react/AssistantApiContext";
-import {
-  MessagePartClientApi,
-  MessagePartClientState,
-} from "../../client/types/Part";
+import { useAui, AuiProvider, type ClientOutput } from "@assistant-ui/store";
+import { PartState } from "../../types/scopes";
 import { resource, tapMemo } from "@assistant-ui/tap";
-import { useResource } from "@assistant-ui/tap/react";
-import { asStore, tapApi } from "../../utils/tap-store";
-import { DerivedScope } from "../../utils/tap-store/derived-scopes";
 
 const TextMessagePartClient = resource(
-  ({ text, isRunning }: { text: string; isRunning: boolean }) => {
-    const state = tapMemo<MessagePartClientState>(
+  ({
+    text,
+    isRunning,
+  }: {
+    text: string;
+    isRunning: boolean;
+  }): ClientOutput<"part"> => {
+    const state = tapMemo<PartState>(
       () => ({
         type: "text",
         text,
@@ -25,15 +22,18 @@ const TextMessagePartClient = resource(
       [text, isRunning],
     );
 
-    return tapApi<MessagePartClientApi>({
-      getState: () => state,
-      addToolResult: () => {
-        throw new Error("Not supported");
-      },
-      resumeToolCall: () => {
-        throw new Error("Not supported");
+    return {
+      state,
+      methods: {
+        getState: () => state,
+        addToolResult: () => {
+          throw new Error("Not supported");
+        },
+        resumeToolCall: () => {
+          throw new Error("Not supported");
+        },
       },
-    });
+    };
   },
 );
 
@@ -43,17 +43,9 @@ export const TextMessagePartProvider: FC<
     isRunning?: boolean;
   }>
 > = ({ text, isRunning = false, children }) => {
-  const store = useResource(
-    asStore(TextMessagePartClient({ text, isRunning })),
-  );
-  const api = useExtendedAssistantApi({
-    part: DerivedScope({
-      source: "root",
-      query: {},
-      get: () => store.getState().api,
-    }),
-    subscribe: store.subscribe,
+  const aui = useAui({
+    part: TextMessagePartClient({ text, isRunning }),
   });
 
-  return <AssistantProvider api={api}>{children}</AssistantProvider>;
+  return <AuiProvider value={aui}>{children}</AuiProvider>;
 };