@apollo/client-ai-apps 0.6.4 → 0.7.0

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

@@ -1,13 +1,22 @@
 import { expect, test, vi, describe } from "vitest";
 import { ApolloClient } from "../ApolloClient.js";
-import { ApolloLink, HttpLink, InMemoryCache, gql } from "@apollo/client";
+import {
+  ApolloLink,
+  HttpLink,
+  InMemoryCache,
+  NetworkStatus,
+  gql,
+  type DocumentNode,
+} from "@apollo/client";
 import { print } from "@apollo/client/utilities";
-import { ToolCallLink } from "../../link/ToolCallLink.js";
+import { ToolCallLink } from "../../../link/ToolCallLink.js";
 import {
   graphqlToolResult,
   minimalHostContextWithToolName,
   mockApplicationManifest,
   mockMcpHost,
+  ObservableStream,
+  parseManifestOperation,
   spyOnConsole,
 } from "../../../testing/internal/index.js";
 
@@ -15,7 +24,8 @@ test("writes tool result data to cache", async () => {
   using _ = spyOnConsole("debug");
 
   const query = gql`
-    query Product($id: ID!) {
+    query Product($id: ID!)
+    @tool(name: "GetProduct", description: "Get a product") {
       product(id: $id) {
         id
         title
@@ -24,27 +34,8 @@ test("writes tool result data to cache", async () => {
     }
   `;
 
-  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());
+  const { client, host } = await setup({ query });
+  using _host = host;
 
   host.sendToolResult({
     content: [],
@@ -60,6 +51,12 @@ test("writes tool result data to cache", async () => {
 
   await client.connect();
 
+  await expect(
+    client.query({ query, variables: { id: "1" } })
+  ).resolves.toStrictEqual({
+    data: { product: { id: "1", title: "Pen", __typename: "Product" } },
+  });
+
   expect(client.extract()).toEqual({
     "Product:1": {
       __typename: "Product",
@@ -140,11 +137,11 @@ test("writes prefetch data to cache", async () => {
   });
 });
 
-test("writes prefetch and tool response data to cache when both are provided", async () => {
+test("writes prefetch response data to cache when both are provided", async () => {
   using _ = spyOnConsole("debug");
 
   const prefetchQuery = gql`
-    query TopProducts {
+    query TopProducts @tool(description: "Shows top products") @prefetch {
       topProducts {
         id
         title
@@ -154,7 +151,7 @@ test("writes prefetch and tool response data to cache when both are provided", a
   `;
 
   const query = gql`
-    query Product($id: ID!) {
+    query Product($id: ID!) @tool(description: "Get a product by id") {
       product(id: $id) {
         id
         title
@@ -167,25 +164,8 @@ test("writes prefetch and tool response data to cache when both are provided", a
     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" }],
-        },
+        parseManifestOperation(prefetchQuery),
+        parseManifestOperation(query),
       ],
     }),
   });
@@ -196,7 +176,6 @@ test("writes prefetch and tool response data to cache when both are provided", a
   host.onCleanup(() => client.stop());
 
   host.sendToolResult({
-    content: [],
     structuredContent: {
       prefetch: {
         __anonymous: {
@@ -216,6 +195,12 @@ test("writes prefetch and tool response data to cache when both are provided", a
 
   await client.connect();
 
+  await expect(
+    client.query({ query, variables: { id: "2" } })
+  ).resolves.toStrictEqual({
+    data: { product: { __typename: "Product", id: "2", title: "iPad" } },
+  });
+
   expect(client.extract()).toEqual({
     "Product:1": {
       __typename: "Product",
@@ -239,7 +224,8 @@ test("excludes extra tool input variables not defined in the operation", async (
   using _ = spyOnConsole("debug");
 
   const query = gql`
-    query Product($id: ID!) {
+    query Product($id: ID!)
+    @tool(name: "GetProduct", description: "Get a product") {
       product(id: $id) {
         id
         title
@@ -248,30 +234,10 @@ test("excludes extra tool input variables not defined in the operation", async (
     }
   `;
 
-  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());
+  const { client, host } = await setup({ query });
+  using _host = host;
 
   host.sendToolResult({
-    content: [],
     structuredContent: {
       result: {
         data: {
@@ -284,7 +250,13 @@ test("excludes extra tool input variables not defined in the operation", async (
 
   await client.connect();
 
-  expect(client.extract()).toEqual({
+  await expect(
+    client.query({ query, variables: { id: "1" } })
+  ).resolves.toStrictEqual({
+    data: { product: { id: "1", title: "Pen", __typename: "Product" } },
+  });
+
+  expect(client.extract()).toStrictEqual({
     "Product:1": {
       __typename: "Product",
       id: "1",
@@ -464,6 +436,511 @@ test("creates a default ToolCallLink when no link is provided", () => {
   }).not.toThrow();
 });
 
+test("reads result data from _meta.structuredContent", async () => {
+  using _ = spyOnConsole("debug");
+
+  const query = gql`
+    query Product($id: ID!)
+    @tool(name: "GetProduct", description: "Get a product") {
+      product(id: $id) @private {
+        id
+        title
+        __typename
+      }
+    }
+  `;
+
+  const { client, host } = await setup({ query });
+  using _host = host;
+
+  host.sendToolResult({
+    structuredContent: {},
+    _meta: {
+      toolName: "GetProduct",
+      structuredContent: {
+        result: {
+          data: {
+            product: { id: "1", title: "Pen", __typename: "Product" },
+          },
+        },
+      },
+    },
+  });
+  host.sendToolInput({ arguments: { id: "1" } });
+
+  await client.connect();
+
+  await expect(
+    client.query({ query, variables: { id: "1" } })
+  ).resolves.toStrictEqual({
+    data: { product: { id: "1", title: "Pen", __typename: "Product" } },
+  });
+
+  expect(client.extract()).toEqual({
+    "Product:1": {
+      __typename: "Product",
+      id: "1",
+      title: "Pen",
+    },
+    ROOT_QUERY: {
+      __typename: "Query",
+      'product({"id":"1"})@private': {
+        __ref: "Product:1",
+      },
+    },
+  });
+});
+
+test("merges prefetch from structuredContent and result from _meta.structuredContent", async () => {
+  using _ = spyOnConsole("debug");
+
+  const prefetchQuery = gql`
+    query TopProducts @tool(description: "Shows top products") @prefetch {
+      topProducts {
+        id
+        title
+        __typename
+      }
+    }
+  `;
+
+  const query = gql`
+    query Product($id: ID!)
+    @tool(name: "GetProduct", description: "Get a product") {
+      product(id: $id) @private {
+        id
+        title
+        __typename
+      }
+    }
+  `;
+
+  const client = new ApolloClient({
+    cache: new InMemoryCache(),
+    manifest: mockApplicationManifest({
+      operations: [
+        parseManifestOperation(prefetchQuery),
+        parseManifestOperation(query),
+      ],
+    }),
+  });
+
+  using host = await mockMcpHost({
+    hostContext: minimalHostContextWithToolName("GetProduct"),
+  });
+  host.onCleanup(() => client.stop());
+
+  host.sendToolResult({
+    structuredContent: {
+      prefetch: {
+        __anonymous: {
+          data: {
+            topProducts: [{ id: "1", title: "iPhone", __typename: "Product" }],
+          },
+        },
+      },
+    },
+    _meta: {
+      toolName: "GetProduct",
+      structuredContent: {
+        result: {
+          data: {
+            product: { id: "2", title: "iPad", __typename: "Product" },
+          },
+        },
+      },
+    },
+  });
+  host.sendToolInput({ arguments: { id: "2" } });
+
+  await client.connect();
+
+  await expect(
+    client.query({ query, variables: { id: "2" } })
+  ).resolves.toStrictEqual({
+    data: { product: { id: "2", title: "iPad", __typename: "Product" } },
+  });
+
+  expect(client.extract()).toStrictEqual({
+    "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"})@private': { __ref: "Product:2" },
+    },
+  });
+});
+
+test("_meta.structuredContent wins over structuredContent", async () => {
+  using _ = spyOnConsole("debug");
+
+  const query = gql`
+    query Product($id: ID!) {
+      product(id: $id) @tool(name: "GetProduct", description: "Get a product") {
+        id
+        title @private
+        __typename
+      }
+    }
+  `;
+
+  const { client, host } = await setup({ query });
+  using _host = host;
+
+  host.sendToolResult({
+    structuredContent: {
+      result: {
+        data: {
+          product: { id: "1", __typename: "Product" },
+        },
+      },
+    },
+    _meta: {
+      toolName: "GetProduct",
+      structuredContent: {
+        result: {
+          data: {
+            product: { id: "1", title: "Meta title", __typename: "Product" },
+          },
+        },
+      },
+    },
+  });
+  host.sendToolInput({ arguments: { id: "1" } });
+
+  await client.connect();
+
+  await expect(
+    client.query({ query, variables: { id: "1" } })
+  ).resolves.toStrictEqual({
+    data: { product: { id: "1", title: "Meta title", __typename: "Product" } },
+  });
+
+  expect(client.extract()).toEqual({
+    "Product:1": {
+      __typename: "Product",
+      id: "1",
+      "title@private": "Meta title",
+    },
+    ROOT_QUERY: {
+      __typename: "Query",
+      'product({"id":"1"})': {
+        __ref: "Product:1",
+      },
+    },
+  });
+});
+
+test("serves tool result data on network-only query without calling execute tool", async () => {
+  using _ = spyOnConsole("debug");
+  const query = gql`
+    query Product($id: ID!)
+    @tool(name: "GetProduct", description: "Get a product") {
+      product(id: $id) {
+        id
+        title
+        __typename
+      }
+    }
+  `;
+
+  const data = {
+    product: { id: "1", title: "Pen", __typename: "Product" },
+  };
+
+  const { client, host } = await setup({ query });
+  using _host = host;
+
+  const execute = vi.fn();
+  host.mockToolCall("execute", execute);
+
+  host.sendToolResult({ structuredContent: { result: { data } } });
+  host.sendToolInput({ arguments: { id: "1" } });
+
+  await client.connect();
+
+  await expect(
+    client.query({
+      query,
+      variables: { id: "1" },
+      fetchPolicy: "network-only",
+    })
+  ).resolves.toStrictEqual({ data });
+  expect(execute).not.toHaveBeenCalled();
+});
+
+test("calls execute tool on second network-only query after hydration is consumed", async () => {
+  using _ = spyOnConsole("debug");
+  const query = gql`
+    query Product($id: ID!)
+    @tool(name: "GetProduct", description: "Get a product") {
+      product(id: $id) {
+        id
+        title
+        __typename
+      }
+    }
+  `;
+
+  const { client, host } = await setup({ query });
+  using _host = host;
+
+  host.mockToolCall("execute", () => ({
+    structuredContent: {
+      data: {
+        product: { id: "1", title: "Updated Pen", __typename: "Product" },
+      },
+    },
+  }));
+
+  host.sendToolResult({
+    structuredContent: {
+      result: {
+        data: {
+          product: { id: "1", title: "Pen", __typename: "Product" },
+        },
+      },
+    },
+  });
+  host.sendToolInput({ arguments: { id: "1" } });
+
+  await client.connect();
+
+  await client.query({
+    query,
+    variables: { id: "1" },
+    fetchPolicy: "network-only",
+  });
+
+  await expect(
+    client.query({
+      query,
+      variables: { id: "1" },
+      fetchPolicy: "network-only",
+    })
+  ).resolves.toStrictEqual({
+    data: {
+      product: { id: "1", title: "Updated Pen", __typename: "Product" },
+    },
+  });
+});
+
+test("serves tool result data on cache-and-network query without calling execute tool", async () => {
+  using _ = spyOnConsole("debug");
+  const query = gql`
+    query Product($id: ID!)
+    @tool(name: "GetProduct", description: "Get a product") {
+      product(id: $id) {
+        id
+        title
+        __typename
+      }
+    }
+  `;
+
+  const { client, host } = await setup({ query });
+  using _host = host;
+
+  const execute = vi.fn();
+  host.mockToolCall("execute", execute);
+
+  host.sendToolResult({
+    structuredContent: {
+      result: {
+        data: {
+          product: { id: "1", title: "Pen", __typename: "Product" },
+        },
+      },
+    },
+  });
+  host.sendToolInput({ arguments: { id: "1" } });
+
+  await client.connect();
+
+  const stream = new ObservableStream(
+    client.watchQuery({
+      query,
+      variables: { id: "1" },
+      fetchPolicy: "cache-and-network",
+    })
+  );
+
+  // The hydrated result is emitted synchronously so we won't observe a loading
+  // state like we normally would with `cache-and-network`
+  await expect(stream).toEmitValue({
+    data: { product: { id: "1", title: "Pen", __typename: "Product" } },
+    dataState: "complete",
+    loading: false,
+    networkStatus: NetworkStatus.ready,
+    partial: false,
+  });
+
+  await expect(stream).not.toEmitAnything();
+
+  expect(execute).not.toHaveBeenCalled();
+});
+
+test("serves tool result data on no-cache query without calling execute tool and does not write to cache", async () => {
+  using _ = spyOnConsole("debug");
+  const query = gql`
+    query Product($id: ID!)
+    @tool(name: "GetProduct", description: "Get a product") {
+      product(id: $id) {
+        id
+        title
+        __typename
+      }
+    }
+  `;
+
+  const data = {
+    product: { id: "1", title: "Pen", __typename: "Product" },
+  };
+
+  const { client, host } = await setup({ query });
+  using _host = host;
+
+  const execute = vi.fn();
+  host.mockToolCall("execute", execute);
+
+  host.sendToolInput({ arguments: { id: "1" } });
+  host.sendToolResult({ structuredContent: { result: { data } } });
+
+  await client.connect();
+
+  await expect(
+    client.query({ query, variables: { id: "1" }, fetchPolicy: "no-cache" })
+  ).resolves.toStrictEqual({ data });
+  expect(execute).not.toHaveBeenCalled();
+  expect(client.extract()).toStrictEqual({});
+});
+
+test("serves hydrated query from tool result while other pending network-only queries call execute", async () => {
+  using _ = spyOnConsole("debug");
+
+  const productQuery = gql`
+    query Product($id: ID!)
+    @tool(name: "GetProduct", description: "Get a product") {
+      product(id: $id) {
+        id
+        title
+        __typename
+      }
+    }
+  `;
+
+  const cartQuery = gql`
+    query Cart @tool(name: "GetCart", description: "Get the cart") {
+      cart {
+        id
+        __typename
+      }
+    }
+  `;
+
+  const productOperation = parseManifestOperation(productQuery);
+  const cartOperation = parseManifestOperation(cartQuery);
+
+  const client = new ApolloClient({
+    cache: new InMemoryCache(),
+    manifest: mockApplicationManifest({
+      operations: [productOperation, cartOperation],
+    }),
+  });
+
+  using host = await mockMcpHost({
+    hostContext: minimalHostContextWithToolName("GetProduct"),
+  });
+  host.onCleanup(() => client.stop());
+
+  const execute = vi.fn(() => ({
+    structuredContent: {
+      data: { cart: { id: "1", __typename: "Cart" } },
+    },
+  }));
+  host.mockToolCall("execute", execute);
+  host.sendToolResult({
+    structuredContent: {
+      result: {
+        data: { product: { id: "1", title: "Pen", __typename: "Product" } },
+      },
+    },
+  });
+  host.sendToolInput({ arguments: { id: "1" } });
+
+  await client.connect();
+
+  const productPromise = client.query({
+    query: productQuery,
+    variables: { id: "1" },
+    fetchPolicy: "network-only",
+  });
+
+  const cartPromise = client.query({
+    query: cartQuery,
+    fetchPolicy: "network-only",
+  });
+
+  await expect(productPromise).resolves.toStrictEqual({
+    data: { product: { id: "1", title: "Pen", __typename: "Product" } },
+  });
+  await expect(cartPromise).resolves.toStrictEqual({
+    data: { cart: { id: "1", __typename: "Cart" } },
+  });
+  expect(execute).toHaveBeenCalledOnce();
+});
+
+test("hydrates prefetch query with network-only fetch policy", async () => {
+  using _ = spyOnConsole("debug");
+
+  const query = gql`
+    query TopProducts @tool(description: "Shows top products") @prefetch {
+      topProducts {
+        id
+        title
+        __typename
+      }
+    }
+  `;
+
+  const data = {
+    topProducts: [{ id: "1", title: "iPhone", __typename: "Product" }],
+  };
+
+  const { client, host } = await setup({ query, toolName: "OtherTool" });
+  using _host = host;
+
+  const execute = vi.fn();
+  host.mockToolCall("execute", execute);
+
+  host.sendToolResult({
+    structuredContent: {
+      prefetch: { __anonymous: { data } },
+    },
+  });
+  host.sendToolInput({ arguments: {} });
+
+  await client.connect();
+
+  await expect(
+    client.query({ query, fetchPolicy: "network-only" })
+  ).resolves.toStrictEqual({
+    data: {
+      topProducts: [{ id: "1", title: "iPhone", __typename: "Product" }],
+    },
+  });
+  expect(execute).not.toHaveBeenCalled();
+});
+
 describe("watchQuery dev warnings", () => {
   const query = gql`
     query Products($category: String!, $page: Int!, $sortBy: String!)
@@ -565,3 +1042,26 @@ describe("watchQuery dev warnings", () => {
     expect(console.warn).toHaveBeenCalledTimes(1);
   });
 });
+
+async function setup({
+  query,
+  toolName,
+}: {
+  query: DocumentNode;
+  toolName?: string;
+}) {
+  const operation = parseManifestOperation(query);
+  const client = new ApolloClient({
+    cache: new InMemoryCache(),
+    manifest: mockApplicationManifest({ operations: [operation] }),
+  });
+
+  const host = await mockMcpHost({
+    hostContext: minimalHostContextWithToolName(
+      toolName ?? operation.tools[0].name
+    ),
+  });
+  host.onCleanup(() => client.stop());
+
+  return { client, host };
+}