@assistant-ui/store 0.0.6 → 0.1.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@assistant-ui/store",
-  "version": "0.0.6",
+  "version": "0.1.0",
   "description": "Tap-based state management for @assistant-ui",
   "keywords": [
     "state-management",
@@ -28,7 +28,7 @@
   ],
   "sideEffects": false,
   "dependencies": {
-    "@assistant-ui/tap": "^0.3.6"
+    "@assistant-ui/tap": "^0.4.0"
   },
   "peerDependencies": {
     "@types/react": "*",

package/src/AuiIf.tsx

@@ -0,0 +1,17 @@
+"use client";
+
+import type { FC, PropsWithChildren } from "react";
+import { useAuiState } from "./useAuiState";
+import type { AssistantState } from "./types/client";
+
+export namespace AuiIf {
+  export type Props = PropsWithChildren<{ condition: AuiIf.Condition }>;
+  export type Condition = (state: AssistantState) => boolean;
+}
+
+export const AuiIf: FC<AuiIf.Props> = ({ children, condition }) => {
+  const result = useAuiState(condition);
+  return result ? children : null;
+};
+
+AuiIf.displayName = "AuiIf";

package/src/Derived.ts

@@ -16,7 +16,7 @@ import type {
  *
  * @example
  * ```typescript
- * const aui = useAssistantClient({
+ * const aui = useAui({
  *   message: Derived({
  *     source: "thread",
  *     query: { index: 0 },

package/src/index.ts

@@ -1,11 +1,11 @@
 // hooks
-export { useAssistantClient } from "./useAssistantClient";
-export { useAssistantState } from "./useAssistantState";
-export { useAssistantEvent } from "./useAssistantEvent";
+export { useAui } from "./useAui";
+export { useAuiState } from "./useAuiState";
+export { useAuiEvent } from "./useAuiEvent";
 
 // components
-export { AssistantIf } from "./AssistantIf";
-export { AssistantProvider } from "./utils/react-assistant-context";
+export { AuiIf } from "./AuiIf";
+export { AuiProvider } from "./utils/react-assistant-context";
 
 // resources
 export { Derived } from "./Derived";

package/src/tapClientList.ts

@@ -1,23 +1,24 @@
-import { tapState } from "@assistant-ui/tap";
+import { tapConst, tapState, withKey } from "@assistant-ui/tap";
 import type { ContravariantResource } from "@assistant-ui/tap";
 import { tapClientLookup } from "./tapClientLookup";
 import type { ClientMethods, ClientOutputOf } from "./types/client";
 
+type DataHandle<TData> = { data: TData | undefined; hasData: boolean };
+
 const createProps = <TData>(
   key: string,
-  data: TData,
+  data: DataHandle<TData>,
   remove: () => void,
 ): tapClientList.ResourceProps<TData> => {
-  let initialData: { data: TData } | undefined = { data };
   return {
     key,
     getInitialData: () => {
-      if (!initialData) {
-        throw new Error("getInitialData may only be called once");
+      if (!data.hasData) {
+        throw new Error(
+          "getInitialData may only be called during initial render",
+        );
       }
-      const data = initialData.data;
-      initialData = undefined;
-      return data;
+      return data.data!;
     },
     remove,
   };
@@ -34,13 +35,16 @@ export const tapClientList = <TData, TState, TMethods extends ClientMethods>(
 
   type Props = tapClientList.ResourceProps<TData>;
 
+  const initialDataHandles: DataHandle<TData>[] = tapConst(() => [], []);
+
   const [items, setItems] = tapState<Record<string, Props>>(() => {
     const entries: [string, Props][] = [];
     for (const data of initialValues) {
       const key = getKey(data);
+      const handle = { data, hasData: true };
       entries.push([
         key,
-        createProps(key, data, () => {
+        createProps(key, handle, () => {
           setItems((items) => {
             const newItems = { ...items };
             delete newItems[key];
@@ -48,16 +52,22 @@ export const tapClientList = <TData, TState, TMethods extends ClientMethods>(
           });
         }),
       ]);
+      initialDataHandles.push(handle);
     }
     return Object.fromEntries(entries);
   });
 
-  const lookup = tapClientLookup<TState, TMethods, Record<string, Props>>(
-    items,
-    Resource,
-    [Resource],
+  const lookup = tapClientLookup<TState, TMethods>(
+    () =>
+      Object.values(items).map((props) => withKey(props.key, Resource(props))),
+    [items, Resource],
   );
 
+  initialDataHandles.forEach((handle) => {
+    handle.data = undefined;
+    handle.hasData = false;
+  });
+
   const add = (data: TData) => {
     const key = getKey(data);
     setItems((items) => {
@@ -67,9 +77,12 @@ export const tapClientList = <TData, TState, TMethods extends ClientMethods>(
         );
       }
 
+      const handle = { data, hasData: true };
+      initialDataHandles.push(handle);
+
       return {
         ...items,
-        [key]: createProps(key, data, () => {
+        [key]: createProps(key, handle, () => {
           setItems((items) => {
             const newItems = { ...items };
             delete newItems[key];

package/src/tapClientLookup.ts

@@ -1,56 +1,76 @@
-import { ResourceElement, tapMemo, tapResources } from "@assistant-ui/tap";
+import {
+  tapInlineResource,
+  tapMemo,
+  tapResources,
+  type ResourceElement,
+} from "@assistant-ui/tap";
 import type { ClientMethods, ClientOutputOf } from "./types/client";
 import { ClientResource } from "./tapClientResource";
+import { wrapperResource } from "./wrapperResource";
 
-export const tapClientLookup = <
-  TState,
-  TMethods extends ClientMethods,
-  M extends Record<string | number | symbol, any>,
->(
-  map: M,
-  getElement: (
-    t: M[keyof M],
-    key: keyof M,
-  ) => ResourceElement<ClientOutputOf<TState, TMethods>>,
-  getElementDeps: any[],
+const ClientResourceWithKey = wrapperResource(
+  <TState, TMethods extends ClientMethods>(
+    el: ResourceElement<ClientOutputOf<TState, TMethods>>,
+  ) => {
+    if (el.key === undefined) {
+      throw new Error("tapClientResource: Element has no key");
+    }
+    return tapInlineResource(ClientResource(el)) as ClientOutputOf<
+      TState,
+      TMethods
+    > & { key: string | number };
+  },
+);
+
+export function tapClientLookup<TState, TMethods extends ClientMethods>(
+  getElements: () => readonly ResourceElement<
+    ClientOutputOf<TState, TMethods>
+  >[],
+  getElementsDeps: readonly unknown[],
 ): {
   state: TState[];
-  get: (lookup: { index: number } | { key: keyof M }) => TMethods;
-} => {
+  get: (lookup: { index: number } | { key: string }) => TMethods;
+} {
   const resources = tapResources(
-    map,
-    (t, key) => ClientResource(getElement(t, key)),
-    getElementDeps,
+    () => getElements().map((el) => ClientResourceWithKey(el)),
+    // biome-ignore lint/correctness/useExhaustiveDependencies: getElementsDeps is passed through from caller
+    getElementsDeps,
   );
-  const keys = tapMemo(() => Object.keys(map) as (keyof M)[], [map]);
+
+  const keys = tapMemo(() => Object.keys(resources), [resources]);
+
+  // For arrays, track element key -> index mapping
+  const keyToIndex = tapMemo(() => {
+    return resources.reduce(
+      (acc, resource, index) => {
+        acc[resource.key] = index;
+        return acc;
+      },
+      {} as Record<string, number>,
+    );
+  }, [resources]);
 
   const state = tapMemo(() => {
-    const result = new Array(keys.length);
-    for (let i = 0; i < keys.length; i++) {
-      result[i] = resources[keys[i]!].state;
-    }
-    return result;
-  }, [keys, resources]);
+    return resources.map((r) => r.state);
+  }, [resources]);
 
   return {
     state,
-    get: (lookup: { index: number } | { key: keyof M }) => {
+    get: (lookup: { index: number } | { key: string }) => {
       if ("index" in lookup) {
         if (lookup.index < 0 || lookup.index >= keys.length) {
           throw new Error(
             `tapClientLookup: Index ${lookup.index} out of bounds (length: ${keys.length})`,
           );
         }
-        return resources[keys[lookup.index]!]!.methods;
+        return resources[lookup.index]!.methods;
       }
 
-      const value = resources[lookup.key];
-      if (!value) {
-        throw new Error(
-          `tapClientLookup: Key "${String(lookup.key)}" not found`,
-        );
+      const index = keyToIndex[lookup.key];
+      if (index === undefined) {
+        throw new Error(`tapClientLookup: Key "${lookup.key}" not found`);
       }
-      return value.methods;
+      return resources[index]!.methods;
     },
   };
-};
+}

package/src/tapClientResource.ts

@@ -4,7 +4,6 @@ import {
   tapRef,
   type ResourceElement,
   tapResource,
-  resource,
   tapInlineResource,
 } from "@assistant-ui/tap";
 import type { ClientMethods, ClientOutputOf } from "./types/client";
@@ -17,6 +16,7 @@ import {
   BaseProxyHandler,
   handleIntrospectionProp,
 } from "./utils/BaseProxyHandler";
+import { wrapperResource } from "./wrapperResource";
 
 /**
  * Symbol used internally to get state from ClientProxy.
@@ -57,6 +57,8 @@ function getOrCreateProxyFn(prop: string | symbol) {
       const method = this[SYMBOL_GET_OUTPUT].methods[prop];
       if (!method)
         throw new Error(`Method "${String(prop)}" is not implemented.`);
+      if (typeof method !== "function")
+        throw new Error(`"${String(prop)}" is not a function.`);
       return method(...args);
     };
     fieldAccessFns.set(prop, template);
@@ -82,7 +84,9 @@ class ClientProxyHandler
     if (prop === SYMBOL_CLIENT_INDEX) return this.index;
     const introspection = handleIntrospectionProp(prop, "ClientProxy");
     if (introspection !== false) return introspection;
-    return getOrCreateProxyFn(prop);
+    const value = this.outputRef.current.methods[prop];
+    if (typeof value === "function") return getOrCreateProxyFn(prop);
+    return value;
   }
 
   ownKeys(): ArrayLike<string | symbol> {
@@ -114,10 +118,12 @@ class ClientProxyHandler
  * });
  * ```
  */
-export const ClientResource = resource(
+export const ClientResource = wrapperResource(
   <TState, TMethods extends ClientMethods>(
     element: ResourceElement<ClientOutputOf<TState, TMethods>>,
-  ): ClientOutputOf<TState, TMethods> => {
+  ): ClientOutputOf<TState, TMethods> & {
+    key: string | number | undefined;
+  } => {
     const valueRef = tapRef(
       null as unknown as ClientOutputOf<TState, TMethods>,
     );
@@ -129,7 +135,7 @@ export const ClientResource = resource(
           {} as TMethods,
           new ClientProxyHandler(valueRef, index),
         ),
-      [],
+      [index],
     );
 
     const value = tapWithClientStack(methods, () => tapResource(element));
@@ -141,7 +147,7 @@ export const ClientResource = resource(
       valueRef.current = value;
     });
 
-    return { methods, state: value.state };
+    return { methods, state: value.state, key: element.key };
   },
 );
 

package/src/types/client.ts

@@ -9,7 +9,7 @@ import type {
  * Base type for methods that can be called on a client.
  */
 export interface ClientMethods {
-  [key: string | symbol]: (...args: any[]) => any;
+  [key: string | symbol]: ((...args: any[]) => any) | ClientMethods;
 }
 
 type ClientMetaType = { source: ClientNames; query: Record<string, unknown> };
@@ -170,7 +170,7 @@ export type AssistantClientAccessor<K extends ClientNames> =
       | ClientMeta<K>
       | { source: "root"; query: Record<string, never> }
       | { source: null; query: null }
-    );
+    ) & { name: K };
 
 /**
  * The assistant client type with all registered clients.

package/src/{useAssistantClient.tsx → useAui.tsx}

@@ -10,6 +10,7 @@ import {
   tapEffect,
   tapRef,
   tapResource,
+  withKey,
 } from "@assistant-ui/tap";
 import type {
   AssistantClient,
@@ -28,7 +29,7 @@ import {
 import {
   DerivedClients,
   RootClients,
-  splitClients,
+  tapSplitClients,
 } from "./utils/splitClients";
 import {
   normalizeEventSelector,
@@ -59,7 +60,7 @@ const RootClientResource = resource(
       { clientRef, emit },
       () => tapClientResource(element),
     );
-    return tapMemo(() => ({ methods }), [state]);
+    return tapMemo(() => ({ state, methods }), [methods, state]);
   },
 );
 
@@ -68,10 +69,12 @@ const RootClientAccessorResource = resource(
     element,
     notifications,
     clientRef,
+    name,
   }: {
     element: ClientElement<K>;
     notifications: NotificationManager;
     clientRef: { parent: AssistantClient; current: AssistantClient | null };
+    name: K;
   }): AssistantClientAccessor<K> => {
     const store = tapInlineResource(
       StoreResource(
@@ -84,17 +87,33 @@ const RootClientAccessorResource = resource(
     }, [store, notifications]);
 
     return tapMemo(() => {
-      const clientFunction = () => store.getState().methods;
-      clientFunction.source = "root" as const;
-      clientFunction.query = {};
-      return clientFunction;
-    }, [store]);
+      const clientFunction = () => store.getValue().methods;
+      Object.defineProperties(clientFunction, {
+        source: {
+          value: "root" as const,
+          writable: false,
+        },
+        query: {
+          value: {} as Record<string, never>,
+          writable: false,
+        },
+        name: {
+          value: name,
+          configurable: true,
+        },
+      });
+      return clientFunction as AssistantClientAccessor<K>;
+    }, [store, name]);
   },
 );
 
 const NoOpRootClientsAccessorsResource = resource(() => {
   return tapMemo(
-    () => ({ clients: {}, subscribe: undefined, on: undefined }),
+    () => ({
+      clients: [] as AssistantClientAccessor<ClientNames>[],
+      subscribe: undefined,
+      on: undefined,
+    }),
     [],
   );
 });
@@ -115,14 +134,19 @@ const RootClientsAccessorsResource = resource(
     );
 
     const results = tapResources(
-      inputClients,
-      (element) =>
-        RootClientAccessorResource({
-          element: element!,
-          notifications,
-          clientRef,
-        }),
-      [notifications, clientRef],
+      () =>
+        Object.keys(inputClients).map((key) =>
+          withKey(
+            key,
+            RootClientAccessorResource({
+              element: inputClients[key as keyof typeof inputClients]!,
+              notifications,
+              clientRef,
+              name: key as keyof typeof inputClients,
+            }),
+          ),
+        ),
+      [inputClients, notifications, clientRef],
     );
 
     return tapMemo(() => {
@@ -136,7 +160,7 @@ const RootClientsAccessorsResource = resource(
         ) {
           if (!this) {
             throw new Error(
-              "const { on } = useAssistantClient() is not supported. Use aui.on() instead.",
+              "const { on } = useAui() is not supported. Use aui.on() instead.",
             );
           }
 
@@ -203,9 +227,11 @@ const DerivedClientAccessorResource = resource(
   <K extends ClientNames>({
     element,
     clientRef,
+    name,
   }: {
     element: DerivedElement<K>;
     clientRef: { parent: AssistantClient; current: AssistantClient | null };
+    name: K;
   }) => {
     const get = tapEffectEvent(() => element.props);
 
@@ -219,9 +245,13 @@ const DerivedClientAccessorResource = resource(
         query: {
           get: () => getMeta(get(), clientRef, metaMemo).query,
         },
+        name: {
+          value: name,
+          configurable: true,
+        },
       });
-      return clientFunction;
-    }, [clientRef]);
+      return clientFunction as AssistantClientAccessor<K>;
+    }, [clientRef, name]);
   },
 );
 
@@ -234,13 +264,18 @@ const DerivedClientsAccessorsResource = resource(
     clientRef: { parent: AssistantClient; current: AssistantClient | null };
   }) => {
     return tapResources(
-      clients,
-      (element) =>
-        DerivedClientAccessorResource({
-          element: element!,
-          clientRef,
-        }),
-      [clientRef],
+      () =>
+        Object.keys(clients).map((key) =>
+          withKey(
+            key,
+            DerivedClientAccessorResource({
+              element: clients[key as keyof typeof clients]!,
+              clientRef,
+              name: key as keyof typeof clients,
+            }),
+          ),
+        ),
+      [clients, clientRef],
     );
   },
 );
@@ -250,19 +285,26 @@ const DerivedClientsAccessorsResource = resource(
  */
 export const AssistantClientResource = resource(
   ({
-    baseClient,
+    parent,
     clients,
   }: {
-    baseClient: AssistantClient;
-    clients: useAssistantClient.Props;
+    parent: AssistantClient;
+    clients: useAui.Props;
   }): AssistantClient => {
-    const { rootClients, derivedClients } = splitClients(clients, baseClient);
+    const { rootClients, derivedClients } = tapSplitClients(clients, parent);
 
     const clientRef = tapRef({
-      parent: baseClient,
+      parent: parent,
       current: null as AssistantClient | null,
     }).current;
 
+    tapEffect(() => {
+      // if (clientRef.current && clientRef.current !== client)
+      //   throw new Error("clientRef.current !== client");
+
+      clientRef.current = client;
+    });
+
     const rootFields = tapResource(
       Object.keys(rootClients).length > 0
         ? RootClientsAccessorsResource({ clients: rootClients, clientRef })
@@ -276,46 +318,47 @@ export const AssistantClientResource = resource(
     const client = tapMemo(() => {
       // Swap DefaultAssistantClient -> createRootAssistantClient at root to change error message
       const proto =
-        baseClient === DefaultAssistantClient
+        parent === DefaultAssistantClient
           ? createRootAssistantClient()
-          : baseClient;
+          : parent;
+
       const client = Object.create(proto) as AssistantClient;
-      Object.assign(client, rootFields.clients, derivedFields, {
-        subscribe: rootFields.subscribe ?? baseClient.subscribe,
-        on: rootFields.on ?? baseClient.on,
+      Object.assign(client, {
+        subscribe: rootFields.subscribe ?? parent.subscribe,
+        on: rootFields.on ?? parent.on,
         [PROXIED_ASSISTANT_STATE_SYMBOL]: createProxiedAssistantState(client),
       });
+
+      for (const field of rootFields.clients) {
+        (client as any)[field.name] = field;
+      }
+      for (const field of derivedFields) {
+        (client as any)[field.name] = field;
+      }
+
       return client;
-    }, [baseClient, rootFields, derivedFields]);
+    }, [parent, rootFields, derivedFields]);
 
     if (clientRef.current === null) {
       clientRef.current = client;
     }
 
-    tapEffect(() => {
-      clientRef.current = client;
-    });
-
     return client;
   },
 );
 
-export namespace useAssistantClient {
+export namespace useAui {
   export type Props = {
     [K in ClientNames]?: ClientElement<K> | DerivedElement<K>;
   };
 }
 
-export function useAssistantClient(): AssistantClient;
-export function useAssistantClient(
-  clients: useAssistantClient.Props,
-): AssistantClient;
-export function useAssistantClient(
-  clients?: useAssistantClient.Props,
-): AssistantClient {
-  const baseClient = useAssistantContextValue();
+export function useAui(): AssistantClient;
+export function useAui(clients: useAui.Props): AssistantClient;
+export function useAui(clients?: useAui.Props): AssistantClient {
+  const parent = useAssistantContextValue();
   if (clients) {
-    return useResource(AssistantClientResource({ baseClient, clients }));
+    return useResource(AssistantClientResource({ parent: parent, clients }));
   }
-  return baseClient;
+  return parent;
 }

package/src/{useAssistantEvent.ts → useAuiEvent.ts}

@@ -1,5 +1,5 @@
 import { useEffect, useEffectEvent } from "react";
-import { useAssistantClient } from "./useAssistantClient";
+import { useAui } from "./useAui";
 import type {
   AssistantEventName,
   AssistantEventCallback,
@@ -7,11 +7,11 @@ import type {
 } from "./types/events";
 import { normalizeEventSelector } from "./types/events";
 
-export const useAssistantEvent = <TEvent extends AssistantEventName>(
+export const useAuiEvent = <TEvent extends AssistantEventName>(
   selector: AssistantEventSelector<TEvent>,
   callback: AssistantEventCallback<TEvent>,
 ) => {
-  const aui = useAssistantClient();
+  const aui = useAui();
   const callbackRef = useEffectEvent(callback);
 
   const { scope, event } = normalizeEventSelector(selector);

package/src/{useAssistantState.tsx → useAuiState.tsx}

@@ -1,6 +1,6 @@
 import { useSyncExternalStore, useDebugValue } from "react";
 import type { AssistantState } from "./types/client";
-import { useAssistantClient } from "./useAssistantClient";
+import { useAui } from "./useAui";
 import { getProxiedAssistantState } from "./utils/proxied-assistant-state";
 
 /**
@@ -11,17 +11,15 @@ import { getProxiedAssistantState } from "./utils/proxied-assistant-state";
  *
  * @example
  * ```typescript
- * const aui = useAssistantClient({
+ * const aui = useAui({
  *   foo: RootScope({ ... }),
  * });
  *
- * const bar = useAssistantState((state) => state.foo.bar);
+ * const bar = useAuiState((state) => state.foo.bar);
  * ```
  */
-export const useAssistantState = <T,>(
-  selector: (state: AssistantState) => T,
-): T => {
-  const aui = useAssistantClient();
+export const useAuiState = <T,>(selector: (state: AssistantState) => T): T => {
+  const aui = useAui();
   const proxiedState = getProxiedAssistantState(aui);
 
   const slice = useSyncExternalStore(