@apollo/client-ai-apps 0.4.0 → 0.5.0

package/src/mcp/core/McpAppManager.ts

@@ -0,0 +1,98 @@
+import { App, PostMessageTransport } from "@modelcontextprotocol/ext-apps";
+import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import type { ApplicationManifest } from "../../types/application-manifest";
+import type { FormattedExecutionResult } from "graphql";
+import type { DocumentNode, OperationVariables } from "@apollo/client";
+import { print } from "@apollo/client/utilities";
+import { cacheAsync, promiseWithResolvers } from "../../utilities";
+import type { ApolloMcpServerApps } from "../../core/types";
+
+type ExecuteQueryCallToolResult = Omit<CallToolResult, "structuredContent"> & {
+  structuredContent: FormattedExecutionResult;
+};
+
+/** @internal */
+export class McpAppManager {
+  readonly app: App;
+
+  #toolName: string | undefined;
+  #toolMetadata: ApolloMcpServerApps.CallToolResult["_meta"];
+  #toolInput: Record<string, unknown> | undefined;
+
+  constructor(manifest: ApplicationManifest) {
+    this.app = new App({ name: manifest.name, version: manifest.appVersion });
+  }
+
+  get toolName() {
+    return this.#toolName;
+  }
+
+  get toolMetadata() {
+    return this.#toolMetadata;
+  }
+
+  get toolInput() {
+    return this.#toolInput;
+  }
+
+  waitForInitialization = cacheAsync(async () => {
+    let toolResult = promiseWithResolvers<ApolloMcpServerApps.CallToolResult>();
+    let toolInput = promiseWithResolvers<Parameters<App["ontoolinput"]>[0]>();
+
+    this.app.ontoolresult = (params) => {
+      toolResult.resolve(
+        params as unknown as ApolloMcpServerApps.CallToolResult
+      );
+    };
+
+    this.app.ontoolinput = (params) => {
+      toolInput.resolve(params);
+    };
+
+    await this.connect();
+
+    const { structuredContent, _meta } = await toolResult.promise;
+    const { arguments: args } = await toolInput.promise;
+
+    this.#toolName = this.app.getHostContext()?.toolInfo?.tool.name;
+    this.#toolMetadata = _meta;
+    this.#toolInput = args;
+
+    return {
+      ...structuredContent,
+      toolName: this.toolName,
+      args,
+    };
+  });
+
+  close() {
+    return this.app.close();
+  }
+
+  async executeQuery({
+    query,
+    variables,
+  }: {
+    query: DocumentNode;
+    variables: OperationVariables | undefined;
+  }) {
+    const result = (await this.app.callServerTool({
+      name: "execute",
+      arguments: { query: print(query), variables },
+    })) as ExecuteQueryCallToolResult;
+
+    return result.structuredContent;
+  }
+
+  private async connect() {
+    try {
+      return await this.app.connect(
+        new PostMessageTransport(window.parent, window.parent)
+      );
+    } catch (e) {
+      const error = e instanceof Error ? e : new Error("Failed to connect");
+
+      throw error;
+    }
+  }
+}

package/src/mcp/core/Platform.ts

@@ -0,0 +1,11 @@
+export const Platform = Object.freeze({
+  target: "mcp",
+  select<
+    T = unknown,
+    TReturn = T extends (...args: any[]) => infer TReturn ? TReturn : T,
+  >(config: { mcp?: T; openai?: T }): TReturn | undefined {
+    const value = config.mcp;
+
+    return typeof value === "function" ? value() : (value as TReturn);
+  },
+});

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

@@ -0,0 +1,464 @@
+import { expect, test, vi } 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 {
+  minimalHostContextWithToolName,
+  mockApplicationManifest,
+  mockMcpHost,
+  spyOnConsole,
+} from "../../../testing/internal/index.js";
+
+test("writes tool result data to cache", async () => {
+  using _ = spyOnConsole("debug");
+
+  const query = gql`
+    query Product($id: ID!) {
+      product(id: $id) {
+        id
+        title
+        __typename
+      }
+    }
+  `;
+
+  const client = new ApolloClient({
+    cache: new InMemoryCache(),
+    manifest: mockApplicationManifest({
+      operations: [
+        {
+          id: "1",
+          name: "Product",
+          body: print(query),
+          type: "query",
+          prefetch: false,
+          variables: { id: "ID" },
+          tools: [{ name: "GetProduct", description: "Get a product" }],
+        },
+      ],
+    }),
+  });
+
+  using host = await mockMcpHost({
+    hostContext: minimalHostContextWithToolName("GetProduct"),
+  });
+  host.onCleanup(() => client.stop());
+
+  host.sendToolResult({
+    content: [],
+    structuredContent: {
+      result: {
+        data: {
+          product: { id: "1", title: "Pen", __typename: "Product" },
+        },
+      },
+    },
+  });
+  host.sendToolInput({ arguments: { id: "1" } });
+
+  await client.waitForInitialization();
+
+  expect(client.extract()).toEqual({
+    "Product:1": {
+      __typename: "Product",
+      id: "1",
+      title: "Pen",
+    },
+    ROOT_QUERY: {
+      __typename: "Query",
+      'product({"id":"1"})': {
+        __ref: "Product:1",
+      },
+    },
+  });
+});
+
+test("writes prefetch data to cache", async () => {
+  using _ = spyOnConsole("debug");
+
+  const query = gql`
+    query TopProducts {
+      topProducts {
+        id
+        title
+        __typename
+      }
+    }
+  `;
+
+  const client = new ApolloClient({
+    cache: new InMemoryCache(),
+    manifest: mockApplicationManifest({
+      operations: [
+        {
+          id: "1",
+          name: "TopProducts",
+          body: print(query),
+          type: "query",
+          prefetch: true,
+          prefetchID: "__anonymous",
+          variables: {},
+          tools: [{ name: "TopProducts", description: "Shows top products" }],
+        },
+      ],
+    }),
+  });
+
+  using host = await mockMcpHost({
+    hostContext: minimalHostContextWithToolName("OtherTool"),
+  });
+  host.onCleanup(() => client.stop());
+
+  host.sendToolResult({
+    content: [],
+    structuredContent: {
+      prefetch: {
+        __anonymous: {
+          data: {
+            topProducts: [{ id: "1", title: "iPhone", __typename: "Product" }],
+          },
+        },
+      },
+    },
+  });
+  host.sendToolInput({ arguments: {} });
+
+  await client.waitForInitialization();
+
+  expect(client.extract()).toEqual({
+    "Product:1": {
+      __typename: "Product",
+      id: "1",
+      title: "iPhone",
+    },
+    ROOT_QUERY: {
+      __typename: "Query",
+      topProducts: [{ __ref: "Product:1" }],
+    },
+  });
+});
+
+test("writes prefetch and tool response data to cache when both are provided", async () => {
+  using _ = spyOnConsole("debug");
+
+  const prefetchQuery = gql`
+    query TopProducts {
+      topProducts {
+        id
+        title
+        __typename
+      }
+    }
+  `;
+
+  const query = gql`
+    query Product($id: ID!) {
+      product(id: $id) {
+        id
+        title
+        __typename
+      }
+    }
+  `;
+
+  const client = new ApolloClient({
+    cache: new InMemoryCache(),
+    manifest: mockApplicationManifest({
+      operations: [
+        {
+          id: "1",
+          name: "TopProducts",
+          body: print(prefetchQuery),
+          type: "query",
+          prefetch: true,
+          prefetchID: "__anonymous",
+          variables: {},
+          tools: [{ name: "TopProducts", description: "Shows top products" }],
+        },
+        {
+          id: "2",
+          name: "Product",
+          body: print(query),
+          type: "query",
+          prefetch: false,
+          variables: { id: "ID" },
+          tools: [{ name: "Product", description: "Get a product by id" }],
+        },
+      ],
+    }),
+  });
+
+  using host = await mockMcpHost({
+    hostContext: minimalHostContextWithToolName("Product"),
+  });
+  host.onCleanup(() => client.stop());
+
+  host.sendToolResult({
+    content: [],
+    structuredContent: {
+      prefetch: {
+        __anonymous: {
+          data: {
+            topProducts: [{ id: "1", title: "iPhone", __typename: "Product" }],
+          },
+        },
+      },
+      result: {
+        data: {
+          product: { __typename: "Product", id: "2", title: "iPad" },
+        },
+      },
+    },
+  });
+  host.sendToolInput({ arguments: { id: "2" } });
+
+  await client.waitForInitialization();
+
+  expect(client.extract()).toEqual({
+    "Product:1": {
+      __typename: "Product",
+      id: "1",
+      title: "iPhone",
+    },
+    "Product:2": {
+      __typename: "Product",
+      id: "2",
+      title: "iPad",
+    },
+    ROOT_QUERY: {
+      __typename: "Query",
+      topProducts: [{ __ref: "Product:1" }],
+      'product({"id":"2"})': { __ref: "Product:2" },
+    },
+  });
+});
+
+test("excludes extra tool input variables not defined in the operation", async () => {
+  using _ = spyOnConsole("debug");
+
+  const query = gql`
+    query Product($id: ID!) {
+      product(id: $id) {
+        id
+        title
+        __typename
+      }
+    }
+  `;
+
+  const client = new ApolloClient({
+    cache: new InMemoryCache(),
+    manifest: mockApplicationManifest({
+      operations: [
+        {
+          id: "1",
+          name: "Product",
+          body: print(query),
+          type: "query",
+          prefetch: false,
+          variables: { id: "ID" },
+          tools: [{ name: "GetProduct", description: "Get a product" }],
+        },
+      ],
+    }),
+  });
+
+  using host = await mockMcpHost({
+    hostContext: minimalHostContextWithToolName("GetProduct"),
+  });
+  host.onCleanup(() => client.stop());
+
+  host.sendToolResult({
+    content: [],
+    structuredContent: {
+      result: {
+        data: {
+          product: { id: "1", title: "Pen", __typename: "Product" },
+        },
+      },
+    },
+  });
+  host.sendToolInput({ arguments: { id: "1", extraParam: "ignored" } });
+
+  await client.waitForInitialization();
+
+  expect(client.extract()).toEqual({
+    "Product:1": {
+      __typename: "Product",
+      id: "1",
+      title: "Pen",
+    },
+    ROOT_QUERY: {
+      __typename: "Query",
+      'product({"id":"1"})': {
+        __ref: "Product:1",
+      },
+    },
+  });
+});
+
+test("allows for custom links provided to the constructor", async () => {
+  using _ = spyOnConsole("debug");
+  const manifest = mockApplicationManifest();
+  const linkHandler = vi.fn<ApolloLink.RequestHandler>((operation, forward) =>
+    forward(operation)
+  );
+
+  const client = new ApolloClient({
+    manifest,
+    cache: new InMemoryCache(),
+    link: ApolloLink.from([new ApolloLink(linkHandler), new ToolCallLink()]),
+  });
+
+  using host = await mockMcpHost({
+    hostContext: minimalHostContextWithToolName("GetProduct"),
+  });
+  host.onCleanup(() => client.stop());
+
+  host.sendToolResult({
+    content: [],
+    structuredContent: {
+      result: {
+        data: {
+          product: { id: "1", title: "Pen", __typename: "Product" },
+        },
+      },
+    },
+  });
+  host.sendToolInput({ arguments: {} });
+
+  host.mockToolCall("execute", () => ({
+    content: [],
+    structuredContent: {
+      data: {
+        product: {
+          id: "1",
+          title: "Pen",
+          rating: 5,
+          price: 1.0,
+          description: "Awesome pen",
+          images: [],
+          __typename: "Product",
+        },
+      },
+    },
+  }));
+
+  await client.waitForInitialization();
+
+  const variables = { id: "1" };
+  const query = gql(manifest.operations[0].body);
+
+  await expect(client.query({ query, variables })).resolves.toStrictEqual({
+    data: {
+      product: {
+        id: "1",
+        title: "Pen",
+        rating: 5,
+        price: 1.0,
+        description: "Awesome pen",
+        images: [],
+        __typename: "Product",
+      },
+    },
+  });
+
+  expect(linkHandler).toHaveBeenCalledOnce();
+  expect(linkHandler).toHaveBeenCalledWith(
+    expect.objectContaining({ query, variables, operationType: "query" }),
+    expect.any(Function)
+  );
+});
+
+test("enforces ToolCallLink as terminating link", () => {
+  const manifest = mockApplicationManifest();
+  const expectedError = 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."
+  );
+
+  expect(() => {
+    new ApolloClient({
+      manifest,
+      cache: new InMemoryCache(),
+      link: new HttpLink(),
+    });
+  }).toThrow(expectedError);
+
+  expect(() => {
+    new ApolloClient({
+      manifest,
+      cache: new InMemoryCache(),
+      link: new ApolloLink(),
+    });
+  }).toThrow(expectedError);
+
+  expect(() => {
+    new ApolloClient({
+      manifest,
+      cache: new InMemoryCache(),
+      link: ApolloLink.from([new ApolloLink(), new HttpLink()]),
+    });
+  }).toThrow(expectedError);
+
+  expect(() => {
+    new ApolloClient({
+      manifest,
+      cache: new InMemoryCache(),
+      link: ApolloLink.split(() => true, new ToolCallLink(), new HttpLink()),
+    });
+  }).toThrow(expectedError);
+
+  // Allow you to use a custom terminating link for `split` links if the
+  // custom link is the `left` branch link.
+  expect(() => {
+    new ApolloClient({
+      manifest,
+      cache: new InMemoryCache(),
+      link: ApolloLink.split(() => true, new HttpLink(), new ToolCallLink()),
+    });
+  }).not.toThrow(expectedError);
+
+  expect(() => {
+    new ApolloClient({
+      manifest,
+      cache: new InMemoryCache(),
+      link: ApolloLink.split(
+        () => true,
+        new ToolCallLink(),
+        new ToolCallLink()
+      ),
+    });
+  }).not.toThrow();
+});
+
+test("allows ToolCallLink as terminating link", () => {
+  const manifest = mockApplicationManifest();
+
+  expect(() => {
+    new ApolloClient({
+      manifest,
+      cache: new InMemoryCache(),
+      link: ApolloLink.from([new ApolloLink(), new ToolCallLink()]),
+    });
+  }).not.toThrow();
+
+  expect(() => {
+    new ApolloClient({
+      manifest,
+      cache: new InMemoryCache(),
+      link: ApolloLink.split(() => true, new HttpLink(), new ToolCallLink()),
+    });
+  }).not.toThrow();
+});
+
+test("creates a default ToolCallLink when no link is provided", () => {
+  const manifest = mockApplicationManifest();
+
+  expect(() => {
+    new ApolloClient({
+      manifest,
+      cache: new InMemoryCache(),
+    });
+  }).not.toThrow();
+});

package/src/mcp/index.ts

@@ -0,0 +1,2 @@
+export { ApolloClient } from "./core/ApolloClient.js";
+export { ToolCallLink } from "./link/ToolCallLink.js";

package/src/mcp/link/ToolCallLink.ts

@@ -0,0 +1,40 @@
+import { ApolloLink, Observable } from "@apollo/client";
+import { from } from "rxjs";
+import type { ApolloClient as McpApolloClient } from "../core/ApolloClient";
+
+/**
+ * A terminating link that sends a GraphQL request through an agent tool call.
+ * When providing a custom link chain to `ApolloClient`, `ApolloClient` will
+ * validate that the terminating link is an instance of this link.
+ *
+ * @example Providing a custom link chain
+ *
+ * ```ts
+ * import { ApolloLink } from "@apollo/client";
+ * import { ApolloClient, ToolCallLink } from "@apollo/client-ai-apps";
+ *
+ * const link = ApolloLink.from([
+ *   ...otherLinks,
+ *   new ToolCallLink()
+ * ]);
+ *
+ * const client = new ApolloClient({
+ *   link,
+ *   // ...
+ * });
+ * ```
+ */
+export class ToolCallLink extends ApolloLink {
+  readonly name = "ToolCallLink";
+
+  request(operation: ApolloLink.Operation): Observable<ApolloLink.Result> {
+    const client = operation.client as McpApolloClient;
+
+    return from(
+      client["appManager"].executeQuery({
+        query: operation.query,
+        variables: operation.variables,
+      })
+    );
+  }
+}

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

@@ -0,0 +1,53 @@
+import { expect, test } from "vitest";
+import { gql, InMemoryCache } from "@apollo/client";
+import { execute } from "@apollo/client/link";
+import { ApolloClient } from "../../core/ApolloClient.js";
+import {
+  mockApplicationManifest,
+  mockMcpHost,
+  ObservableStream,
+  spyOnConsole,
+} from "../../../testing/internal/index.js";
+import { ToolCallLink } from "../ToolCallLink.js";
+
+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 host = await mockMcpHost();
+  host.onCleanup(() => client.stop());
+
+  host.sendToolResult({
+    _meta: { toolName: "GetProduct" },
+    content: [],
+    structuredContent: {},
+  });
+  host.sendToolInput({ arguments: {} });
+
+  host.mockToolCall("execute", () => ({
+    content: [],
+    structuredContent: {
+      data: { greeting: "Hello, world" },
+    },
+  }));
+
+  await client.waitForInitialization();
+
+  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/react/hooks/__tests__/useApp.test.tsx

@@ -0,0 +1,46 @@
+import { test, expect } from "vitest";
+import {
+  disableActEnvironment,
+  renderHookToSnapshotStream,
+} from "@testing-library/react-render-stream";
+import { Suspense } from "react";
+import { InMemoryCache } from "@apollo/client";
+import { App } from "@modelcontextprotocol/ext-apps";
+
+import { useApp } from "../useApp.js";
+import { ApolloClient } from "../../../core/ApolloClient.js";
+import {
+  mockApplicationManifest,
+  mockMcpHost,
+  spyOnConsole,
+} from "../../../../testing/internal/index.js";
+import { ApolloProvider } from "../../../../react/ApolloProvider.js";
+
+test("returns app instance created by ApolloClient", async () => {
+  using _ = spyOnConsole("debug");
+  const client = new ApolloClient({
+    cache: new InMemoryCache(),
+    manifest: mockApplicationManifest(),
+  });
+
+  using host = await mockMcpHost();
+  host.onCleanup(() => client.stop());
+
+  host.sendToolInput({ arguments: {} });
+  host.sendToolResult({
+    _meta: { toolName: "Test" },
+    content: [],
+    structuredContent: {},
+  });
+
+  using _disabledAct = disableActEnvironment();
+  const { takeSnapshot } = await renderHookToSnapshotStream(() => useApp(), {
+    wrapper: ({ children }) => (
+      <Suspense>
+        <ApolloProvider client={client}>{children}</ApolloProvider>
+      </Suspense>
+    ),
+  });
+
+  await expect(takeSnapshot()).resolves.toBeInstanceOf(App);
+});