@apollo/client-ai-apps 0.6.5 → 0.7.0

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

@@ -0,0 +1,99 @@
+import { expect, test } from "vitest";
+import { gql, InMemoryCache } from "@apollo/client";
+import { execute } from "@apollo/client/link";
+import {
+  eachHostEnv,
+  mockApplicationManifest,
+  ObservableStream,
+  spyOnConsole,
+} from "../../testing/internal/index.js";
+import { ToolCallLink } from "../ToolCallLink.js";
+
+eachHostEnv((setupHost, ApolloClient) => {
+  test("merges _meta.structuredContent into result for @private fields", async () => {
+    using _ = spyOnConsole("debug");
+    const query = gql`
+      query GreetingQuery {
+        greeting @private
+      }
+    `;
+
+    const client = new ApolloClient({
+      cache: new InMemoryCache(),
+      manifest: mockApplicationManifest(),
+    });
+
+    using env = await setupHost({
+      client,
+      toolCall: {
+        name: "GetProduct",
+        result: { structuredContent: { result: { data: { product: null } } } },
+      },
+    });
+    const { host, params } = env;
+
+    host.sendToolInput(params.toolInput);
+    host.sendToolResult(params.toolResult);
+
+    host.mockToolCall("execute", () => ({
+      structuredContent: {},
+      _meta: {
+        structuredContent: { data: { greeting: "Hello, private world" } },
+      },
+    }));
+
+    await client.connect();
+
+    const observable = execute(new ToolCallLink(), { query }, { client });
+    const stream = new ObservableStream(observable);
+
+    await expect(stream).toEmitValue({
+      data: { greeting: "Hello, private world" },
+    });
+
+    await expect(stream).toComplete();
+  });
+
+  test("delegates query execution to MCP host", async () => {
+    using _ = spyOnConsole("debug");
+    const query = gql`
+      query GreetingQuery {
+        greeting
+      }
+    `;
+
+    const client = new ApolloClient({
+      cache: new InMemoryCache(),
+      manifest: mockApplicationManifest(),
+    });
+
+    using env = await setupHost({
+      client,
+      toolCall: {
+        name: "GetProduct",
+        result: { structuredContent: { result: { data: { product: null } } } },
+      },
+    });
+    const { host, params } = env;
+
+    host.sendToolInput(params.toolInput);
+    host.sendToolResult(params.toolResult);
+
+    host.mockToolCall("execute", () => ({
+      structuredContent: {
+        data: { greeting: "Hello, world" },
+      },
+    }));
+
+    await client.connect();
+
+    const observable = execute(new ToolCallLink(), { query }, { client });
+    const stream = new ObservableStream(observable);
+
+    await expect(stream).toEmitValue({
+      data: { greeting: "Hello, world" },
+    });
+
+    await expect(stream).toComplete();
+  });
+});

package/src/mcp/core/ApolloClient.ts

@@ -1,178 +1,40 @@
-import {
-  ApolloLink,
-  ApolloClient as BaseApolloClient,
-  DocumentTransform,
-} from "@apollo/client";
-import type {
-  WatchQueryOptions,
-  ObservableQuery,
-  OperationVariables,
-} from "@apollo/client";
-import { removeDirectivesFromDocument } from "@apollo/client/utilities/internal";
-import { parse, visit } 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";
+import type { App } from "@modelcontextprotocol/ext-apps";
+import { AbstractApolloClient } from "../../core/AbstractApolloClient.js";
+import { connectToHost, promiseWithResolvers } from "../../utilities/index.js";
+import type { ApolloMcpServerApps } from "../../core/types.js";
 
 export declare namespace ApolloClient {
-  export interface Options extends Omit<BaseApolloClient.Options, "link"> {
-    link?: BaseApolloClient.Options["link"];
-    manifest: ApplicationManifest;
-  }
+  export interface Options extends AbstractApolloClient.Options {}
 }
 
-export class ApolloClient extends BaseApolloClient {
-  manifest: ApplicationManifest;
-  private readonly appManager: McpAppManager;
-
-  /** @internal */
-  readonly [aiClientSymbol] = true;
-
-  #toolInput: Record<string, unknown> | undefined;
-
+export class ApolloClient extends AbstractApolloClient {
   constructor(options: ApolloClient.Options) {
-    const link = options.link ?? new ToolCallLink();
-
-    if (__DEV__) {
-      validateTerminatingLink(link);
-    }
-
-    super({
-      ...options,
-      link,
-      // Strip out the prefetch/tool directives so they don't get sent with the operation to the server
-      documentTransform: new DocumentTransform((document) => {
-        const serverDocument = removeDirectivesFromDocument(
-          [{ name: "prefetch" }, { name: "tool" }],
-          document
-        )!;
-
-        return visit(serverDocument, {
-          OperationDefinition(node) {
-            return { ...node, description: undefined };
-          },
-        });
-      }).concat(options.documentTransform ?? DocumentTransform.identity()),
-    });
-
-    this.manifest = options.manifest;
-    this.appManager = new McpAppManager(this.manifest);
-  }
-
-  stop() {
-    super.stop();
-    this.appManager.close().catch(() => {});
-  }
-
-  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 { structuredContent, toolName, args } =
-      await this.appManager.connect();
-
-    this.#toolInput = args;
-
-    this.manifest.operations.forEach((operation) => {
-      if (
-        operation.prefetchID &&
-        structuredContent.prefetch?.[operation.prefetchID]
-      ) {
-        this.writeQuery({
-          query: parse(operation.body),
-          data: structuredContent.prefetch[operation.prefetchID].data,
-        });
-      }
-
-      if (operation.tools.find((tool) => tool.name === toolName)) {
-        if (structuredContent.result?.data) {
-          this.writeQuery({
-            query: parse(operation.body),
-            data: structuredContent.result.data,
-            variables: getVariablesForOperationFromToolInput(operation, args),
-          });
-        }
-      }
+    super(options, async (app) => {
+      const toolResult =
+        promiseWithResolvers<ApolloMcpServerApps.CallToolResult>();
+      const toolInput =
+        promiseWithResolvers<Parameters<App["ontoolinput"]>[0]>();
+
+      app.ontoolresult = (params) => {
+        toolResult.resolve(
+          params as unknown as ApolloMcpServerApps.CallToolResult
+        );
+      };
+
+      app.ontoolinput = (params) => {
+        toolInput.resolve(params);
+      };
+
+      await connectToHost(app);
+
+      const { structuredContent, _meta } = await toolResult.promise;
+      const { arguments: args } = await toolInput.promise;
+
+      return {
+        structuredContent,
+        toolInput: args,
+        _meta,
+      };
     });
-  });
-}
-
-function validateTerminatingLink(link: ApolloLink) {
-  let terminatingLink = link;
-
-  while (terminatingLink.right) {
-    terminatingLink = terminatingLink.right;
   }
-
-  if (
-    !isNamedLink(terminatingLink) ||
-    terminatingLink.name !== "ToolCallLink"
-  ) {
-    throw new Error(
-      "The terminating link must be a `ToolCallLink`. If you are using a `split` link, ensure the `right` branch uses a `ToolCallLink` as the terminating link."
-    );
-  }
-}
-
-function isNamedLink(link: ApolloLink): link is ApolloLink & { name: string } {
-  return "name" in link;
 }