@assistant-ui/store 0.0.0

package/src/types.ts

@@ -0,0 +1,120 @@
+import type { ResourceElement } from "@assistant-ui/tap";
+
+/**
+ * Definition of a scope in the assistant client (internal type)
+ * @template TValue - The API type (must include getState() and any actions)
+ * @template TSource - The parent scope name (or "root" for top-level scopes)
+ * @template TQuery - The query parameters needed to access this scope from its source
+ * @internal
+ */
+export type ScopeDefinition<
+  TValue = any,
+  TSource extends string | "root" = any,
+  TQuery = any,
+> = {
+  value: TValue;
+  source: TSource;
+  query: TQuery;
+};
+
+/**
+ * Module augmentation interface for assistant-ui store type extensions.
+ *
+ * @example
+ * ```typescript
+ * declare module "@assistant-ui/store" {
+ *   interface AssistantScopeRegistry {
+ *     foo: {
+ *       value: { getState: () => { bar: string }; updateBar: (bar: string) => void };
+ *       source: "root";
+ *       query: Record<string, never>;
+ *     };
+ *   }
+ * }
+ * ```
+ */
+// eslint-disable-next-line @typescript-eslint/no-empty-object-type
+export interface AssistantScopeRegistry {}
+
+export type AssistantScopes = keyof AssistantScopeRegistry extends never
+  ? Record<"ERROR: No scopes were defined", ScopeDefinition>
+  : { [K in keyof AssistantScopeRegistry]: AssistantScopeRegistry[K] };
+
+/**
+ * Helper type to extract the value type from a scope definition
+ */
+export type ScopeValue<T extends ScopeDefinition> = T["value"];
+
+/**
+ * Helper type to extract the source type from a scope definition
+ */
+export type ScopeSource<T extends ScopeDefinition> = T["source"];
+
+/**
+ * Helper type to extract the query type from a scope definition
+ */
+export type ScopeQuery<T extends ScopeDefinition> = T["query"];
+
+/**
+ * Type for a scope field - a function that returns the current API value,
+ * with source and query metadata attached
+ */
+export type ScopeField<T extends ScopeDefinition> = (() => ScopeValue<T>) &
+  (
+    | {
+        source: ScopeSource<T>;
+        query: ScopeQuery<T>;
+      }
+    | {
+        source: null;
+        query: null;
+      }
+  );
+
+/**
+ * Props passed to a derived scope resource element
+ */
+export type DerivedScopeProps<T extends ScopeDefinition> = {
+  get: (parent: AssistantClient) => ScopeValue<T>;
+  source: ScopeSource<T>;
+  query: ScopeQuery<T>;
+};
+
+/**
+ * Input type for scope definitions - ResourceElement that returns the API value
+ * Can optionally include source/query metadata via DerivedScope
+ */
+export type ScopeInput<T extends ScopeDefinition> = ResourceElement<{
+  api: ScopeValue<T>;
+}>;
+
+/**
+ * Map of scope names to their input definitions
+ */
+export type ScopesInput = {
+  [K in keyof AssistantScopes]?: ScopeInput<AssistantScopes[K]>;
+};
+
+/**
+ * Unsubscribe function type
+ */
+export type Unsubscribe = () => void;
+
+/**
+ * State type extracted from all scopes
+ */
+export type AssistantState = {
+  [K in keyof AssistantScopes]: ReturnType<
+    AssistantScopes[K]["value"]["getState"]
+  >;
+};
+
+/**
+ * The assistant client type with all registered scopes
+ */
+export type AssistantClient = {
+  [K in keyof AssistantScopes]: ScopeField<AssistantScopes[K]>;
+} & {
+  subscribe(listener: () => void): Unsubscribe;
+  flushSync(): void;
+};

package/src/useAssistantClient.tsx

@@ -0,0 +1,250 @@
+import { useMemo } from "react";
+import { useResource } from "@assistant-ui/tap/react";
+import {
+  resource,
+  tapMemo,
+  tapResource,
+  tapResources,
+  tapEffectEvent,
+  ResourceElement,
+} from "@assistant-ui/tap";
+import type {
+  AssistantClient,
+  AssistantScopes,
+  ScopesInput,
+  ScopeField,
+  ScopeInput,
+  DerivedScopeProps,
+} from "./types";
+import { asStore } from "./asStore";
+import { useAssistantContextValue } from "./AssistantContext";
+import { splitScopes } from "./utils/splitScopes";
+
+/**
+ * Resource for a single root scope
+ * Returns a tuple of [scopeName, {scopeFunction, subscribe, flushSync}]
+ */
+const RootScopeResource = resource(
+  <K extends keyof AssistantScopes>({
+    scopeName,
+    element,
+  }: {
+    scopeName: K;
+    element: ScopeInput<AssistantScopes[K]>;
+  }) => {
+    const store = tapResource(asStore(element));
+
+    return tapMemo(() => {
+      const scopeFunction = (() => store.getState().api) as ScopeField<
+        AssistantScopes[K]
+      >;
+      scopeFunction.source = "root";
+      scopeFunction.query = {} as AssistantScopes[K]["query"];
+
+      return [
+        scopeName,
+        {
+          scopeFunction,
+          subscribe: store.subscribe,
+          flushSync: store.flushSync,
+        },
+      ] as const;
+    }, [scopeName, store]);
+  },
+);
+
+/**
+ * Resource for all root scopes
+ * Mounts each root scope and returns an object mapping scope names to their stores
+ */
+const RootScopesResource = resource((scopes: ScopesInput) => {
+  const resultEntries = tapResources(
+    Object.entries(scopes).map(([scopeName, element]) =>
+      RootScopeResource(
+        {
+          scopeName: scopeName as keyof AssistantScopes,
+          element: element as ScopeInput<
+            AssistantScopes[keyof AssistantScopes]
+          >,
+        },
+        { key: scopeName },
+      ),
+    ),
+  );
+
+  return tapMemo(() => {
+    if (resultEntries.length === 0) {
+      return {
+        scopes: {},
+      };
+    }
+
+    return {
+      scopes: Object.fromEntries(
+        resultEntries.map(([scopeName, { scopeFunction }]) => [
+          scopeName,
+          scopeFunction,
+        ]),
+      ) as {
+        [K in keyof typeof scopes]: ScopeField<AssistantScopes[K]>;
+      },
+      subscribe: (callback: () => void) => {
+        const unsubscribes = resultEntries.map(([, { subscribe }]) => {
+          return subscribe(() => {
+            console.log("Callback called for");
+            callback();
+          });
+        });
+        return () => {
+          unsubscribes.forEach((unsubscribe) => unsubscribe());
+        };
+      },
+      flushSync: () => {
+        resultEntries.forEach(([, { flushSync }]) => {
+          flushSync();
+        });
+      },
+    };
+  }, [...resultEntries]);
+});
+
+/**
+ * Hook to mount and access root scopes
+ */
+export const useRootScopes = (rootScopes: ScopesInput) => {
+  return useResource(RootScopesResource(rootScopes));
+};
+
+/**
+ * Resource for a single derived scope
+ * Returns a tuple of [scopeName, scopeFunction] where scopeFunction has source and query
+ */
+const DerivedScopeResource = resource(
+  <K extends keyof AssistantScopes>({
+    scopeName,
+    element,
+    parentClient,
+  }: {
+    scopeName: K;
+    element: ResourceElement<
+      AssistantScopes[K],
+      DerivedScopeProps<AssistantScopes[K]>
+    >;
+    parentClient: AssistantClient;
+  }) => {
+    const get = tapEffectEvent(element.props.get);
+    const source = element.props.source;
+    const query = element.props.query;
+    return tapMemo(() => {
+      const scopeFunction = (() => get(parentClient)) as ScopeField<
+        AssistantScopes[K]
+      >;
+      scopeFunction.source = source;
+      scopeFunction.query = query;
+
+      return [scopeName, scopeFunction] as const;
+    }, [scopeName, get, source, JSON.stringify(query), parentClient]);
+  },
+);
+
+/**
+ * Resource for all derived scopes
+ * Builds stable scope functions with source and query metadata
+ */
+const DerivedScopesResource = resource(
+  ({
+    scopes,
+    parentClient,
+  }: {
+    scopes: ScopesInput;
+    parentClient: AssistantClient;
+  }) => {
+    const resultEntries = tapResources(
+      Object.entries(scopes).map(([scopeName, element]) =>
+        DerivedScopeResource(
+          {
+            scopeName: scopeName as keyof AssistantScopes,
+            element: element as ScopeInput<
+              AssistantScopes[keyof AssistantScopes]
+            >,
+            parentClient,
+          },
+          { key: scopeName },
+        ),
+      ),
+    );
+
+    return tapMemo(() => {
+      return Object.fromEntries(resultEntries) as {
+        [K in keyof typeof scopes]: ScopeField<AssistantScopes[K]>;
+      };
+    }, [...resultEntries]);
+  },
+);
+
+/**
+ * Hook to mount and access derived scopes
+ */
+export const useDerivedScopes = (
+  derivedScopes: ScopesInput,
+  parentClient: AssistantClient,
+) => {
+  return useResource(
+    DerivedScopesResource({ scopes: derivedScopes, parentClient }),
+  );
+};
+
+const useExtendedAssistantClientImpl = (
+  scopes: ScopesInput,
+): AssistantClient => {
+  const baseClient = useAssistantContextValue();
+  const { rootScopes, derivedScopes } = splitScopes(scopes);
+
+  // Mount the scopes to keep them alive
+  const rootFields = useRootScopes(rootScopes);
+  const derivedFields = useDerivedScopes(derivedScopes, baseClient);
+
+  return useMemo(() => {
+    // Merge base client with extended client
+    // If baseClient is the default proxy, spreading it will be a no-op
+    return {
+      ...baseClient,
+      ...rootFields.scopes,
+      ...derivedFields,
+      subscribe: rootFields.subscribe ?? baseClient.subscribe,
+      flushSync: rootFields.flushSync ?? baseClient.flushSync,
+    } as AssistantClient;
+  }, [baseClient, rootFields, derivedFields]);
+};
+
+/**
+ * Hook to access or extend the AssistantClient
+ *
+ * @example Without config - returns the client from context:
+ * ```typescript
+ * const client = useAssistantClient();
+ * const fooState = client.foo.getState();
+ * ```
+ *
+ * @example With config - creates a new client with additional scopes:
+ * ```typescript
+ * const client = useAssistantClient({
+ *   message: DerivedScope({
+ *     source: "thread",
+ *     query: { type: "index", index: 0 },
+ *     get: () => messageApi,
+ *   }),
+ * });
+ * ```
+ */
+export function useAssistantClient(): AssistantClient;
+export function useAssistantClient(scopes: ScopesInput): AssistantClient;
+export function useAssistantClient(scopes?: ScopesInput): AssistantClient {
+  if (scopes) {
+    // eslint-disable-next-line react-hooks/rules-of-hooks
+    return useExtendedAssistantClientImpl(scopes);
+  } else {
+    // eslint-disable-next-line react-hooks/rules-of-hooks
+    return useAssistantContextValue();
+  }
+}

package/src/useAssistantState.tsx

@@ -0,0 +1,80 @@
+import { useMemo, useSyncExternalStore, useDebugValue } from "react";
+import type { AssistantClient, AssistantState } from "./types";
+import { useAssistantClient } from "./useAssistantClient";
+
+/**
+ * Proxied state that lazily accesses scope states
+ */
+class ProxiedAssistantState {
+  #client: AssistantClient;
+
+  constructor(client: AssistantClient) {
+    this.#client = client;
+  }
+
+  #getScope<K extends keyof AssistantState>(key: K): AssistantState[K] {
+    const scopeField = this.#client[key];
+    if (!scopeField) {
+      throw new Error(`Scope "${String(key)}" not found in client`);
+    }
+
+    const api = scopeField();
+    const state = api.getState();
+    return state as AssistantState[K];
+  }
+
+  // Create a Proxy to dynamically handle property access
+  static create(client: AssistantClient): AssistantState {
+    const instance = new ProxiedAssistantState(client);
+    return new Proxy(instance, {
+      get(target, prop) {
+        if (typeof prop === "string" && prop in client) {
+          return target.#getScope(prop as keyof AssistantState);
+        }
+        return undefined;
+      },
+    }) as unknown as AssistantState;
+  }
+}
+
+/**
+ * Hook to access a slice of the assistant state with automatic subscription
+ *
+ * @param selector - Function to select a slice of the state
+ * @returns The selected state slice
+ *
+ * @example
+ * ```typescript
+ * const client = useAssistantClient({
+ *   foo: RootScope({ ... }),
+ * });
+ *
+ * const bar = useAssistantState((state) => state.foo.bar);
+ * ```
+ */
+export const useAssistantState = <T,>(
+  selector: (state: AssistantState) => T,
+): T => {
+  const client = useAssistantClient();
+
+  const proxiedState = useMemo(
+    () => ProxiedAssistantState.create(client),
+    [client],
+  );
+
+  const slice = useSyncExternalStore(
+    client.subscribe,
+    () => selector(proxiedState),
+    () => selector(proxiedState),
+  );
+
+  useDebugValue(slice);
+
+  if (slice instanceof ProxiedAssistantState) {
+    throw new Error(
+      "You tried to return the entire AssistantState. This is not supported due to technical limitations.",
+    );
+  }
+
+  return slice;
+};

package/src/utils/splitScopes.ts

@@ -0,0 +1,38 @@
+import { DerivedScope } from "../DerivedScope";
+import type { AssistantScopes, ScopeInput, ScopesInput } from "../types";
+
+/**
+ * Splits a scopes object into root scopes and derived scopes.
+ *
+ * @param scopes - The scopes input object to split
+ * @returns An object with { rootScopes, derivedScopes }
+ *
+ * @example
+ * ```typescript
+ * const scopes = {
+ *   foo: RootScope({ ... }),
+ *   bar: DerivedScope({ ... }),
+ * };
+ *
+ * const { rootScopes, derivedScopes } = splitScopes(scopes);
+ * // rootScopes = { foo: ... }
+ * // derivedScopes = { bar: ... }
+ * ```
+ */
+export function splitScopes(scopes: ScopesInput) {
+  const rootScopes: ScopesInput = {};
+  const derivedScopes: ScopesInput = {};
+
+  for (const [key, scopeElement] of Object.entries(scopes) as [
+    keyof ScopesInput,
+    ScopeInput<AssistantScopes[keyof ScopesInput]>,
+  ][]) {
+    if (scopeElement.type === DerivedScope) {
+      derivedScopes[key] = scopeElement;
+    } else {
+      rootScopes[key] = scopeElement;
+    }
+  }
+
+  return { rootScopes, derivedScopes };
+}