Improved hydration when using network-only, cache-and-network, and no-cache fetch policies

.changeset/hydrate_network_only_queries_from_tool_result.md

@@ -0,0 +1,7 @@
+---
+default: patch
+---
+
+# Hydrate `network-only`, `cache-and-network`, and `no-cache` queries from tool result data
+
+Queries using `network-only`, `cache-and-network`, or `no-cache` fetch policies previously called the `execute` tool to fetch data, even when the result was already available from the tool call that initiated the app. These queries are now served directly from the tool result on first load, avoiding a redundant `execute` call. Subsequent queries will continue to call `execute` using the configured fetch policy as expected.

package.json

@@ -123,7 +123,7 @@
     "@types/node": "^24.10.4",
     "@types/react": "^19.2.2",
     "@vitejs/plugin-react": "^5.1.1",
-    "@vitest/coverage-v8": "^4.0.13",
+    "@vitest/coverage-v8": "^4.1.0",
     "copyfiles": "^2.4.1",
     "esbuild": "^0.25.12",
     "graphql": "^16.12.0",
@@ -136,7 +136,7 @@
     "rxjs": "^7.8.1",
     "typescript": "^5.9.3",
     "vite": "^7.3.1",
-    "vitest": "^4.0.13"
+    "vitest": "^4.1.0"
   },
   "peerDependencies": {
     "@apollo/client": "^4.0.0",

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

src/mcp/core/ApolloClient.ts

@@ -22,6 +22,7 @@ import {
   getVariablesForOperationFromToolInput,
   warnOnVariableMismatch,
 } from "../../utilities/index.js";
+import { ToolHydrationLink } from "../../link/ToolHydrationLink.js";
 import { McpAppManager } from "./McpAppManager.js";
 
 export declare namespace ApolloClient {
@@ -39,8 +40,10 @@ export class ApolloClient extends BaseApolloClient {
   readonly [aiClientSymbol] = true;
 
   #toolInput: Record<string, unknown> | undefined;
+  #toolHydrationLink: ToolHydrationLink;
 
   constructor(options: ApolloClient.Options) {
+    const toolHydrationLink = new ToolHydrationLink();
     const link = options.link ?? new ToolCallLink();
 
     if (__DEV__) {
@@ -49,7 +52,7 @@ export class ApolloClient extends BaseApolloClient {
 
     super({
       ...options,
-      link,
+      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(
@@ -65,10 +68,15 @@ export class ApolloClient extends BaseApolloClient {
       }).concat(options.documentTransform ?? DocumentTransform.identity()),
     });
 
+    this.#toolHydrationLink = toolHydrationLink;
     this.manifest = options.manifest;
     this.appManager = new McpAppManager(this.manifest);
   }
 
+  setLink(newLink: ApolloLink): void {
+    super.setLink(this.#toolHydrationLink.concat(newLink));
+  }
+
   stop() {
     super.stop();
     this.appManager.close().catch(() => {});
@@ -141,18 +149,31 @@ export class ApolloClient extends BaseApolloClient {
           query: parse(operation.body),
           data: structuredContent.prefetch[operation.prefetchID].data,
         });
+        this.#toolHydrationLink.hydrate(operation, {
+          result: structuredContent.prefetch[operation.prefetchID],
+          variables: {},
+        });
       }
 
       if (operation.tools.find((tool) => tool.name === toolName)) {
         if (structuredContent.result?.data) {
+          const variables = getVariablesForOperationFromToolInput(
+            operation,
+            args
+          );
+          const result = structuredContent.result;
           this.writeQuery({
             query: parse(operation.body),
-            data: structuredContent.result.data,
-            variables: getVariablesForOperationFromToolInput(operation, args),
+            data: result.data,
+            variables,
           });
+
+          this.#toolHydrationLink.hydrate(operation, { result, variables });
         }
       }
     });
+
+    this.#toolHydrationLink.complete();
   });
 }