@apollo/client-ai-apps 0.5.2 → 0.5.3

package/src/mcp/core/ApolloClient.ts

@@ -3,15 +3,24 @@ import {
   ApolloClient as BaseApolloClient,
   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,
+  getVariableNamesFromDocument,
   getVariablesForOperationFromToolInput,
+  warnOnVariableMismatch,
 } from "../../utilities/index.js";
 import { McpAppManager } from "./McpAppManager.js";
 
@@ -29,6 +38,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();
 
@@ -57,9 +68,63 @@ 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/mcp/core/McpAppManager.ts

@@ -35,7 +35,7 @@ export class McpAppManager {
     return this.#toolInput;
   }
 
-  waitForInitialization = cacheAsync(async () => {
+  connect = cacheAsync(async () => {
     let toolResult = promiseWithResolvers<ApolloMcpServerApps.CallToolResult>();
     let toolInput = promiseWithResolvers<Parameters<App["ontoolinput"]>[0]>();
 
@@ -49,7 +49,7 @@ export class McpAppManager {
       toolInput.resolve(params);
     };
 
-    await this.connect();
+    await this.connectToHost();
 
     const { structuredContent, _meta } = await toolResult.promise;
     const { arguments: args } = await toolInput.promise;
@@ -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/mcp/core/__tests__/ApolloClient.test.ts

@@ -1,9 +1,10 @@
-import { expect, test, vi } from "vitest";
+import { expect, test, vi, describe } from "vitest";
 import { ApolloClient } from "../ApolloClient.js";
 import { ApolloLink, HttpLink, InMemoryCache, gql } from "@apollo/client";
 import { print } from "@apollo/client/utilities";
 import { ToolCallLink } from "../../link/ToolCallLink.js";
 import {
+  graphqlToolResult,
   minimalHostContextWithToolName,
   mockApplicationManifest,
   mockMcpHost,
@@ -57,7 +58,7 @@ test("writes tool result data to cache", async () => {
   });
   host.sendToolInput({ arguments: { id: "1" } });
 
-  await client.waitForInitialization();
+  await client.connect();
 
   expect(client.extract()).toEqual({
     "Product:1": {
@@ -124,7 +125,7 @@ test("writes prefetch data to cache", async () => {
   });
   host.sendToolInput({ arguments: {} });
 
-  await client.waitForInitialization();
+  await client.connect();
 
   expect(client.extract()).toEqual({
     "Product:1": {
@@ -213,7 +214,7 @@ test("writes prefetch and tool response data to cache when both are provided", a
   });
   host.sendToolInput({ arguments: { id: "2" } });
 
-  await client.waitForInitialization();
+  await client.connect();
 
   expect(client.extract()).toEqual({
     "Product:1": {
@@ -281,7 +282,7 @@ test("excludes extra tool input variables not defined in the operation", async (
   });
   host.sendToolInput({ arguments: { id: "1", extraParam: "ignored" } });
 
-  await client.waitForInitialization();
+  await client.connect();
 
   expect(client.extract()).toEqual({
     "Product:1": {
@@ -345,7 +346,7 @@ test("allows for custom links provided to the constructor", async () => {
     },
   }));
 
-  await client.waitForInitialization();
+  await client.connect();
 
   const variables = { id: "1" };
   const query = gql(manifest.operations[0].body);
@@ -462,3 +463,105 @@ test("creates a default ToolCallLink when no link is provided", () => {
     });
   }).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>;
+  }) {
+    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);
+  });
+});

package/src/mcp/link/__tests__/ToolCallLink.test.ts

@@ -49,7 +49,7 @@ test("delegates query execution to MCP host", async () => {
     },
   }));
 
-  await client.waitForInitialization();
+  await client.connect();
 
   const observable = execute(new ToolCallLink(), { query }, { client });
   const stream = new ObservableStream(observable);