@assistant-ui/react 0.11.5 → 0.11.6

package/src/context/providers/PartProvider.tsx

@@ -3,7 +3,7 @@
 import { useMemo, type FC, type PropsWithChildren } from "react";
 import {
   AssistantApi,
-  AssistantApiProvider,
+  AssistantProvider,
   useAssistantApi,
   createAssistantApiField,
 } from "../react/AssistantApiContext";
@@ -24,5 +24,5 @@ export const PartByIndexProvider: FC<
     } satisfies Partial<AssistantApi>;
   }, [api, index]);
 
-  return <AssistantApiProvider api={api2}>{children}</AssistantApiProvider>;
+  return <AssistantProvider api={api2}>{children}</AssistantProvider>;
 };

package/src/context/providers/TextMessagePartProvider.tsx

@@ -2,7 +2,7 @@
 
 import { useMemo, type FC, type PropsWithChildren } from "react";
 import {
-  AssistantApiProvider,
+  AssistantProvider,
   AssistantApi,
   createAssistantApiField,
 } from "../react/AssistantApiContext";
@@ -25,15 +25,13 @@ const TextMessagePartClient = resource(
       [text, isRunning],
     );
 
-    const api = tapApi<MessagePartClientApi>({
+    return tapApi<MessagePartClientApi>({
       getState: () => state,
       addToolResult: () => {
         throw new Error("Not supported");
       },
       __internal_getRuntime: () => null,
     });
-
-    return api;
   },
 );
 
@@ -51,12 +49,12 @@ export const TextMessagePartProvider: FC<
       part: createAssistantApiField({
         source: "root",
         query: {},
-        get: () => store.getState(),
+        get: () => store.getState().api,
       }),
       subscribe: store.subscribe,
-      // flushSync: store.flushSync,
+      flushSync: store.flushSync,
     } satisfies Partial<AssistantApi>;
   }, [store]);
 
-  return <AssistantApiProvider api={api}>{children}</AssistantApiProvider>;
+  return <AssistantProvider api={api}>{children}</AssistantProvider>;
 };

package/src/context/providers/ThreadListItemProvider.tsx

@@ -3,7 +3,7 @@
 import { useMemo, type FC, type PropsWithChildren } from "react";
 import {
   AssistantApi,
-  AssistantApiProvider,
+  AssistantProvider,
   useAssistantApi,
   createAssistantApiField,
 } from "../react/AssistantApiContext";
@@ -42,7 +42,7 @@ export const ThreadListItemByIndexProvider: FC<
     } satisfies Partial<AssistantApi>;
   }, [api, index, archived]);
 
-  return <AssistantApiProvider api={api2}>{children}</AssistantApiProvider>;
+  return <AssistantProvider api={api2}>{children}</AssistantProvider>;
 };
 
 export const ThreadListItemByIdProvider: FC<
@@ -73,5 +73,5 @@ export const ThreadListItemByIdProvider: FC<
     } satisfies Partial<AssistantApi>;
   }, [api, id]);
 
-  return <AssistantApiProvider api={api2}>{children}</AssistantApiProvider>;
+  return <AssistantProvider api={api2}>{children}</AssistantProvider>;
 };

package/src/context/react/AssistantApiContext.tsx

@@ -43,6 +43,7 @@ import {
   ThreadListClientApi,
   ThreadListClientState,
 } from "../../client/types/ThreadList";
+import { ThreadViewportProvider } from "../providers/ThreadViewportProvider";
 
 export type AssistantState = {
   readonly threads: ThreadListClientState;
@@ -280,13 +281,17 @@ const extendApi = (
   };
 };
 
-export const AssistantApiProvider: FC<
+export const AssistantProvider: FC<
   PropsWithChildren<{ api: Partial<AssistantApi> }>
 > = ({ api: api2, children }) => {
   const api = useAssistantApi();
   const extendedApi = useMemo(() => extendApi(api, api2), [api, api2]);
 
   return (
-    <AssistantApiContext value={extendedApi}>{children}</AssistantApiContext>
+    <AssistantApiContext value={extendedApi}>
+      {/* TODO temporarily allow accessing viewport state from outside the viewport */}
+      {/* TODO figure out if this behavior should be deprecated, since it is quite hacky */}
+      <ThreadViewportProvider>{children}</ThreadViewportProvider>
+    </AssistantApiContext>
   );
 };

package/src/context/react/index.ts

@@ -54,3 +54,5 @@ export {
   useThreadListItemRuntime,
   useThreadListItem,
 } from "../../legacy-runtime/hooks/ThreadListItemContext";
+
+export { AssistantProvider } from "./AssistantApiContext";

package/src/legacy-runtime/AssistantRuntimeProvider.tsx

@@ -1,12 +1,10 @@
 "use client";
 
 import { FC, memo, PropsWithChildren } from "react";
-import { AssistantApiProvider } from "../context/react/AssistantApiContext";
+import { AssistantProvider } from "../context/react/AssistantApiContext";
 import { AssistantRuntime } from "./runtime/AssistantRuntime";
 import { AssistantRuntimeCore } from "./runtime-cores/core/AssistantRuntimeCore";
 import { useAssistantClient } from "../client/AssistantClient";
-
-import { ThreadViewportProvider } from "../context/providers/ThreadViewportProvider";
 import { ThreadListClient } from "./client/ThreadListRuntimeClient";
 
 export namespace AssistantProvider {
@@ -46,15 +44,3 @@ export const AssistantRuntimeProviderImpl: FC<AssistantProvider.Props> = ({
 };
 
 export const AssistantRuntimeProvider = memo(AssistantRuntimeProviderImpl);
-
-const AssistantProvider: FC<
-  PropsWithChildren<{ api: ReturnType<typeof useAssistantClient> }>
-> = ({ children, api }) => {
-  return (
-    <AssistantApiProvider api={api}>
-      {/* TODO temporarily allow accessing viewport state from outside the viewport */}
-      {/* TODO figure out if this behavior should be deprecated, since it is quite hacky */}
-      <ThreadViewportProvider>{children}</ThreadViewportProvider>
-    </AssistantApiProvider>
-  );
-};

package/src/legacy-runtime/client/AttachmentRuntimeClient.ts

@@ -7,16 +7,15 @@ import { tapSubscribable } from "../util-hooks/tapSubscribable";
 export const AttachmentRuntimeClient = resource(
   ({ runtime }: { runtime: AttachmentRuntime }) => {
     const state = tapSubscribable(runtime);
-    const api = tapApi<AttachmentClientApi>({
-      getState: () => state,
-      remove: runtime.remove,
-      __internal_getRuntime: () => runtime,
-    });
-
-    return {
-      state: state,
-      api,
-      key: state.id,
-    };
+    return tapApi<AttachmentClientApi>(
+      {
+        getState: () => state,
+        remove: runtime.remove,
+        __internal_getRuntime: () => runtime,
+      },
+      {
+        key: state.id,
+      },
+    );
   },
 );

package/src/legacy-runtime/client/ComposerRuntimeClient.ts

@@ -98,7 +98,7 @@ export const ComposerClient = resource(
       };
     }, [runtimeState, attachments.state]);
 
-    const api = tapApi<ComposerClientApi>({
+    return tapApi<ComposerClientApi>({
       getState: () => state,
 
       setText: runtime.setText,
@@ -126,10 +126,5 @@ export const ComposerClient = resource(
 
       __internal_getRuntime: () => runtime,
     });
-
-    return {
-      state,
-      api,
-    };
   },
 );

package/src/legacy-runtime/client/MessagePartRuntimeClient.ts

@@ -7,21 +7,20 @@ export const MessagePartClient = resource(
   ({ runtime }: { runtime: MessagePartRuntime }) => {
     const runtimeState = tapSubscribable(runtime);
 
-    const api = tapApi<MessagePartClientApi>({
-      getState: () => runtimeState,
+    return tapApi<MessagePartClientApi>(
+      {
+        getState: () => runtimeState,
 
-      addToolResult: (result) => runtime.addToolResult(result),
+        addToolResult: (result) => runtime.addToolResult(result),
 
-      __internal_getRuntime: () => runtime,
-    });
-
-    return {
-      key:
-        runtimeState.type === "tool-call"
-          ? "toolCallId-" + runtimeState.toolCallId
-          : undefined,
-      state: runtimeState,
-      api,
-    };
+        __internal_getRuntime: () => runtime,
+      },
+      {
+        key:
+          runtimeState.type === "tool-call"
+            ? "toolCallId-" + runtimeState.toolCallId
+            : undefined,
+      },
+    );
   },
 );

package/src/legacy-runtime/client/MessageRuntimeClient.ts

@@ -99,44 +99,43 @@ export const MessageClient = resource(
       isHoveringState,
     ]);
 
-    const api = tapApi<MessageClientApi>({
-      getState: () => state,
-
-      composer: composer.api,
-
-      reload: (config) => runtime.reload(config),
-      speak: () => runtime.speak(),
-      stopSpeaking: () => runtime.stopSpeaking(),
-      submitFeedback: (feedback) => runtime.submitFeedback(feedback),
-      switchToBranch: (options) => runtime.switchToBranch(options),
-      getCopyText: () => runtime.unstable_getCopyText(),
-
-      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);
-        }
-      },
+    return tapApi<MessageClientApi>(
+      {
+        getState: () => state,
+
+        composer: composer.api,
+
+        reload: (config) => runtime.reload(config),
+        speak: () => runtime.speak(),
+        stopSpeaking: () => runtime.stopSpeaking(),
+        submitFeedback: (feedback) => runtime.submitFeedback(feedback),
+        switchToBranch: (options) => runtime.switchToBranch(options),
+        getCopyText: () => runtime.unstable_getCopyText(),
+
+        part: (selector) => {
+          if ("index" in selector) {
+            return parts.api({ index: selector.index });
+          } else {
+            return parts.api({ key: "toolCallId-" + selector.toolCallId });
+          }
+        },
 
-      setIsCopied,
-      setIsHovering,
+        attachment: (selector) => {
+          if ("id" in selector) {
+            return attachments.api({ key: selector.id });
+          } else {
+            return attachments.api(selector);
+          }
+        },
 
-      __internal_getRuntime: () => runtime,
-    });
+        setIsCopied,
+        setIsHovering,
 
-    return {
-      key: runtimeState.id,
-      state,
-      api,
-    };
+        __internal_getRuntime: () => runtime,
+      },
+      {
+        key: runtimeState.id,
+      },
+    );
   },
 );

package/src/legacy-runtime/client/ThreadListItemRuntimeClient.ts

@@ -38,23 +38,22 @@ export const ThreadListItemClient = resource(
       };
     }, [runtime, events]);
 
-    const api = tapApi<ThreadListItemClientApi>({
-      getState: () => runtimeState,
-      switchTo: runtime.switchTo,
-      rename: runtime.rename,
-      archive: runtime.archive,
-      unarchive: runtime.unarchive,
-      delete: runtime.delete,
-      generateTitle: runtime.generateTitle,
-      initialize: runtime.initialize,
-      detach: runtime.detach,
-      __internal_getRuntime: () => runtime,
-    });
-
-    return {
-      state: runtimeState,
-      api,
-      key: runtimeState.id,
-    };
+    return tapApi<ThreadListItemClientApi>(
+      {
+        getState: () => runtimeState,
+        switchTo: runtime.switchTo,
+        rename: runtime.rename,
+        archive: runtime.archive,
+        unarchive: runtime.unarchive,
+        delete: runtime.delete,
+        generateTitle: runtime.generateTitle,
+        initialize: runtime.initialize,
+        detach: runtime.detach,
+        __internal_getRuntime: () => runtime,
+      },
+      {
+        key: runtimeState.id,
+      },
+    );
   },
 );

package/src/legacy-runtime/client/ThreadListRuntimeClient.ts

@@ -53,7 +53,7 @@ export const ThreadListClient = resource(
       };
     }, [runtimeState, threadItems.state, main.state]);
 
-    const api = tapApi<ThreadListClientApi>({
+    return tapApi<ThreadListClientApi>({
       getState: () => state,
 
       thread: () => main.api,
@@ -81,10 +81,5 @@ export const ThreadListClient = resource(
         runtime.switchToNewThread();
       },
     });
-
-    return {
-      state,
-      api,
-    };
   },
 );

package/src/legacy-runtime/client/ThreadRuntimeClient.ts

@@ -111,7 +111,7 @@ export const ThreadClient = resource(
       };
     }, [runtimeState, messages, composer.state]);
 
-    const api = tapApi<ThreadClientApi>({
+    return tapApi<ThreadClientApi>({
       getState: () => state,
 
       composer: composer.api,
@@ -136,10 +136,5 @@ export const ThreadClient = resource(
 
       __internal_getRuntime: () => runtime,
     });
-
-    return {
-      state,
-      api,
-    };
   },
 );

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

@@ -34,15 +34,32 @@ class ReadonlyApiHandler<TApi extends ApiObject> implements ProxyHandler<TApi> {
   }
 }
 
-export const tapApi = <TApi extends ApiObject>(api: TApi) => {
+export const tapApi = <TApi extends ApiObject & { getState: () => any }>(
+  api: TApi,
+  options?: {
+    key?: string | undefined;
+  },
+) => {
   const ref = tapRef(() => api);
   tapEffect(() => {
     ref.current = api;
   });
 
-  return tapMemo(
+  const apiProxy = tapMemo(
     () =>
       new Proxy<TApi>({} as TApi, new ReadonlyApiHandler(() => ref.current)),
     [],
   );
+
+  const key = options?.key;
+  const state = api.getState();
+
+  return tapMemo(
+    () => ({
+      key,
+      state,
+      api: apiProxy,
+    }),
+    [state, key],
+  );
 };