@apollo/client-ai-apps 0.5.2 → 0.5.3

package/src/openai/react/hooks/createHydrationUtils.ts

@@ -0,0 +1,182 @@
+import { useState, useCallback, useRef, useMemo, useLayoutEffect } from "react";
+import type {
+  DocumentNode,
+  OperationVariables,
+  TypedDocumentNode,
+} from "@apollo/client";
+import { useApolloClient } from "./useApolloClient.js";
+import { useToolName } from "./useToolName.js";
+import { isReactive } from "../../../react/reactive.js";
+import type { Reactive } from "../../../react/reactive.js";
+import { equal } from "@wry/equality";
+import { __DEV__ } from "@apollo/client/utilities/environment";
+import {
+  getToolNamesFromDocument,
+  getVariableNamesFromDocument,
+} from "../../../utilities/index.js";
+
+type HydratedVariablesInput<TVariables> = {
+  [K in keyof TVariables]: TVariables[K] | Reactive<TVariables[K]>;
+};
+
+type StateVariables<TVariables, Input> = {
+  [K in keyof TVariables as K extends keyof Input ?
+    Input[K] extends Reactive<any> ?
+      never
+    : K
+  : K]: K extends keyof TVariables ? TVariables[K] : never;
+};
+
+type SetVariables<TState> = (
+  update: Partial<TState> | ((prev: TState) => Partial<TState>)
+) => void;
+
+/** @experimental */
+export function createHydrationUtils<
+  TVariables extends OperationVariables = OperationVariables,
+>(document: TypedDocumentNode<any, TVariables> | DocumentNode) {
+  const documentToolNames = getToolNamesFromDocument(document);
+  const variableNames = getVariableNamesFromDocument(document);
+
+  function useHydratedVariables<
+    TInputVariables extends HydratedVariablesInput<TVariables>,
+  >(
+    variables: TInputVariables &
+      Record<Exclude<keyof TInputVariables, keyof TVariables>, never>
+  ): [
+    variables: TVariables,
+    setVariables: SetVariables<StateVariables<TVariables, TInputVariables>>,
+  ] {
+    const client = useApolloClient();
+    const toolName = useToolName();
+    const [toolInput] = useState(() => client.toolInput);
+
+    const toolMatches =
+      toolInput !== undefined &&
+      toolName !== undefined &&
+      documentToolNames.has(toolName);
+
+    const [stateVars, setStateVars] = useState<Record<string, unknown>>(() => {
+      const values: Record<string, unknown> = {};
+
+      for (const [key, value] of Object.entries(
+        toolMatches ? toolInput : variables
+      )) {
+        if (variableNames.has(key) && !isReactive(value)) {
+          values[key] = value;
+        }
+      }
+
+      return values;
+    });
+
+    const [initialReactiveValues] = useState<Record<string, unknown>>(() => {
+      const values: Record<string, unknown> = {};
+
+      for (const [key, value] of Object.entries(variables)) {
+        if (variableNames.has(key) && isReactive(value)) {
+          values[key] = value.value;
+        }
+      }
+
+      return values;
+    });
+
+    const [reactiveVars, setReactiveVars] = useState(() => {
+      const values: Record<string, unknown> = {};
+
+      for (const [key, value] of Object.entries(initialReactiveValues)) {
+        if (toolMatches && key in toolInput) {
+          values[key] = toolInput[key];
+        } else if (!toolMatches) {
+          values[key] = value;
+        }
+      }
+
+      return values;
+    });
+
+    const changedKeysRef = useRef(new Set<string>());
+    const nextReactiveVars: Record<string, unknown> = {};
+
+    for (const [key, value] of Object.entries(variables)) {
+      if (!variableNames.has(key) || !isReactive(value)) continue;
+
+      const hasChanged =
+        changedKeysRef.current.has(key) ||
+        !equal(value.value, initialReactiveValues[key]);
+
+      if (toolMatches && !hasChanged) {
+        if (key in toolInput) {
+          nextReactiveVars[key] = toolInput[key];
+        }
+      } else {
+        nextReactiveVars[key] = value.value;
+      }
+    }
+
+    if (!equal(nextReactiveVars, reactiveVars)) {
+      setReactiveVars(nextReactiveVars);
+    }
+
+    // Clear the tool input after first mount so that remounting the component
+    // uses the user-provided variables rather than the hydrated tool input.
+    // This runs once on mount; watchQuery also clears it when useQuery is
+    // present, so both paths are idempotent.
+    useLayoutEffect(() => {
+      if (toolMatches) {
+        client.clearToolInput();
+      }
+      // eslint-disable-next-line react-hooks/exhaustive-deps
+    }, []);
+
+    useLayoutEffect(() => {
+      for (const [key, value] of Object.entries(variables)) {
+        if (
+          variableNames.has(key) &&
+          isReactive(value) &&
+          !changedKeysRef.current.has(key) &&
+          !equal(value.value, initialReactiveValues[key])
+        ) {
+          changedKeysRef.current.add(key);
+        }
+      }
+    });
+
+    const resolvedVariables = useMemo(() => {
+      return { ...stateVars, ...reactiveVars } as TVariables;
+    }, [stateVars, reactiveVars]);
+
+    const setVariables = useCallback<
+      SetVariables<StateVariables<TVariables, TInputVariables>>
+    >((update) => {
+      setStateVars((prev) => {
+        const updates =
+          typeof update === "function" ? update(prev as any) : update;
+
+        const filtered = Object.fromEntries(
+          Object.entries(updates).filter(([key]) => {
+            if (key in initialReactiveValues) {
+              if (__DEV__) {
+                console.warn(
+                  `Attempted to set reactive variable "${key}" via setVariables. ` +
+                    `Reactive variables are read-only and are ignored. `
+                );
+              }
+              return false;
+            }
+            return true;
+          })
+        );
+
+        if (Object.keys(filtered).length === 0) return prev;
+
+        return { ...prev, ...filtered };
+      });
+    }, []);
+
+    return [resolvedVariables, setVariables];
+  }
+
+  return { useHydratedVariables };
+}

package/src/openai/react/index.ts

@@ -2,4 +2,5 @@ export { useApp } from "./hooks/useApp.js";
 export { useToolName } from "./hooks/useToolName.js";
 export { useToolMetadata } from "./hooks/useToolMetadata.js";
 export { useToolInput } from "./hooks/useToolInput.js";
+export { createHydrationUtils } from "./hooks/createHydrationUtils.js";
 export { useWidgetState } from "./hooks/useWidgetState.js";

package/src/react/ApolloProvider.tsx

@@ -25,7 +25,7 @@ export function ApolloProvider({ children, client }: ApolloProvider.Props) {
     );
   }
 
-  use(client.waitForInitialization());
+  use(client.connect());
 
   return (
     <BaseApolloProvider client={client as BaseApolloClient}>

package/src/react/index.mcp.ts

@@ -2,6 +2,7 @@ export * from "./index.js";
 export {
   useApp,
   useToolInput,
+  createHydrationUtils,
   useToolMetadata,
   useToolName,
 } from "../mcp/react/index.js";

package/src/react/index.openai.ts

@@ -2,6 +2,7 @@ export * from "./index.js";
 export {
   useApp,
   useToolInput,
+  createHydrationUtils,
   useToolMetadata,
   useToolName,
 } from "../openai/react/index.js";

package/src/react/index.ts

@@ -1,6 +1,8 @@
 import { missingHook } from "./missingHook.js";
 
 export { ApolloProvider } from "./ApolloProvider.js";
+export { reactive } from "./reactive.js";
+export type { Reactive } from "./reactive.js";
 
 // Use `mcp` related types since these are the most common between the two
 // targets
@@ -17,3 +19,8 @@ export const useToolMetadata =
 
 export const useToolName =
   missingHook<typeof import("./index.mcp.js").useToolName>("useToolName");
+
+/** @experimental */
+export const createHydrationUtils = missingHook<
+  typeof import("./index.mcp.js").createHydrationUtils
+>("createHydrationUtils");

package/src/react/reactive.ts

@@ -0,0 +1,19 @@
+const REACTIVE = Symbol("apollo.reactive");
+
+export interface Reactive<T> {
+  readonly [REACTIVE]: true;
+  readonly value: T;
+}
+
+export function reactive<T>(value: T): Reactive<T> {
+  return { [REACTIVE]: true, value };
+}
+
+export function isReactive(value: unknown): value is Reactive<unknown> {
+  return (
+    typeof value === "object" &&
+    value !== null &&
+    REACTIVE in value &&
+    (value as Record<symbol, unknown>)[REACTIVE] === true
+  );
+}

package/src/testing/internal/mcp/graphqlToolResult.ts

@@ -2,20 +2,20 @@ import type { McpUiToolResultNotification } from "@modelcontextprotocol/ext-apps
 import type { FormattedExecutionResult } from "graphql";
 import type { ApolloMcpServerApps } from "../../../core/types";
 
-export function graphqlToolResult(
+export function graphqlToolResult<TData = Record<string, unknown>>(
   options:
-    | FormattedExecutionResult
+    | FormattedExecutionResult<TData>
     | {
         prefetch?: Record<string, FormattedExecutionResult>;
-        result: FormattedExecutionResult;
+        result: FormattedExecutionResult<TData>;
       }
 ): McpUiToolResultNotification["params"] {
   let structuredContent: ApolloMcpServerApps.StructuredContent;
 
   if ("data" in options) {
-    structuredContent = { result: options };
+    structuredContent = { result: options as FormattedExecutionResult };
   } else if ("result" in options) {
-    structuredContent = { result: options.result };
+    structuredContent = { result: options.result as FormattedExecutionResult };
 
     if (options.prefetch) {
       structuredContent.prefetch = options.prefetch;

package/src/utilities/getToolNamesFromDocument.ts

@@ -0,0 +1,15 @@
+import { getOperationDefinition } from "@apollo/client/utilities/internal";
+import { Kind, type DocumentNode } from "graphql";
+
+export function getToolNamesFromDocument(document: DocumentNode) {
+  const operationDef = getOperationDefinition(document);
+
+  return new Set(
+    operationDef?.directives
+      ?.filter((d) => d.name.value === "tool")
+      .flatMap((d) => {
+        const nameArg = d.arguments?.find((arg) => arg.name.value === "name");
+        return nameArg?.value.kind === Kind.STRING ? [nameArg.value.value] : [];
+      })
+  );
+}

package/src/utilities/getVariableNamesFromDocument.ts

@@ -0,0 +1,9 @@
+import { getOperationDefinition } from "@apollo/client/utilities/internal";
+import type { DocumentNode } from "graphql";
+
+export function getVariableNamesFromDocument(document: DocumentNode) {
+  const operationDef = getOperationDefinition(document);
+  return new Set(
+    operationDef?.variableDefinitions?.map((v) => v.variable.name.value) ?? []
+  );
+}

package/src/utilities/index.ts

@@ -1,5 +1,8 @@
 export { aiClientSymbol } from "./constants.js";
 export { getVariablesForOperationFromToolInput } from "./getVariablesForOperationFromToolInput.js";
+export { getToolNamesFromDocument } from "./getToolNamesFromDocument.js";
+export { getVariableNamesFromDocument } from "./getVariableNamesFromDocument.js";
 export { cacheAsync } from "./cacheAsync.js";
 export { invariant } from "./invariant.js";
 export { promiseWithResolvers } from "./promiseWithResolvers.js";
+export { warnOnVariableMismatch } from "./warnOnVariableMismatch.js";

package/src/utilities/warnOnVariableMismatch.ts

@@ -0,0 +1,20 @@
+import type { DocumentNode, OperationVariables } from "@apollo/client";
+import { getOperationName } from "@apollo/client/utilities/internal/internal.cjs";
+
+export function warnOnVariableMismatch(
+  document: DocumentNode,
+  toolInput: OperationVariables,
+  actualVariables: OperationVariables | undefined
+) {
+  const operationName = getOperationName(document, "(anonymous)");
+
+  console.warn(
+    `The operation "${operationName}" has a @tool directive matching the current ` +
+      "tool call, but the variables provided to the query don't match the tool " +
+      "input. Use the `useHydratedVariables` hook returned from `createHydrationUtils` " +
+      "to provide the hydrated variables to the query. " +
+      "\n\nExpected variables:\n%o\n\nReceived:\n%o",
+    toolInput,
+    actualVariables
+  );
+}