@apollo/client-ai-apps 0.5.2 → 0.5.3

package/src/mcp/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 [initialReactiveVars] = 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(initialReactiveVars)) {
+        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, initialReactiveVars[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, initialReactiveVars[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 initialReactiveVars) {
+              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/mcp/react/index.ts

@@ -2,3 +2,4 @@ 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";

package/src/openai/core/ApolloClient.ts

@@ -1,17 +1,26 @@
 import type { ApolloLink } from "@apollo/client";
 import { ApolloClient as BaseApolloClient } from "@apollo/client";
 import { DocumentTransform } from "@apollo/client";
+import type {
+  WatchQueryOptions,
+  ObservableQuery,
+  OperationVariables,
+} from "@apollo/client";
 import { removeDirectivesFromDocument } from "@apollo/client/utilities/internal";
 import { parse } from "graphql";
+import { equal } from "@wry/equality";
 import { __DEV__ } from "@apollo/client/utilities/environment";
 import type { ApplicationManifest } from "../../types/application-manifest.js";
 import { ToolCallLink } from "../link/ToolCallLink.js";
 import {
   aiClientSymbol,
   cacheAsync,
+  getToolNamesFromDocument,
   getVariablesForOperationFromToolInput,
+  warnOnVariableMismatch,
 } from "../../utilities/index.js";
 import { McpAppManager } from "./McpAppManager.js";
+import { getVariableNamesFromDocument } from "../../utilities/getVariableNamesFromDocument.js";
 
 export declare namespace ApolloClient {
   export interface Options extends Omit<BaseApolloClient.Options, "link"> {
@@ -27,6 +36,8 @@ export class ApolloClient extends BaseApolloClient {
   /** @internal */
   readonly [aiClientSymbol] = true;
 
+  #toolInput: Record<string, unknown> | undefined;
+
   constructor(options: ApolloClient.Options) {
     const link = options.link ?? new ToolCallLink();
 
@@ -55,9 +66,64 @@ export class ApolloClient extends BaseApolloClient {
     this.appManager.close().catch(() => {});
   }
 
-  waitForInitialization = cacheAsync(async () => {
+  get toolInput() {
+    return this.#toolInput;
+  }
+
+  clearToolInput() {
+    this.#toolInput = undefined;
+  }
+
+  watchQuery<
+    T = any,
+    TVariables extends OperationVariables = OperationVariables,
+  >(options: WatchQueryOptions<TVariables, T>): ObservableQuery<T, TVariables> {
+    if (__DEV__) {
+      const toolInput = this.#toolInput;
+
+      if (toolInput) {
+        const toolName = this.appManager.toolName;
+        const hasMatchingTool =
+          !!toolName && getToolNamesFromDocument(options.query).has(toolName);
+
+        if (hasMatchingTool) {
+          // Clear after first matching comparison so this only fires once and
+          // remounting doesn't produce spurious warnings.
+          this.#toolInput = undefined;
+
+          const variableNames = getVariableNamesFromDocument(options.query);
+
+          if (variableNames.size > 0) {
+            const { variables } = options;
+
+            const toolInputVariables = Object.entries(toolInput).filter(
+              ([key]) => variableNames.has(key)
+            );
+
+            const hasToolInputMismatch = toolInputVariables.some(
+              ([key, value]) => !equal(value, variables?.[key])
+            );
+
+            if (hasToolInputMismatch) {
+              warnOnVariableMismatch(
+                options.query,
+                Object.fromEntries(toolInputVariables),
+                variables
+              );
+            }
+          }
+        }
+      }
+    }
+
+    return super.watchQuery(options);
+  }
+
+  connect = cacheAsync(async () => {
     const { prefetch, result, toolName, args } =
-      await this.appManager.waitForInitialization();
+      await this.appManager.connect();
+
+    this.#toolInput = args;
 
     this.manifest.operations.forEach((operation) => {
       if (operation.prefetchID && prefetch?.[operation.prefetchID]) {

package/src/openai/core/McpAppManager.ts

@@ -35,7 +35,7 @@ export class McpAppManager {
     return this.#toolInput;
   }
 
-  waitForInitialization = cacheAsync(async () => {
+  connect = cacheAsync(async () => {
     let toolResult = promiseWithResolvers<ApolloMcpServerApps.CallToolResult>();
 
     this.app.ontoolresult = (params) => {
@@ -44,7 +44,7 @@ export class McpAppManager {
       );
     };
 
-    await this.connect();
+    await this.connectToHost();
 
     const { structuredContent } = await toolResult.promise;
 
@@ -55,12 +55,12 @@ export class McpAppManager {
     // before we get the tool result (which should always happen and at most
     // once according to the spec). Rather than relying on the
     // `ui/notifications/tool-input` notification to set the tool input value,
-    // we read from `window.openai.toolInput so that we hvae the most recent
+    // we read from `window.openai.toolInput so that we have the most recent
     // set value.
     //
     // When OpenAI fixes this issue and sends `ui/notifications/tool-input`
     // consistently according to the MCP Apps specification, this can be
-    // reverrted to use the `app.ontoolinput` callback.
+    // reverted to use the `app.ontoolinput` callback.
     this.#toolInput = window.openai.toolInput;
 
     // OpenAI doesn't provide access to `_meta`, so we need to use
@@ -93,7 +93,7 @@ export class McpAppManager {
     return result.structuredContent;
   }
 
-  private async connect() {
+  private async connectToHost() {
     try {
       return await this.app.connect(
         new PostMessageTransport(window.parent, window.parent)

package/src/openai/core/__tests__/ApolloClient.test.ts

@@ -1,9 +1,10 @@
 import { expect, test, describe, vi } from "vitest";
 import { ApolloClient } from "../ApolloClient.js";
 import { parse } from "graphql";
-import { ApolloLink, HttpLink, InMemoryCache } from "@apollo/client";
+import { ApolloLink, HttpLink, InMemoryCache, gql } from "@apollo/client";
 import { ToolCallLink } from "../../link/ToolCallLink.js";
 import {
+  graphqlToolResult,
   minimalHostContextWithToolName,
   mockApplicationManifest,
   mockMcpHost,
@@ -30,7 +31,7 @@ describe("Client Basics", () => {
       structuredContent: {},
     });
 
-    await client.waitForInitialization();
+    await client.connect();
 
     host.mockToolCall("execute", () => ({
       content: [],
@@ -124,7 +125,7 @@ describe("prefetchData", () => {
       },
     });
 
-    await client.waitForInitialization();
+    await client.connect();
 
     expect(client.extract()).toMatchInlineSnapshot(`
       {
@@ -202,7 +203,7 @@ describe("prefetchData", () => {
       },
     });
 
-    await client.waitForInitialization();
+    await client.connect();
 
     expect(client.extract()).toMatchInlineSnapshot(`
       {
@@ -307,7 +308,7 @@ describe("prefetchData", () => {
       },
     });
 
-    await client.waitForInitialization();
+    await client.connect();
 
     expect(client.extract()).toMatchInlineSnapshot(`
       {
@@ -375,7 +376,7 @@ describe("prefetchData", () => {
       },
     });
 
-    await client.waitForInitialization();
+    await client.connect();
 
     expect(client.extract()).toMatchInlineSnapshot(`
       {
@@ -440,7 +441,7 @@ describe("custom links", () => {
       },
     }));
 
-    await client.waitForInitialization();
+    await client.connect();
 
     const variables = { id: "1" };
     const query = parse(manifest.operations[0].body);
@@ -535,3 +536,108 @@ describe("custom links", () => {
     }).not.toThrow();
   });
 });
+
+describe("watchQuery dev warnings", () => {
+  const query = gql`
+    query Products($category: String!, $page: Int!, $sortBy: String!)
+    @tool(name: "GetProductsByCategory") {
+      products(category: $category, page: $page, sortBy: $sortBy) {
+        id
+      }
+    }
+  `;
+
+  async function setupClient({
+    toolInput,
+  }: {
+    toolInput: Record<string, unknown>;
+  }) {
+    stubOpenAiGlobals({ toolInput });
+    const client = new ApolloClient({
+      cache: new InMemoryCache(),
+      manifest: mockApplicationManifest(),
+    });
+    using host = await mockMcpHost({
+      hostContext: minimalHostContextWithToolName("GetProductsByCategory"),
+    });
+    host.onCleanup(() => client.stop());
+    host.sendToolResult(graphqlToolResult({ data: { products: [] } }));
+    host.sendToolInput({
+      arguments: toolInput,
+    });
+    await client.connect();
+    return client;
+  }
+
+  test("warns when variables don't match tool input", async () => {
+    using _ = spyOnConsole("debug", "warn");
+    const client = await setupClient({
+      toolInput: { category: "electronics", page: 1, sortBy: "title" },
+    });
+
+    client.watchQuery({
+      query,
+      variables: { category: "music", page: 1, sortBy: "name" },
+    });
+
+    expect(console.warn).toHaveBeenCalledWith(
+      expect.stringContaining("useHydratedVariables"),
+      { category: "electronics", page: 1, sortBy: "title" },
+      { category: "music", page: 1, sortBy: "name" }
+    );
+  });
+
+  test("does not warn when variables match tool input", async () => {
+    using _ = spyOnConsole("debug", "warn");
+    const client = await setupClient({
+      toolInput: { category: "electronics", page: 1, sortBy: "title" },
+    });
+
+    client.watchQuery({
+      query,
+      variables: { category: "electronics", page: 1, sortBy: "title" },
+    });
+
+    expect(console.warn).not.toHaveBeenCalled();
+  });
+
+  test("does not warn when query has no matching @tool directive", async () => {
+    using _ = spyOnConsole("debug", "warn");
+    const client = await setupClient({
+      toolInput: { category: "electronics", page: 1, sortBy: "title" },
+    });
+
+    const queryWithoutTool = gql`
+      query Products($category: String!) {
+        products(category: $category) {
+          id
+        }
+      }
+    `;
+
+    client.watchQuery({
+      query: queryWithoutTool,
+      variables: { category: "music" },
+    });
+
+    expect(console.warn).not.toHaveBeenCalled();
+  });
+
+  test("warning fires at most once (subsequent calls don't re-warn)", async () => {
+    using _ = spyOnConsole("debug", "warn");
+    const client = await setupClient({
+      toolInput: { category: "electronics", page: 1, sortBy: "title" },
+    });
+
+    client.watchQuery({
+      query,
+      variables: { category: "music", page: 1, sortBy: "name" },
+    });
+    client.watchQuery({
+      query,
+      variables: { category: "music", page: 1, sortBy: "name" },
+    });
+
+    expect(console.warn).toHaveBeenCalledTimes(1);
+  });
+});