@assistant-ui/store 0.2.0 → 0.2.1

package/src/tapClientResource.ts

@@ -0,0 +1,187 @@
+import {
+  tapEffect,
+  tapMemo,
+  tapRef,
+  type ResourceElement,
+  tapResource,
+} from "@assistant-ui/tap";
+import type { ClientMethods } from "./types/client";
+import {
+  tapClientStack,
+  tapWithClientStack,
+  SYMBOL_CLIENT_INDEX,
+} from "./utils/tap-client-stack-context";
+import {
+  BaseProxyHandler,
+  handleIntrospectionProp,
+} from "./utils/BaseProxyHandler";
+import { wrapperResource } from "./wrapperResource";
+
+/**
+ * Symbol used internally to get state from ClientProxy.
+ * This allows getState() to be optional in the user-facing client.
+ */
+const SYMBOL_GET_OUTPUT = Symbol("assistant-ui.store.getValue");
+
+type ClientInternal = {
+  [SYMBOL_GET_OUTPUT]: ClientMethods;
+};
+
+export const getClientState = (client: ClientMethods) => {
+  const output = (client as unknown as ClientInternal)[SYMBOL_GET_OUTPUT];
+  if (!output) {
+    throw new Error(
+      "Client scope contains a non-client resource. " +
+        "Ensure your Derived get() returns a client created with tapClientResource(), not a plain resource.",
+    );
+  }
+  return (output as any).getState?.();
+};
+
+// Global cache for function templates by field name
+const fieldAccessFns = new Map<
+  string | symbol,
+  (this: unknown, ...args: unknown[]) => unknown
+>();
+
+function getOrCreateProxyFn(prop: string | symbol) {
+  let template = fieldAccessFns.get(prop);
+  if (!template) {
+    template = function (this: unknown, ...args: unknown[]) {
+      if (!this || typeof this !== "object") {
+        throw new Error(
+          `Method "${String(prop)}" called without proper context. ` +
+            `This may indicate the function was called incorrectly.`,
+        );
+      }
+
+      const output = (this as ClientInternal)[SYMBOL_GET_OUTPUT];
+      if (!output) {
+        throw new Error(
+          `Method "${String(prop)}" called on invalid client proxy. ` +
+            `Ensure you are calling this method on a valid client instance.`,
+        );
+      }
+
+      const method = output[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);
+  }
+  return template;
+}
+
+class ClientProxyHandler
+  extends BaseProxyHandler
+  implements ProxyHandler<object>
+{
+  private boundFns: Map<string | symbol, Function> | undefined;
+  private cachedReceiver: unknown;
+
+  constructor(
+    private readonly outputRef: {
+      current: ClientMethods;
+    },
+    private readonly index: number,
+  ) {
+    super();
+  }
+
+  get(_: unknown, prop: string | symbol, receiver: unknown) {
+    if (prop === SYMBOL_GET_OUTPUT) return this.outputRef.current;
+    if (prop === SYMBOL_CLIENT_INDEX) return this.index;
+    const introspection = handleIntrospectionProp(prop, "ClientProxy");
+    if (introspection !== false) return introspection;
+    const value = this.outputRef.current[prop];
+    if (typeof value === "function") {
+      if (this.cachedReceiver !== receiver) {
+        this.boundFns = new Map();
+        this.cachedReceiver = receiver;
+      }
+      let bound = this.boundFns!.get(prop);
+      if (!bound) {
+        bound = getOrCreateProxyFn(prop).bind(receiver);
+        this.boundFns!.set(prop, bound);
+      }
+      return bound;
+    }
+    return value;
+  }
+
+  ownKeys(): ArrayLike<string | symbol> {
+    return Object.keys(this.outputRef.current);
+  }
+
+  has(_: unknown, prop: string | symbol) {
+    if (prop === SYMBOL_GET_OUTPUT) return true;
+    if (prop === SYMBOL_CLIENT_INDEX) return true;
+    return prop in this.outputRef.current;
+  }
+}
+
+/**
+ * Resource that wraps a plain resource element to create a stable client proxy.
+ *
+ * Takes a ResourceElement that returns methods (with optional getState()) and
+ * wraps it to produce a stable client proxy. This adds the client to the
+ * client stack, enabling event scoping.
+ *
+ * @internal
+ */
+export const ClientResource = wrapperResource(
+  <TMethods extends ClientMethods>(
+    element: ResourceElement<TMethods>,
+  ): {
+    methods: TMethods;
+    state: unknown;
+    key: string | number | undefined;
+  } => {
+    const valueRef = tapRef(null as unknown as TMethods);
+
+    const index = tapClientStack().length;
+    const methods = tapMemo(
+      () =>
+        new Proxy<TMethods>(
+          {} as TMethods,
+          new ClientProxyHandler(valueRef, index),
+        ),
+      [index],
+    );
+
+    const value = tapWithClientStack(methods, () => tapResource(element));
+    if (!valueRef.current) {
+      valueRef.current = value;
+    }
+
+    tapEffect(() => {
+      valueRef.current = value;
+    });
+
+    const state = (value as any).getState?.();
+    return { methods, state, key: element.key };
+  },
+);
+
+type InferClientState<TMethods> = TMethods extends {
+  getState: () => infer S;
+}
+  ? S
+  : undefined;
+
+export const tapClientResource = <TMethods extends ClientMethods>(
+  element: ResourceElement<TMethods>,
+): {
+  state: InferClientState<TMethods>;
+  methods: TMethods;
+  key: string | number | undefined;
+} => {
+  return tapResource(ClientResource(element)) as {
+    state: InferClientState<TMethods>;
+    methods: TMethods;
+    key: string | number | undefined;
+  };
+};

package/src/types/client.ts

@@ -0,0 +1,180 @@
+import type { ResourceElement } from "@assistant-ui/tap";
+import type {
+  AssistantEventName,
+  AssistantEventCallback,
+  AssistantEventSelector,
+} from "./events";
+
+/**
+ * Base type for methods that can be called on a client.
+ */
+export interface ClientMethods {
+  [key: string | symbol]: (...args: any[]) => any;
+}
+
+type ClientMetaType = { source: ClientNames; query: Record<string, unknown> };
+
+/**
+ * Schema of a client in the assistant system (internal type).
+ * @template TState - The state type for this client
+ * @template TMethods - The methods available on this client
+ * @template TMeta - Source/query metadata (optional)
+ * @template TEvents - Events that this client can emit (optional)
+ * @internal
+ */
+export type ClientSchema<
+  TMethods extends ClientMethods = ClientMethods,
+  TMeta extends ClientMetaType = never,
+  TEvents extends Record<string, unknown> = never,
+> = {
+  methods: TMethods;
+  meta?: TMeta;
+  events?: TEvents;
+};
+
+/**
+ * Module augmentation interface for assistant-ui store type extensions.
+ *
+ * @example
+ * ```typescript
+ * declare module "@assistant-ui/store" {
+ *   interface ScopeRegistry {
+ *     // Simple client (meta and events are optional)
+ *     foo: {
+ *       methods: {
+ *         getState: () => { bar: string };
+ *         updateBar: (bar: string) => void;
+ *       };
+ *     };
+ *     // Full client with meta and events
+ *     bar: {
+ *       methods: {
+ *         getState: () => { id: string };
+ *         update: () => void;
+ *       };
+ *       meta: { source: "fooList"; query: { index: number } };
+ *       events: {
+ *         "bar.updated": { id: string };
+ *       };
+ *     };
+ *   }
+ * }
+ * ```
+ */
+export interface ScopeRegistry {}
+
+type ClientEventsType<K extends ClientNames> = Record<
+  `${K}.${string}`,
+  unknown
+>;
+
+type ClientError<E extends string> = {
+  methods: Record<E, () => E>;
+  meta: { source: ClientNames; query: Record<E, E> };
+  events: Record<`${E}.`, E>;
+};
+
+type ValidateClient<K extends keyof ScopeRegistry> = ScopeRegistry[K] extends {
+  methods: ClientMethods;
+}
+  ? "meta" extends keyof ScopeRegistry[K]
+    ? ScopeRegistry[K]["meta"] extends ClientMetaType
+      ? "events" extends keyof ScopeRegistry[K]
+        ? ScopeRegistry[K]["events"] extends ClientEventsType<K>
+          ? ScopeRegistry[K]
+          : ClientError<`ERROR: ${K & string} has invalid events type`>
+        : ScopeRegistry[K]
+      : ClientError<`ERROR: ${K & string} has invalid meta type`>
+    : "events" extends keyof ScopeRegistry[K]
+      ? ScopeRegistry[K]["events"] extends ClientEventsType<K>
+        ? ScopeRegistry[K]
+        : ClientError<`ERROR: ${K & string} has invalid events type`>
+      : ScopeRegistry[K]
+  : ClientError<`ERROR: ${K & string} has invalid methods type`>;
+
+type ClientSchemas = keyof ScopeRegistry extends never
+  ? {
+      "ERROR: No clients were defined": ClientError<"ERROR: No clients were defined">;
+    }
+  : { [K in keyof ScopeRegistry]: ValidateClient<K> };
+
+/**
+ * Output type that client resources return (just methods).
+ *
+ * @example
+ * ```typescript
+ * const FooResource = resource((): ClientResourceOutput<"foo"> => {
+ *   const [state, setState] = tapState({ bar: "hello" });
+ *   return {
+ *     getState: () => state,
+ *     updateBar: (b) => setState({ bar: b }),
+ *   };
+ * });
+ * ```
+ */
+export type ClientOutput<K extends ClientNames> = ClientSchemas[K]["methods"] &
+  ClientMethods;
+
+export type ClientNames = keyof ClientSchemas extends infer U ? U : never;
+
+export type ClientEvents<K extends ClientNames> =
+  "events" extends keyof ClientSchemas[K]
+    ? ClientSchemas[K]["events"] extends ClientEventsType<K>
+      ? ClientSchemas[K]["events"]
+      : never
+    : never;
+
+export type ClientMeta<K extends ClientNames> =
+  "meta" extends keyof ClientSchemas[K]
+    ? Pick<
+        ClientSchemas[K]["meta"] extends ClientMetaType
+          ? ClientSchemas[K]["meta"]
+          : never,
+        "source" | "query"
+      >
+    : never;
+
+export type ClientElement<K extends ClientNames> = ResourceElement<
+  ClientOutput<K>
+>;
+
+/**
+ * Unsubscribe function type.
+ */
+export type Unsubscribe = () => void;
+
+/**
+ * State type extracted from all clients via their getState() methods.
+ */
+export type AssistantState = {
+  [K in ClientNames]: ClientSchemas[K]["methods"] extends {
+    getState: () => infer S;
+  }
+    ? S
+    : never;
+};
+
+/**
+ * Type for a client accessor - a function that returns the methods,
+ * with source/query metadata attached (derived from meta).
+ */
+export type AssistantClientAccessor<K extends ClientNames> =
+  (() => ClientSchemas[K]["methods"]) &
+    (
+      | ClientMeta<K>
+      | { source: "root"; query: Record<string, never> }
+      | { source: null; query: null }
+    ) & { name: K };
+
+/**
+ * The assistant client type with all registered clients.
+ */
+export type AssistantClient = {
+  [K in ClientNames]: AssistantClientAccessor<K>;
+} & {
+  subscribe(listener: () => void): Unsubscribe;
+  on<TEvent extends AssistantEventName>(
+    selector: AssistantEventSelector<TEvent>,
+    callback: AssistantEventCallback<TEvent>,
+  ): Unsubscribe;
+};

package/src/types/events.ts

@@ -0,0 +1,77 @@
+import type {
+  AssistantClientAccessor,
+  ClientEvents,
+  ClientNames,
+} from "./client";
+
+// --- Event Map Construction ---
+type UnionToIntersection<U> = (
+  U extends unknown
+    ? (x: U) => void
+    : never
+) extends (x: infer I) => void
+  ? I
+  : never;
+
+type ClientEventMap = UnionToIntersection<
+  { [K in ClientNames]: ClientEvents<K> }[ClientNames]
+>;
+
+// --- Core Types ---
+
+type WildcardPayload = {
+  [K in keyof ClientEventMap]: { event: K; payload: ClientEventMap[K] };
+}[keyof ClientEventMap];
+
+export type AssistantEventPayload = ClientEventMap & { "*": WildcardPayload };
+
+export type AssistantEventName = keyof AssistantEventPayload;
+
+type EventSource<T extends AssistantEventName> =
+  T extends `${infer Source}.${string}` ? Source : never;
+
+// --- Scoping ---
+
+type ParentOf<K extends ClientNames> =
+  AssistantClientAccessor<K> extends { source: infer S }
+    ? S extends ClientNames
+      ? S
+      : never
+    : never;
+
+type AncestorsOf<
+  K extends ClientNames,
+  Seen extends ClientNames = never,
+> = K extends Seen
+  ? never
+  : ParentOf<K> extends never
+    ? never
+    : ParentOf<K> | AncestorsOf<ParentOf<K>, Seen | K>;
+
+/** Valid scopes: `"*"` | event source | ancestors of event source */
+export type AssistantEventScope<TEvent extends AssistantEventName> =
+  | "*"
+  | EventSource<TEvent>
+  | (EventSource<TEvent> extends ClientNames
+      ? AncestorsOf<EventSource<TEvent>>
+      : never);
+
+// --- Selection & Callbacks ---
+
+export type AssistantEventSelector<TEvent extends AssistantEventName> =
+  | TEvent
+  | { scope: AssistantEventScope<TEvent>; event: TEvent };
+
+export const normalizeEventSelector = <TEvent extends AssistantEventName>(
+  selector: AssistantEventSelector<TEvent>,
+) => {
+  if (typeof selector === "string") {
+    const source = selector.split(".")[0] as AssistantEventScope<TEvent>;
+    return { scope: source, event: selector };
+  }
+  return { scope: selector.scope, event: selector.event };
+};
+
+export type AssistantEventCallback<TEvent extends AssistantEventName> = (
+  payload: AssistantEventPayload[TEvent],
+) => void;

package/src/useAui.ts

@@ -18,32 +18,32 @@ import type {
   ClientNames,
   ClientElement,
   ClientMeta,
-} from "@assistant-ui/core/store";
+} from "./types/client";
+import { Derived, DerivedElement } from "./Derived";
+import {
+  useAssistantContextValue,
+  DefaultAssistantClient,
+  createRootAssistantClient,
+} from "./utils/react-assistant-context";
+import {
+  DerivedClients,
+  RootClients,
+  tapSplitClients,
+} from "./utils/splitClients";
 import {
-  Derived,
-  type DerivedElement,
-  type ScopesConfig,
   normalizeEventSelector,
   type AssistantEventName,
   type AssistantEventCallback,
   type AssistantEventSelector,
-  tapClientResource,
-} from "@assistant-ui/core/store";
+} from "./types/events";
+import { NotificationManager } from "./utils/NotificationManager";
+import { withAssistantTapContextProvider } from "./utils/tap-assistant-context";
+import { tapClientResource } from "./tapClientResource";
+import { getClientIndex } from "./utils/tap-client-stack-context";
 import {
-  NotificationManager,
-  withAssistantTapContextProvider,
-  getClientIndex,
   PROXIED_ASSISTANT_STATE_SYMBOL,
   createProxiedAssistantState,
-  type RootClients,
-  type DerivedClients,
-  tapSplitClients,
-} from "@assistant-ui/core/store/internal";
-import {
-  useAssistantContextValue,
-  DefaultAssistantClient,
-  createRootAssistantClient,
-} from "./utils/react-assistant-context";
+} from "./utils/proxied-assistant-state";
 
 const tapShallowMemoArray = <T>(array: readonly T[]) => {
   // biome-ignore lint/correctness/useExhaustiveDependencies: shallow memo
@@ -354,7 +354,9 @@ export const AssistantClientResource = resource(
 );
 
 export namespace useAui {
-  export type Props = ScopesConfig;
+  export type Props = {
+    [K in ClientNames]?: ClientElement<K> | DerivedElement<K>;
+  };
 }
 
 export function useAui(): AssistantClient;

package/src/useAuiEvent.ts

@@ -5,8 +5,8 @@ import type {
   AssistantEventName,
   AssistantEventCallback,
   AssistantEventSelector,
-} from "@assistant-ui/core/store";
-import { normalizeEventSelector } from "@assistant-ui/core/store";
+} from "./types/events";
+import { normalizeEventSelector } from "./types/events";
 
 export const useAuiEvent = <TEvent extends AssistantEventName>(
   selector: AssistantEventSelector<TEvent>,

package/src/useAuiState.ts

@@ -1,7 +1,7 @@
 import { useSyncExternalStore, useDebugValue } from "react";
-import type { AssistantState } from "@assistant-ui/core/store";
-import { getProxiedAssistantState } from "@assistant-ui/core/store/internal";
+import type { AssistantState } from "./types/client";
 import { useAui } from "./useAui";
+import { getProxiedAssistantState } from "./utils/proxied-assistant-state";
 
 /**
  * Hook to access a slice of the assistant state with automatic subscription
@@ -15,7 +15,7 @@ import { useAui } from "./useAui";
  *   foo: RootScope({ ... }),
  * });
  *
- * const bar = useAuiState((state) => state.foo.bar);
+ * const bar = useAuiState((s) => s.foo.bar);
  * ```
  */
 export const useAuiState = <T>(selector: (state: AssistantState) => T): T => {

package/src/utils/BaseProxyHandler.ts

@@ -0,0 +1,50 @@
+const INTROSPECTION_PROPS = new Set(["$$typeof", "nodeType", "then"]);
+
+/**
+ * Handles common proxy introspection properties.
+ * Returns the appropriate value for toStringTag, toJSON, and props that should return undefined.
+ * Returns `false` if the prop should be handled by the subclass.
+ */
+export const handleIntrospectionProp = (
+  prop: string | symbol,
+  name: string,
+): unknown | false => {
+  if (prop === Symbol.toStringTag) return name;
+  if (typeof prop === "symbol") return undefined;
+  if (prop === "toJSON") return () => name;
+  if (INTROSPECTION_PROPS.has(prop)) return undefined;
+  return false;
+};
+
+export abstract class BaseProxyHandler implements ProxyHandler<object> {
+  abstract get(_: unknown, prop: string | symbol, receiver?: unknown): unknown;
+  abstract ownKeys(): ArrayLike<string | symbol>;
+  abstract has(_: unknown, prop: string | symbol): boolean;
+
+  getOwnPropertyDescriptor(_: unknown, prop: string | symbol) {
+    const value = this.get(_, prop);
+    if (value === undefined) return undefined;
+    return {
+      value,
+      writable: false,
+      enumerable: true,
+      configurable: false,
+    };
+  }
+
+  set() {
+    return false;
+  }
+  setPrototypeOf() {
+    return false;
+  }
+  defineProperty() {
+    return false;
+  }
+  deleteProperty() {
+    return false;
+  }
+  preventExtensions(): boolean {
+    return false;
+  }
+}

package/src/utils/NotificationManager.ts

@@ -0,0 +1,114 @@
+import { resource } from "@assistant-ui/tap";
+import type { ClientStack } from "./tap-client-stack-context";
+import type {
+  AssistantEventName,
+  AssistantEventPayload,
+} from "../types/events";
+import { Unsubscribe } from "../types/client";
+import { tapConst } from "@assistant-ui/tap";
+
+type InternalCallback = (payload: unknown, clientStack: ClientStack) => void;
+
+export type NotificationManager = {
+  on<TEvent extends AssistantEventName>(
+    event: TEvent,
+    callback: (
+      payload: AssistantEventPayload[TEvent],
+      clientStack: ClientStack,
+    ) => void,
+  ): Unsubscribe;
+  emit<TEvent extends Exclude<AssistantEventName, "*">>(
+    event: TEvent,
+    payload: AssistantEventPayload[TEvent],
+    clientStack: ClientStack,
+  ): void;
+  subscribe(callback: () => void): Unsubscribe;
+  notifySubscribers(): void;
+};
+
+export const NotificationManager = resource((): NotificationManager => {
+  return tapConst(() => {
+    const listeners = new Map<string, Set<InternalCallback>>();
+    const wildcardListeners = new Set<InternalCallback>();
+    const subscribers = new Set<() => void>();
+
+    return {
+      on(event, callback) {
+        const cb = callback as InternalCallback;
+        if (event === "*") {
+          wildcardListeners.add(cb);
+          return () => wildcardListeners.delete(cb);
+        }
+
+        let set = listeners.get(event);
+        if (!set) {
+          set = new Set();
+          listeners.set(event, set);
+        }
+        set.add(cb);
+
+        return () => {
+          set!.delete(cb);
+          if (set!.size === 0) listeners.delete(event);
+        };
+      },
+
+      emit(event, payload, clientStack) {
+        const eventListeners = listeners.get(event);
+        if (!eventListeners && wildcardListeners.size === 0) return;
+
+        queueMicrotask(() => {
+          const errors = [];
+          if (eventListeners) {
+            for (const cb of eventListeners) {
+              try {
+                cb(payload, clientStack);
+              } catch (e) {
+                errors.push(e);
+              }
+            }
+          }
+          if (wildcardListeners.size > 0) {
+            const wrapped = { event, payload };
+            for (const cb of wildcardListeners) {
+              try {
+                cb(wrapped, clientStack);
+              } catch (e) {
+                errors.push(e);
+              }
+            }
+          }
+
+          if (errors.length > 0) {
+            if (errors.length === 1) {
+              throw errors[0];
+            } else {
+              for (const error of errors) {
+                console.error(error);
+              }
+              throw new AggregateError(
+                errors,
+                "Errors occurred during event emission",
+              );
+            }
+          }
+        });
+      },
+
+      subscribe(callback) {
+        subscribers.add(callback);
+        return () => subscribers.delete(callback);
+      },
+
+      notifySubscribers() {
+        for (const cb of subscribers) {
+          try {
+            cb();
+          } catch (e) {
+            console.error("NotificationManager: subscriber callback error", e);
+          }
+        }
+      },
+    };
+  }, []);
+});