@assistant-ui/react 0.5.91 → 0.5.93

package/src/api/subscribable/ShallowMemoizeSubject.ts

@@ -0,0 +1,47 @@
+import { shallowEqual } from "./shallowEqual";
+import { BaseSubject } from "./BaseSubject";
+import { SubscribableWithState } from "./Subscribable";
+import { SKIP_UPDATE } from "./SKIP_UPDATE";
+
+export class ShallowMemoizeSubject<TState extends object, TPath>
+  extends BaseSubject
+  implements SubscribableWithState<TState, TPath>
+{
+  public get path() {
+    return this.binding.path;
+  }
+
+  constructor(
+    private binding: SubscribableWithState<TState | SKIP_UPDATE, TPath>,
+  ) {
+    super();
+    const state = binding.getState();
+    if (state === SKIP_UPDATE)
+      throw new Error("Entry not available in the store");
+    this._previousState = state;
+  }
+
+  private _previousState: TState;
+  public getState = () => {
+    if (!this.isConnected) this._syncState();
+    return this._previousState;
+  };
+
+  private _syncState() {
+    const state = this.binding.getState();
+    if (state === SKIP_UPDATE) return false;
+    if (shallowEqual(state, this._previousState)) return false;
+    this._previousState = state;
+    return true;
+  }
+
+  protected _connect() {
+    const callback = () => {
+      if (this._syncState()) {
+        this.notifySubscribers();
+      }
+    };
+
+    return this.binding.subscribe(callback);
+  }
+}

package/src/api/subscribable/Subscribable.ts

@@ -0,0 +1,15 @@
+import { Unsubscribe } from "../../types";
+
+export type Subscribable = {
+  subscribe: (callback: () => void) => Unsubscribe;
+};
+
+export type SubscribableWithState<TState, TPath> = Subscribable & {
+  path: TPath;
+  getState: () => TState;
+};
+
+export type NestedSubscribable<
+  TState extends Subscribable | undefined,
+  TPath,
+> = SubscribableWithState<TState, TPath>;

package/src/api/subscribable/shallowEqual.ts

@@ -0,0 +1,16 @@
+export function shallowEqual<T extends object>(
+  objA: T | undefined,
+  objB: T | undefined,
+) {
+  if (objA === undefined && objB === undefined) return true;
+  if (objA === undefined) return false;
+  if (objB === undefined) return false;
+
+  for (const key of Object.keys(objA)) {
+    const valueA = objA[key as keyof T];
+    const valueB = objB[key as keyof T];
+    if (!Object.is(valueA, valueB)) return false;
+  }
+
+  return true;
+}

package/src/context/ReadonlyStore.ts

@@ -0,0 +1,7 @@
+import type { StoreApi } from "zustand";
+
+export type ReadonlyStore<T> = Omit<StoreApi<T>, "setState" | "destroy">;
+
+export const writableStore = <T>(store: ReadonlyStore<T> | undefined) => {
+  return store as unknown as StoreApi<T>;
+};

package/src/context/index.ts

@@ -0,0 +1,3 @@
+export * from "./providers";
+export * from "./stores";
+export * from "./react";

package/src/context/providers/AssistantRuntimeProvider.tsx

@@ -0,0 +1,74 @@
+"use client";
+
+import type { FC, PropsWithChildren } from "react";
+import { memo, useEffect, useMemo, useState } from "react";
+import { AssistantContext } from "../react/AssistantContext";
+import { makeAssistantToolUIsStore } from "../stores/AssistantToolUIs";
+import { ThreadRuntimeProvider } from "./ThreadRuntimeProvider";
+import { AssistantRuntime } from "../../api/AssistantRuntime";
+import { create } from "zustand";
+import { writableStore } from "../ReadonlyStore";
+
+export namespace AssistantRuntimeProvider {
+  export type Props = PropsWithChildren<{
+    /**
+     * The runtime to provide to the rest of your app.
+     */
+    runtime: AssistantRuntime;
+  }>;
+}
+
+const useAssistantRuntimeStore = (runtime: AssistantRuntime) => {
+  const [store] = useState(() => create(() => runtime));
+
+  useEffect(() => {
+    writableStore(store).setState(runtime, true);
+  }, [runtime, store]);
+
+  return store;
+};
+
+const useAssistantToolUIsStore = () => {
+  return useMemo(() => makeAssistantToolUIsStore(), []);
+};
+
+const useThreadManagerStore = (runtime: AssistantRuntime) => {
+  const [store] = useState(() =>
+    create(() => runtime.threadManager.getState()),
+  );
+
+  useEffect(() => {
+    const updateState = () =>
+      writableStore(store).setState(runtime.threadManager.getState(), true);
+    updateState();
+    return runtime.threadManager.subscribe(updateState);
+  }, [runtime, store]);
+
+  return store;
+};
+
+export const AssistantRuntimeProviderImpl: FC<
+  AssistantRuntimeProvider.Props
+> = ({ children, runtime }) => {
+  const useAssistantRuntime = useAssistantRuntimeStore(runtime);
+  const useToolUIs = useAssistantToolUIsStore();
+  const useThreadManager = useThreadManagerStore(runtime);
+  const context = useMemo(() => {
+    return {
+      useToolUIs,
+      useAssistantRuntime,
+      useAssistantActions: useAssistantRuntime,
+      useThreadManager,
+    };
+  }, [useAssistantRuntime, useToolUIs, useThreadManager]);
+
+  return (
+    <AssistantContext.Provider value={context}>
+      <ThreadRuntimeProvider runtime={runtime.thread}>
+        {children}
+      </ThreadRuntimeProvider>
+    </AssistantContext.Provider>
+  );
+};
+
+export const AssistantRuntimeProvider = memo(AssistantRuntimeProviderImpl);

package/src/context/providers/AttachmentRuntimeProvider.tsx

@@ -0,0 +1,60 @@
+"use client";
+
+import {
+  type FC,
+  type PropsWithChildren,
+  useEffect,
+  useMemo,
+  useState,
+} from "react";
+import { create } from "zustand";
+import { AttachmentContext } from "../react/AttachmentContext";
+import { writableStore } from "../ReadonlyStore";
+import { AttachmentRuntime } from "../../api/AttachmentRuntime";
+
+type AttachmentRuntimeProviderProps = PropsWithChildren<{
+  runtime: AttachmentRuntime;
+}>;
+
+const useAttachmentRuntimeStore = (runtime: AttachmentRuntime) => {
+  const [store] = useState(() => create(() => runtime));
+
+  useEffect(() => {
+    writableStore(store).setState(runtime, true);
+  }, [runtime, store]);
+
+  return store;
+};
+const useAttachmentStore = (runtime: AttachmentRuntime) => {
+  const [store] = useState(() => create(() => runtime.getState()));
+  useEffect(() => {
+    const updateState = () =>
+      writableStore(store).setState(runtime.getState(), true);
+    updateState();
+    return runtime.subscribe(updateState);
+  }, [runtime, store]);
+
+  return store;
+};
+
+export const AttachmentRuntimeProvider: FC<AttachmentRuntimeProviderProps> = ({
+  runtime,
+  children,
+}) => {
+  const useAttachmentRuntime = useAttachmentRuntimeStore(runtime);
+  const useAttachment = useAttachmentStore(runtime);
+  const source = useAttachment((s) => s.source);
+  const context = useMemo(() => {
+    return {
+      source,
+      useAttachmentRuntime,
+      useAttachment,
+    };
+  }, [useAttachmentRuntime, useAttachment]);
+
+  return (
+    <AttachmentContext.Provider value={context}>
+      {children}
+    </AttachmentContext.Provider>
+  );
+};

package/src/context/providers/ContentPartRuntimeProvider.tsx

@@ -0,0 +1,51 @@
+"use client";
+
+import { type FC, type PropsWithChildren, useEffect, useState } from "react";
+import { create } from "zustand";
+import { ContentPartContext } from "../react/ContentPartContext";
+import type { ContentPartContextValue } from "../react/ContentPartContext";
+import { writableStore } from "../ReadonlyStore";
+import { ContentPartRuntime } from "../../api/ContentPartRuntime";
+
+type ContentPartProviderProps = PropsWithChildren<{
+  runtime: ContentPartRuntime;
+}>;
+
+const useContentPartRuntimeStore = (runtime: ContentPartRuntime) => {
+  const [store] = useState(() => create(() => runtime));
+
+  useEffect(() => {
+    writableStore(store).setState(runtime, true);
+  }, [runtime, store]);
+
+  return store;
+};
+
+export const useContentPartStore = (runtime: ContentPartRuntime) => {
+  const [store] = useState(() => create(() => runtime.getState()));
+  useEffect(() => {
+    const updateState = () =>
+      writableStore(store).setState(runtime.getState(), true);
+    updateState();
+    return runtime.subscribe(updateState);
+  }, [runtime, store]);
+
+  return store;
+};
+
+export const ContentPartRuntimeProvider: FC<ContentPartProviderProps> = ({
+  runtime,
+  children,
+}) => {
+  const useContentPartRuntime = useContentPartRuntimeStore(runtime);
+  const useContentPart = useContentPartStore(runtime);
+  const [context] = useState<ContentPartContextValue>(() => {
+    return { useContentPartRuntime, useContentPart };
+  });
+
+  return (
+    <ContentPartContext.Provider value={context}>
+      {children}
+    </ContentPartContext.Provider>
+  );
+};

package/src/context/providers/MessageRuntimeProvider.tsx

@@ -0,0 +1,74 @@
+"use client";
+
+import { type FC, type PropsWithChildren, useEffect, useState } from "react";
+import { create } from "zustand";
+import { MessageContext } from "../react/MessageContext";
+import type { MessageContextValue } from "../react/MessageContext";
+import { makeMessageUtilsStore } from "../stores/MessageUtils";
+import { ReadonlyStore, writableStore } from "../ReadonlyStore";
+import { MessageRuntime } from "../../api/MessageRuntime";
+
+type MessageProviderProps = PropsWithChildren<{
+  runtime: MessageRuntime;
+}>;
+
+const useMessageRuntimeStore = (runtime: MessageRuntime) => {
+  const [store] = useState(() => create(() => runtime));
+
+  useEffect(() => {
+    writableStore(store).setState(runtime, true);
+  }, [runtime, store]);
+
+  return store;
+};
+
+const useMessageStore = (runtime: MessageRuntime) => {
+  const [store] = useState(() => create(() => runtime.getState()));
+  useEffect(() => {
+    const updateState = () =>
+      writableStore(store).setState(runtime.getState(), true);
+    updateState();
+    return runtime.subscribe(updateState);
+  }, [runtime, store]);
+
+  return store;
+};
+
+const useMessageUtilsStore = () => {
+  const [store] = useState(() => makeMessageUtilsStore());
+  return store;
+};
+
+const useEditComposerStore = (
+  useMessageRuntime: ReadonlyStore<MessageRuntime>,
+) => {
+  const runtime = useMessageRuntime.getState().composer;
+  const [store] = useState(() => create(() => runtime.getState()));
+
+  useEffect(() => {
+    const updateState = () => writableStore(store).setState(runtime.getState());
+    updateState();
+    return runtime.subscribe(updateState);
+  }, [runtime, store]);
+
+  return store;
+};
+
+export const MessageRuntimeProvider: FC<MessageProviderProps> = ({
+  runtime,
+  children,
+}) => {
+  const useMessageRuntime = useMessageRuntimeStore(runtime);
+  const useMessage = useMessageStore(runtime);
+  const useMessageUtils = useMessageUtilsStore();
+  const useEditComposer = useEditComposerStore(useMessageRuntime);
+  const [context] = useState<MessageContextValue>(() => {
+    return { useMessageRuntime, useMessage, useMessageUtils, useEditComposer };
+  });
+
+  return (
+    <MessageContext.Provider value={context}>
+      {children}
+    </MessageContext.Provider>
+  );
+};

package/src/context/providers/TextContentPartProvider.tsx

@@ -0,0 +1,75 @@
+"use client";
+
+import type { FC, PropsWithChildren } from "react";
+import { useEffect, useState } from "react";
+import { create } from "zustand";
+import {
+  ContentPartContext,
+  ContentPartContextValue,
+} from "../react/ContentPartContext";
+import { ContentPartStatus, TextContentPart } from "../../types/AssistantTypes";
+import { writableStore } from "../ReadonlyStore";
+import {
+  ContentPartRuntimeImpl,
+  ContentPartState,
+} from "../../api/ContentPartRuntime";
+
+type TextContentPartProviderProps = {
+  text: string;
+  isRunning?: boolean | undefined;
+};
+
+const COMPLETE_STATUS: ContentPartStatus = {
+  type: "complete",
+};
+
+const RUNNING_STATUS: ContentPartStatus = {
+  type: "running",
+};
+
+export const TextContentPartProvider: FC<
+  PropsWithChildren<TextContentPartProviderProps>
+> = ({ children, text, isRunning }) => {
+  const [context] = useState<ContentPartContextValue>(() => {
+    const useContentPartRuntime = create(
+      // TODO
+      () => new ContentPartRuntimeImpl(null as any, null as any, null as any),
+    );
+    const useContentPart = create<ContentPartState>(() => ({
+      status: isRunning ? RUNNING_STATUS : COMPLETE_STATUS,
+      part: { type: "text", text },
+      type: "text",
+      text,
+    }));
+
+    return { useContentPartRuntime, useContentPart };
+  });
+
+  useEffect(() => {
+    const state = context.useContentPart.getState() as ContentPartState & {
+      type: "text";
+    };
+
+    const textUpdated = (state as TextContentPart).text !== text;
+    const targetStatus = isRunning ? RUNNING_STATUS : COMPLETE_STATUS;
+    const statusUpdated = state.status !== targetStatus;
+
+    if (!textUpdated && !statusUpdated) return;
+
+    writableStore(context.useContentPart).setState(
+      {
+        type: "text",
+        text,
+        part: { type: "text", text },
+        status: targetStatus,
+      } satisfies ContentPartState,
+      true,
+    );
+  }, [context, isRunning, text]);
+
+  return (
+    <ContentPartContext.Provider value={context}>
+      {children}
+    </ContentPartContext.Provider>
+  );
+};

package/src/context/providers/ThreadRuntimeProvider.tsx

@@ -0,0 +1,87 @@
+import type { FC, PropsWithChildren } from "react";
+import { useEffect, useMemo, useState } from "react";
+import type { ThreadContextValue } from "../react/ThreadContext";
+import { ThreadContext } from "../react/ThreadContext";
+import { makeThreadViewportStore } from "../stores/ThreadViewport";
+import { writableStore } from "../ReadonlyStore";
+import { ThreadRuntime } from "../../api/ThreadRuntime";
+import { create } from "zustand";
+import { ThreadComposerRuntime } from "../../api/ComposerRuntime";
+
+type ThreadProviderProps = {
+  runtime: ThreadRuntime;
+};
+
+const useThreadRuntimeStore = (runtime: ThreadRuntime) => {
+  const [store] = useState(() => create(() => runtime));
+
+  useEffect(() => {
+    writableStore(store).setState(runtime, true);
+  }, [runtime, store]);
+
+  return store;
+};
+
+const useThreadStore = (runtime: ThreadRuntime) => {
+  const [store] = useState(() => create(() => runtime.getState()));
+  useEffect(() => {
+    const updateState = () =>
+      writableStore(store).setState(runtime.getState(), true);
+    updateState();
+    return runtime.subscribe(updateState);
+  }, [runtime, store]);
+
+  return store;
+};
+
+const useThreadMessagesStore = (runtime: ThreadRuntime) => {
+  const [store] = useState(() => create(() => runtime.messages));
+
+  useEffect(() => {
+    const updateState = () =>
+      writableStore(store).setState(runtime.messages, true);
+    updateState();
+    return runtime.subscribe(updateState);
+  }, [runtime, store]);
+
+  return store;
+};
+
+const useThreadComposerStore = (runtime: ThreadComposerRuntime) => {
+  const [store] = useState(() => create(() => runtime.getState()));
+
+  useEffect(() => {
+    const updateState = () =>
+      writableStore(store).setState(runtime.getState(), true);
+    updateState();
+    return runtime.subscribe(updateState);
+  }, [runtime, store]);
+
+  return store;
+};
+
+export const ThreadRuntimeProvider: FC<
+  PropsWithChildren<ThreadProviderProps>
+> = ({ children, runtime }) => {
+  const useThreadRuntime = useThreadRuntimeStore(runtime);
+  const useThread = useThreadStore(runtime);
+  const useThreadMessages = useThreadMessagesStore(runtime);
+  const useThreadComposer = useThreadComposerStore(runtime.composer);
+
+  const context = useMemo<ThreadContextValue>(() => {
+    const useViewport = makeThreadViewportStore();
+
+    return {
+      useThread,
+      useThreadRuntime,
+      useThreadMessages,
+      useThreadActions: useThreadRuntime,
+      useComposer: useThreadComposer,
+      useViewport,
+    };
+  }, [useThread, useThreadRuntime, useThreadMessages, useThreadComposer]);
+
+  return (
+    <ThreadContext.Provider value={context}>{children}</ThreadContext.Provider>
+  );
+};

package/src/context/providers/index.ts

@@ -0,0 +1,2 @@
+export { AssistantRuntimeProvider } from "./AssistantRuntimeProvider";
+export { TextContentPartProvider } from "./TextContentPartProvider";

package/src/context/react/AssistantContext.ts

@@ -0,0 +1,74 @@
+"use client";
+
+import { createContext } from "react";
+import type { AssistantToolUIsState } from "../stores/AssistantToolUIs";
+import { ReadonlyStore } from "../ReadonlyStore";
+import { createContextHook } from "./utils/createContextHook";
+import { createContextStoreHook } from "./utils/createContextStoreHook";
+import { UseBoundStore } from "zustand";
+import { AssistantRuntime } from "../../api/AssistantRuntime";
+import { ThreadManagerState } from "../../api/ThreadManagerRuntime";
+
+export type AssistantContextValue = {
+  useToolUIs: UseBoundStore<ReadonlyStore<AssistantToolUIsState>>;
+  useAssistantRuntime: UseBoundStore<ReadonlyStore<AssistantRuntime>>;
+  useThreadManager: UseBoundStore<ReadonlyStore<ThreadManagerState>>;
+
+  /**
+   * @deprecated Use `useAssistantRuntime` instead. This will be removed in 0.6.0.
+   */
+  useAssistantActions: UseBoundStore<ReadonlyStore<AssistantRuntime>>;
+};
+
+export const AssistantContext = createContext<AssistantContextValue | null>(
+  null,
+);
+
+export const useAssistantContext = createContextHook(
+  AssistantContext,
+  "AssistantRuntimeProvider",
+);
+
+export function useAssistantRuntime(options?: {
+  optional?: false | undefined;
+}): AssistantRuntime;
+export function useAssistantRuntime(options?: {
+  optional?: boolean | undefined;
+}): AssistantRuntime | null;
+export function useAssistantRuntime(options?: {
+  optional?: boolean | undefined;
+}) {
+  const context = useAssistantContext(options);
+  if (!context) return null;
+  return context.useAssistantRuntime();
+}
+
+export const actions = createContextStoreHook(
+  useAssistantContext,
+  "useAssistantActions",
+);
+
+/**
+ * @deprecated Use `useAssistantRuntime` instead. This will be removed in 0.6.0.
+ */
+export const useAssistantActionsStore = actions.useAssistantActionsStore;
+
+/**
+ * @deprecated Use `useAssistantRuntime` instead. This will be removed in 0.6.0.
+ */
+export const useAssistantActions = actions.useAssistantActions;
+
+/**
+ * @deprecated Use `useAssistantRuntime` instead. This will be removed in 0.6.0.
+ */
+export const useAssistantRuntimeStore = useAssistantActionsStore;
+
+export const { useToolUIs, useToolUIsStore } = createContextStoreHook(
+  useAssistantContext,
+  "useToolUIs",
+);
+
+export const { useThreadManager } = createContextStoreHook(
+  useAssistantContext,
+  "useThreadManager",
+);