@apollo/client-ai-apps 0.4.0 → 0.5.1

package/src/mcp/react/hooks/__tests__/useToolInput.test.tsx

@@ -0,0 +1,50 @@
+import { test, expect } from "vitest";
+import {
+  disableActEnvironment,
+  renderHookToSnapshotStream,
+} from "@testing-library/react-render-stream";
+import { Suspense } from "react";
+import { InMemoryCache } from "@apollo/client";
+
+import { useToolInput } from "../useToolInput.js";
+import { ApolloClient } from "../../../core/ApolloClient.js";
+import {
+  graphqlToolResult,
+  minimalHostContextWithToolName,
+  mockApplicationManifest,
+  mockMcpHost,
+  spyOnConsole,
+} from "../../../../testing/internal/index.js";
+import { ApolloProvider } from "../../../../react/ApolloProvider.js";
+
+test("returns the tool input from the MCP host", async () => {
+  using _ = spyOnConsole("debug");
+  const client = new ApolloClient({
+    cache: new InMemoryCache(),
+    manifest: mockApplicationManifest(),
+  });
+
+  using host = await mockMcpHost({
+    hostContext: minimalHostContextWithToolName("GetProduct"),
+  });
+  host.onCleanup(() => client.stop());
+
+  host.sendToolInput({ arguments: { id: "1" } });
+  host.sendToolResult(graphqlToolResult({ data: { product: null } }));
+
+  using _disabledAct = disableActEnvironment();
+  const { takeSnapshot } = await renderHookToSnapshotStream(
+    () => useToolInput(),
+    {
+      wrapper: ({ children }) => (
+        <Suspense>
+          <ApolloProvider client={client}>{children}</ApolloProvider>
+        </Suspense>
+      ),
+    }
+  );
+
+  await expect(takeSnapshot()).resolves.toEqual({ id: "1" });
+
+  await expect(takeSnapshot).not.toRerender();
+});

package/src/mcp/react/hooks/__tests__/useToolMetadata.test.tsx

@@ -0,0 +1,53 @@
+import { test, expect } from "vitest";
+import {
+  disableActEnvironment,
+  renderHookToSnapshotStream,
+} from "@testing-library/react-render-stream";
+import { Suspense } from "react";
+import { InMemoryCache } from "@apollo/client";
+
+import { useToolMetadata } from "../useToolMetadata.js";
+import { ApolloClient } from "../../../core/ApolloClient.js";
+import {
+  mockApplicationManifest,
+  mockMcpHost,
+  spyOnConsole,
+} from "../../../../testing/internal/index.js";
+import { ApolloProvider } from "../../../../react/ApolloProvider.js";
+
+test("returns the tool metadata from the MCP host", 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: "TestTool", customField: "customValue" },
+    content: [],
+    structuredContent: {},
+  });
+
+  using _disabledAct = disableActEnvironment();
+  const { takeSnapshot } = await renderHookToSnapshotStream(
+    () => useToolMetadata(),
+    {
+      wrapper: ({ children }) => (
+        <Suspense>
+          <ApolloProvider client={client}>{children}</ApolloProvider>
+        </Suspense>
+      ),
+    }
+  );
+
+  await expect(takeSnapshot()).resolves.toEqual({
+    toolName: "TestTool",
+    customField: "customValue",
+  });
+
+  await expect(takeSnapshot).not.toRerender();
+});

package/src/mcp/react/hooks/__tests__/useToolName.test.tsx

@@ -0,0 +1,50 @@
+import { test, expect } from "vitest";
+import {
+  disableActEnvironment,
+  renderHookToSnapshotStream,
+} from "@testing-library/react-render-stream";
+import { Suspense } from "react";
+import { InMemoryCache } from "@apollo/client";
+
+import { useToolName } from "../useToolName.js";
+import { ApolloClient } from "../../../core/ApolloClient.js";
+import {
+  graphqlToolResult,
+  minimalHostContextWithToolName,
+  mockApplicationManifest,
+  mockMcpHost,
+  spyOnConsole,
+} from "../../../../testing/internal/index.js";
+import { ApolloProvider } from "../../../../react/ApolloProvider.js";
+
+test("returns the tool name from the MCP host", async () => {
+  using _ = spyOnConsole("debug");
+  const client = new ApolloClient({
+    cache: new InMemoryCache(),
+    manifest: mockApplicationManifest(),
+  });
+
+  using host = await mockMcpHost({
+    hostContext: minimalHostContextWithToolName("GetProduct"),
+  });
+  host.onCleanup(() => client.stop());
+
+  host.sendToolInput({ arguments: {} });
+  host.sendToolResult(graphqlToolResult({ data: { product: null } }));
+
+  using _disabledAct = disableActEnvironment();
+  const { takeSnapshot } = await renderHookToSnapshotStream(
+    () => useToolName(),
+    {
+      wrapper: ({ children }) => (
+        <Suspense>
+          <ApolloProvider client={client}>{children}</ApolloProvider>
+        </Suspense>
+      ),
+    }
+  );
+
+  await expect(takeSnapshot()).resolves.toBe("GetProduct");
+
+  await expect(takeSnapshot).not.toRerender();
+});

package/src/mcp/react/hooks/useApolloClient.ts

@@ -0,0 +1,14 @@
+import { useApolloClient as useBaseApolloClient } from "@apollo/client/react";
+import { ApolloClient } from "../../core/ApolloClient.js";
+import { aiClientSymbol, invariant } from "../../../utilities/index.js";
+
+export function useApolloClient(override?: ApolloClient) {
+  const client = useBaseApolloClient(override) as ApolloClient;
+
+  invariant(
+    client[aiClientSymbol],
+    'The "client" instance provided to <ApolloProvider /> is the wrong instance. You might have imported `ApolloClient` from `@apollo/client`. Please import `ApolloClient` from `@apollo/client-ai-apps` instead.'
+  );
+
+  return client;
+}

package/src/mcp/react/hooks/useApp.ts

@@ -0,0 +1,5 @@
+import { useApolloClient } from "./useApolloClient.js";
+
+export function useApp() {
+  return useApolloClient()["appManager"].app;
+}

package/src/mcp/react/hooks/useToolInput.ts

@@ -0,0 +1,5 @@
+import { useApolloClient } from "./useApolloClient.js";
+
+export function useToolInput() {
+  return useApolloClient()["appManager"].toolInput;
+}

package/src/mcp/react/hooks/useToolMetadata.ts

@@ -0,0 +1,5 @@
+import { useApolloClient } from "./useApolloClient.js";
+
+export function useToolMetadata() {
+  return useApolloClient()["appManager"].toolMetadata;
+}

package/src/mcp/react/hooks/useToolName.ts

@@ -0,0 +1,5 @@
+import { useApolloClient } from "./useApolloClient.js";
+
+export function useToolName() {
+  return useApolloClient()["appManager"].toolName;
+}

package/src/mcp/react/index.ts

@@ -0,0 +1,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";

package/src/openai/core/ApolloClient.ts

@@ -0,0 +1,100 @@
+import type { ApolloLink } from "@apollo/client";
+import { ApolloClient as BaseApolloClient } from "@apollo/client";
+import { DocumentTransform } from "@apollo/client";
+import { removeDirectivesFromDocument } from "@apollo/client/utilities/internal";
+import { parse } from "graphql";
+import { __DEV__ } from "@apollo/client/utilities/environment";
+import type { ApplicationManifest } from "../../types/application-manifest.js";
+import { ToolCallLink } from "../link/ToolCallLink.js";
+import {
+  aiClientSymbol,
+  cacheAsync,
+  getVariablesForOperationFromToolInput,
+} from "../../utilities/index.js";
+import { McpAppManager } from "./McpAppManager.js";
+
+export declare namespace ApolloClient {
+  export interface Options extends Omit<BaseApolloClient.Options, "link"> {
+    link?: BaseApolloClient.Options["link"];
+    manifest: ApplicationManifest;
+  }
+}
+
+export class ApolloClient extends BaseApolloClient {
+  manifest: ApplicationManifest;
+  private readonly appManager: McpAppManager;
+
+  /** @internal */
+  readonly [aiClientSymbol] = true;
+
+  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) => {
+        return removeDirectivesFromDocument(
+          [{ name: "prefetch" }, { name: "tool" }],
+          document
+        )!;
+      }).concat(options.documentTransform ?? DocumentTransform.identity()),
+    });
+
+    this.manifest = options.manifest;
+    this.appManager = new McpAppManager(this.manifest);
+  }
+
+  stop() {
+    super.stop();
+    this.appManager.close().catch(() => {});
+  }
+
+  waitForInitialization = cacheAsync(async () => {
+    const { prefetch, result, toolName, args } =
+      await this.appManager.waitForInitialization();
+
+    this.manifest.operations.forEach((operation) => {
+      if (operation.prefetchID && prefetch?.[operation.prefetchID]) {
+        this.writeQuery({
+          query: parse(operation.body),
+          data: prefetch[operation.prefetchID].data,
+        });
+      }
+
+      if (operation.tools.find((tool) => tool.name === toolName)) {
+        this.writeQuery({
+          query: parse(operation.body),
+          data: result.data,
+          variables: getVariablesForOperationFromToolInput(operation, args),
+        });
+      }
+    });
+  });
+}
+
+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;
+}

package/src/openai/core/McpAppManager.ts

@@ -0,0 +1,110 @@
+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: Record<string, unknown> | null = null;
+  #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
+      );
+
+      // OpenAI is not consistent about sending `ui/notifications/tool-input`
+      // before we get the rool result (which should happen according to the
+      // spec). We resolve this promise in case it wasn't sent to avoid stalling
+      // initialization indefinitely.
+      //
+      // When OpenAI fixes this issue and sends `ui/notifications/tool-input`
+      // consistently, this can be removed.
+      toolInput.resolve({});
+    };
+
+    this.app.ontoolinput = (params) => {
+      toolInput.resolve(params);
+    };
+
+    await this.connect();
+
+    const { structuredContent } = await toolResult.promise;
+    const { arguments: args } = await toolInput.promise;
+
+    this.#toolName = this.app.getHostContext()?.toolInfo?.tool.name;
+    this.#toolInput = args;
+
+    // OpenAI doesn't provide access to `_meta`, so we need to use
+    // window.openai.toolResponseMetadata directly
+    this.#toolMetadata = window.openai.toolResponseMetadata;
+
+    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/openai/core/Platform.ts

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