@apollo/client-ai-apps 0.6.4 → 0.7.0

package/src/core/AbstractApolloClient.ts

@@ -0,0 +1,217 @@
+import {
+  ApolloLink,
+  ApolloClient as NativeApolloClient,
+  DocumentTransform,
+  ObservableQuery,
+  type OperationVariables,
+  type WatchQueryOptions,
+} from "@apollo/client";
+import { __DEV__ } from "@apollo/client/utilities/environment";
+import { removeDirectivesFromDocument } from "@apollo/client/utilities/internal";
+import { parse, visit } from "graphql";
+import { equal } from "@wry/equality";
+import type { ApplicationManifest } from "../types/application-manifest.js";
+import { aiClientSymbol } from "../utilities/constants.js";
+import {
+  McpAppManager,
+  type ConnectToHostImplementation,
+} from "./McpAppManager.js";
+import { ToolHydrationLink } from "../link/ToolHydrationLink.js";
+import { ToolCallLink } from "../link/ToolCallLink.js";
+import {
+  cacheAsync,
+  getToolNamesFromDocument,
+  getVariableNamesFromDocument,
+  getVariablesForOperationFromToolInput,
+  warnOnVariableMismatch,
+} from "../utilities/index.js";
+import type { ApolloMcpServerApps } from "./types.js";
+import type { ToolInfo } from "./typeRegistration.js";
+
+export declare namespace AbstractApolloClient {
+  export interface Options extends Omit<NativeApolloClient.Options, "link"> {
+    link?: NativeApolloClient.Options["link"];
+    manifest: ApplicationManifest;
+  }
+}
+
+export class AbstractApolloClient extends NativeApolloClient {
+  manifest: ApplicationManifest;
+  private readonly appManager: McpAppManager;
+
+  /** @internal */
+  readonly [aiClientSymbol] = true;
+
+  #hydratedToolInput: Record<string, unknown> | undefined;
+  #toolHydrationLink: ToolHydrationLink;
+
+  #toolInfo: ToolInfo | undefined;
+  #toolMetadata: ApolloMcpServerApps.CallToolResult["_meta"] | undefined;
+
+  constructor(
+    options: AbstractApolloClient.Options,
+    connectToHost: ConnectToHostImplementation
+  ) {
+    const toolHydrationLink = new ToolHydrationLink();
+    const link = options.link ?? new ToolCallLink();
+
+    if (__DEV__) {
+      validateTerminatingLink(link);
+    }
+
+    super({
+      ...options,
+      link: toolHydrationLink.concat(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.#toolHydrationLink = toolHydrationLink;
+    this.manifest = options.manifest;
+    this.appManager = new McpAppManager(this.manifest, connectToHost);
+  }
+
+  get toolInfo() {
+    return this.#toolInfo;
+  }
+
+  get toolMetadata() {
+    return this.#toolMetadata;
+  }
+
+  setLink(newLink: ApolloLink): void {
+    super.setLink(this.#toolHydrationLink.concat(newLink));
+  }
+
+  stop() {
+    super.stop();
+    this.appManager.close().catch(() => {});
+  }
+
+  protected get hydratedToolInput() {
+    return this.#hydratedToolInput;
+  }
+
+  protected clearHydratedToolInput() {
+    this.#hydratedToolInput = undefined;
+  }
+
+  watchQuery<
+    T = any,
+    TVariables extends OperationVariables = OperationVariables,
+  >(options: WatchQueryOptions<TVariables, T>): ObservableQuery<T, TVariables> {
+    if (__DEV__) {
+      const toolInput = this.#hydratedToolInput;
+
+      if (toolInput) {
+        const toolName = this.toolInfo?.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.#hydratedToolInput = 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, toolInput, _meta } =
+      await this.appManager.connect();
+
+    this.#hydratedToolInput = toolInput;
+    this.#toolMetadata = _meta;
+
+    if (toolName != null) {
+      this.#toolInfo = { toolName, toolInput };
+    }
+
+    this.manifest.operations.forEach((operation) => {
+      if (
+        operation.prefetchID &&
+        structuredContent.prefetch?.[operation.prefetchID]
+      ) {
+        this.writeQuery({
+          query: parse(operation.body),
+          data: structuredContent.prefetch[operation.prefetchID].data,
+        });
+        this.#toolHydrationLink.hydrate(operation, {
+          result: structuredContent.prefetch[operation.prefetchID],
+          variables: {},
+        });
+      }
+
+      if (
+        structuredContent.result &&
+        operation.tools.find((tool) => tool.name === toolName)
+      ) {
+        this.#toolHydrationLink.hydrate(operation, {
+          result: structuredContent.result,
+          variables: getVariablesForOperationFromToolInput(
+            operation,
+            toolInput
+          ),
+        });
+      }
+    });
+
+    this.#toolHydrationLink.complete();
+  });
+}
+
+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/core/ApolloClient.ts

@@ -1,25 +1,23 @@
-import { ApolloClient as BaseApolloClient } from "@apollo/client";
-import type { ApplicationManifest } from "../types/application-manifest.js";
+import { AbstractApolloClient } from "./AbstractApolloClient.js";
 import { aiClientSymbol } from "../utilities/constants.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 {
+export class ApolloClient extends AbstractApolloClient {
   /** @internal */
   readonly [aiClientSymbol] = true;
 
   constructor(options: ApolloClient.Options) {
-    super(options as any);
+    super(options as any, () => {
+      throw new Error(
+        "Cannot connect an `ApolloClient` instance from `@apollo/client-ai-apps` without export conditions. Please set conditions or import from the `/openai` or `/mcp` subpath directly."
+      );
+    });
 
     throw new Error(
       "Cannot construct an `ApolloClient` instance from `@apollo/client-ai-apps` without export conditions. Please set conditions or import from the `/openai` or `/mcp` subpath directly."
     );
   }
-
-  async connect() {}
 }

package/src/core/McpAppManager.ts

@@ -0,0 +1,106 @@
+import {
+  App,
+  type McpUiHostContextChangedNotification,
+} from "@modelcontextprotocol/ext-apps";
+import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import type { ApplicationManifest } from "../types/application-manifest.js";
+import type { FormattedExecutionResult } from "graphql";
+import type { DocumentNode, OperationVariables } from "@apollo/client";
+import { print } from "@apollo/client/utilities";
+import { cacheAsync } from "../utilities/index.js";
+import type { ApolloMcpServerApps } from "./types.js";
+
+type ExecuteQueryCallToolResult = Omit<CallToolResult, "structuredContent"> & {
+  structuredContent: FormattedExecutionResult;
+  _meta?: { structuredContent?: FormattedExecutionResult };
+};
+
+interface ConnectToHostResult {
+  structuredContent: ApolloMcpServerApps.StructuredContent;
+  toolInput: Record<string, unknown> | undefined;
+  _meta: ApolloMcpServerApps.Meta | undefined;
+}
+
+export type ConnectToHostImplementation = (
+  app: App
+) => Promise<ConnectToHostResult>;
+
+/** @internal */
+export class McpAppManager {
+  readonly app: App;
+
+  #connectToHost: ConnectToHostImplementation;
+
+  #hostContextCallbacks = new Set<
+    (params: McpUiHostContextChangedNotification["params"]) => void
+  >();
+
+  constructor(
+    manifest: ApplicationManifest,
+    connectToHost: ConnectToHostImplementation
+  ) {
+    this.app = new App({ name: manifest.name, version: manifest.appVersion });
+    this.#connectToHost = connectToHost;
+  }
+
+  onHostContextChanged(
+    cb: (params: McpUiHostContextChangedNotification["params"]) => void
+  ) {
+    this.#hostContextCallbacks.add(cb);
+
+    return () => {
+      this.#hostContextCallbacks.delete(cb);
+    };
+  }
+
+  connect = cacheAsync(async () => {
+    this.app.onhostcontextchanged = (params) => {
+      this.#hostContextCallbacks.forEach((cb) => cb(params));
+    };
+
+    const { structuredContent, _meta, toolInput } = await this.#connectToHost(
+      this.app
+    );
+
+    return {
+      structuredContent: {
+        ...structuredContent,
+        ..._meta?.structuredContent,
+      },
+      toolInput,
+      // Some hosts do not provide toolInfo in the ui/initialize response, so we
+      // fallback to `_meta.toolName` provided by Apollo MCP server if the value
+      // is not available.
+      toolName:
+        this.app.getHostContext()?.toolInfo?.tool.name ??
+        _meta?.toolName ??
+        // Some hosts do not forward `_meta` nor do they provide `toolInfo`. Our
+        // MCP server provides `toolName` in `structuredContent` as a workaround
+        // that we can use if all else fails
+        structuredContent.toolName,
+      _meta,
+    };
+  });
+
+  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,
+      ...result._meta?.structuredContent,
+    };
+  }
+}

package/src/core/typeRegistration.ts

@@ -1,24 +1,9 @@
 export interface Register {}
 
-/**
- * @deprecated Please use the `ToolInfo` type instead. `ToolName` will be
- * removed in the next major version.
- */
-export type ToolName =
-  Register extends { toolName: infer T extends string } ? T : string;
-
 type RegisteredToolInputs =
   Register extends { toolInputs: infer T extends Record<string, unknown> } ? T
   : never;
 
-/**
- * @deprecated Please use the `ToolInfo` type instead. `ToolInput` will be
- * removed in the next major version.
- */
-export type ToolInput =
-  [RegisteredToolInputs] extends [never] ? Record<string, unknown>
-  : RegisteredToolInputs[keyof RegisteredToolInputs];
-
 type ToolInfoFromInputs<T extends Record<string, unknown>> = {
   [K in keyof T]: {
     toolName: K;

package/src/core/types.ts

@@ -4,11 +4,12 @@ import type { FormattedExecutionResult } from "graphql";
 export namespace ApolloMcpServerApps {
   export interface Meta {
     toolName: string;
+    structuredContent?: ApolloMcpServerApps.StructuredContent;
     [x: string]: unknown;
   }
 
   export interface StructuredContent {
-    result: FormattedExecutionResult;
+    result?: FormattedExecutionResult;
     prefetch?: Record<string, FormattedExecutionResult>;
     toolName?: string;
     [x: string]: unknown;

package/src/index.mcp.ts

@@ -2,4 +2,3 @@ export * from "./index.js";
 
 export { Platform } from "./mcp/core/Platform.js";
 export { ApolloClient } from "./mcp/core/ApolloClient.js";
-export { ToolCallLink } from "./mcp/link/ToolCallLink.js";

package/src/index.openai.ts

@@ -2,4 +2,3 @@ export * from "./index.js";
 
 export { Platform } from "./openai/core/Platform.js";
 export { ApolloClient } from "./openai/core/ApolloClient.js";
-export { ToolCallLink } from "./openai/link/ToolCallLink.js";

package/src/index.ts

@@ -8,12 +8,7 @@ export type {
   ManifestWidgetSettings,
 } from "./types/application-manifest.js";
 
-export type {
-  Register,
-  ToolInfo,
-  ToolInput,
-  ToolName,
-} from "./core/typeRegistration.js";
+export type { Register, ToolInfo } from "./core/typeRegistration.js";
 
 export { ApolloClient } from "./core/ApolloClient.js";
 export { ToolCallLink } from "./link/ToolCallLink.js";

package/src/link/ToolCallLink.ts

@@ -1,4 +1,9 @@
-import { ApolloLink } from "@apollo/client";
+import { ApolloClient, ApolloLink } from "@apollo/client";
+import { __DEV__ } from "@apollo/client/utilities/environment";
+import { from } from "rxjs";
+
+import type { McpAppManager } from "../core/McpAppManager";
+import { aiClientSymbol, invariant } from "../utilities/index.js";
 
 /**
  * A terminating link that sends a GraphQL request through an agent tool call.
@@ -23,11 +28,28 @@ import { ApolloLink } from "@apollo/client";
  * ```
  */
 export class ToolCallLink extends ApolloLink {
-  constructor() {
-    super();
+  readonly name = "ToolCallLink";
+  request(operation: ApolloLink.Operation) {
+    const client = getPrivateAccess(operation.client);
+
+    return from(
+      client.appManager.executeQuery({
+        query: operation.query,
+        variables: operation.variables,
+      })
+    );
+  }
+}
 
-    throw new Error(
-      "Cannot construct a `ToolCallLink` from `@apollo/client-ai-apps` without export conditions. Please set export conditions or import from  the `/openai` or `/mcp` subpath directly."
+function getPrivateAccess(
+  client: ApolloClient
+): ApolloClient & { appManager: McpAppManager } {
+  if (__DEV__) {
+    invariant(
+      (client as any)[aiClientSymbol],
+      'The "client" instance used with `ToolCallLink` is the wrong instance. You might have imported `ApolloClient` from `@apollo/client`. Please import `ApolloClient` from `@apollo/client-ai-apps` instead.'
     );
   }
+
+  return client as any;
 }

package/src/link/ToolHydrationLink.ts

@@ -0,0 +1,90 @@
+import { ApolloLink, Observable } from "@apollo/client";
+import type { OperationVariables } from "@apollo/client";
+import { canonicalStringify } from "@apollo/client/utilities";
+import type { FormattedExecutionResult } from "graphql";
+import { of } from "rxjs";
+import type { ManifestOperation } from "../types/application-manifest.js";
+
+type ResolveFn = () => void;
+type OperationKey = string & { __type: "PendingKey" };
+
+/**
+ * @internal
+ * Holds requests until the client is fully hydrated after a tool call. It is
+ * particularly useful for `network-only` and `cache-and-network` fetch policies
+ * where the request makes it to the link chain because it prevents a followup
+ * tool call to the MCP server for the just fetched data from the MCP server.
+ */
+export class ToolHydrationLink extends ApolloLink {
+  #hydrated = false;
+  #pending: ResolveFn[] = [];
+  #operations = new Map<OperationKey, FormattedExecutionResult>();
+
+  hydrate(
+    operation: ManifestOperation,
+    {
+      result,
+      variables,
+    }: { result: FormattedExecutionResult; variables: OperationVariables }
+  ) {
+    this.#operations.set(
+      getKey({ operationName: operation.name, variables }),
+      result
+    );
+  }
+
+  complete(): void {
+    this.#hydrated = true;
+
+    for (const resolve of this.#pending.splice(0)) {
+      resolve();
+    }
+  }
+
+  request(
+    operation: ApolloLink.Operation,
+    forward: ApolloLink.ForwardFunction
+  ): Observable<ApolloLink.Result> {
+    const maybeSendToTerminatingLink = (): Observable<ApolloLink.Result> => {
+      const key = getKey(operation);
+      const result = this.#operations.get(key);
+
+      if (result) {
+        this.#operations.delete(key);
+        return of(result);
+      }
+
+      return forward(operation);
+    };
+
+    if (this.#hydrated) {
+      return maybeSendToTerminatingLink();
+    }
+
+    return new Observable((observer) => {
+      let active = true;
+
+      const resolve = () => {
+        if (!active) return;
+        maybeSendToTerminatingLink().subscribe(observer);
+      };
+      this.#pending.push(resolve);
+
+      return () => {
+        active = false;
+        const idx = this.#pending.indexOf(resolve);
+        if (idx !== -1) this.#pending.splice(idx, 1);
+      };
+    });
+  }
+}
+
+function getKey({
+  operationName,
+  variables,
+}: {
+  operationName: string | undefined;
+  variables?: OperationVariables;
+}): OperationKey {
+  return `${operationName}:${canonicalStringify(variables ?? {})}` as OperationKey;
+}

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();
+  });
+});