@apollo/client-ai-apps 0.5.2 → 0.5.4

package/src/openai/react/hooks/createHydrationUtils.ts

@@ -0,0 +1,182 @@
+import { useState, useCallback, useRef, useMemo, useLayoutEffect } from "react";
+import type {
+  DocumentNode,
+  OperationVariables,
+  TypedDocumentNode,
+} from "@apollo/client";
+import { useApolloClient } from "./useApolloClient.js";
+import { useToolName } from "./useToolName.js";
+import { isReactive } from "../../../react/reactive.js";
+import type { Reactive } from "../../../react/reactive.js";
+import { equal } from "@wry/equality";
+import { __DEV__ } from "@apollo/client/utilities/environment";
+import {
+  getToolNamesFromDocument,
+  getVariableNamesFromDocument,
+} from "../../../utilities/index.js";
+
+type HydratedVariablesInput<TVariables> = {
+  [K in keyof TVariables]: TVariables[K] | Reactive<TVariables[K]>;
+};
+
+type StateVariables<TVariables, Input> = {
+  [K in keyof TVariables as K extends keyof Input ?
+    Input[K] extends Reactive<any> ?
+      never
+    : K
+  : K]: K extends keyof TVariables ? TVariables[K] : never;
+};
+
+type SetVariables<TState> = (
+  update: Partial<TState> | ((prev: TState) => Partial<TState>)
+) => void;
+
+/** @experimental */
+export function createHydrationUtils<
+  TVariables extends OperationVariables = OperationVariables,
+>(document: TypedDocumentNode<any, TVariables> | DocumentNode) {
+  const documentToolNames = getToolNamesFromDocument(document);
+  const variableNames = getVariableNamesFromDocument(document);
+
+  function useHydratedVariables<
+    TInputVariables extends HydratedVariablesInput<TVariables>,
+  >(
+    variables: TInputVariables &
+      Record<Exclude<keyof TInputVariables, keyof TVariables>, never>
+  ): [
+    variables: TVariables,
+    setVariables: SetVariables<StateVariables<TVariables, TInputVariables>>,
+  ] {
+    const client = useApolloClient();
+    const toolName = useToolName();
+    const [toolInput] = useState(() => client.toolInput);
+
+    const toolMatches =
+      toolInput !== undefined &&
+      toolName !== undefined &&
+      documentToolNames.has(toolName);
+
+    const [stateVars, setStateVars] = useState<Record<string, unknown>>(() => {
+      const values: Record<string, unknown> = {};
+
+      for (const [key, value] of Object.entries(
+        toolMatches ? toolInput : variables
+      )) {
+        if (variableNames.has(key) && !isReactive(value)) {
+          values[key] = value;
+        }
+      }
+
+      return values;
+    });
+
+    const [initialReactiveValues] = useState<Record<string, unknown>>(() => {
+      const values: Record<string, unknown> = {};
+
+      for (const [key, value] of Object.entries(variables)) {
+        if (variableNames.has(key) && isReactive(value)) {
+          values[key] = value.value;
+        }
+      }
+
+      return values;
+    });
+
+    const [reactiveVars, setReactiveVars] = useState(() => {
+      const values: Record<string, unknown> = {};
+
+      for (const [key, value] of Object.entries(initialReactiveValues)) {
+        if (toolMatches && key in toolInput) {
+          values[key] = toolInput[key];
+        } else if (!toolMatches) {
+          values[key] = value;
+        }
+      }
+
+      return values;
+    });
+
+    const changedKeysRef = useRef(new Set<string>());
+    const nextReactiveVars: Record<string, unknown> = {};
+
+    for (const [key, value] of Object.entries(variables)) {
+      if (!variableNames.has(key) || !isReactive(value)) continue;
+
+      const hasChanged =
+        changedKeysRef.current.has(key) ||
+        !equal(value.value, initialReactiveValues[key]);
+
+      if (toolMatches && !hasChanged) {
+        if (key in toolInput) {
+          nextReactiveVars[key] = toolInput[key];
+        }
+      } else {
+        nextReactiveVars[key] = value.value;
+      }
+    }
+
+    if (!equal(nextReactiveVars, reactiveVars)) {
+      setReactiveVars(nextReactiveVars);
+    }
+
+    // Clear the tool input after first mount so that remounting the component
+    // uses the user-provided variables rather than the hydrated tool input.
+    // This runs once on mount; watchQuery also clears it when useQuery is
+    // present, so both paths are idempotent.
+    useLayoutEffect(() => {
+      if (toolMatches) {
+        client.clearToolInput();
+      }
+      // eslint-disable-next-line react-hooks/exhaustive-deps
+    }, []);
+
+    useLayoutEffect(() => {
+      for (const [key, value] of Object.entries(variables)) {
+        if (
+          variableNames.has(key) &&
+          isReactive(value) &&
+          !changedKeysRef.current.has(key) &&
+          !equal(value.value, initialReactiveValues[key])
+        ) {
+          changedKeysRef.current.add(key);
+        }
+      }
+    });
+
+    const resolvedVariables = useMemo(() => {
+      return { ...stateVars, ...reactiveVars } as TVariables;
+    }, [stateVars, reactiveVars]);
+
+    const setVariables = useCallback<
+      SetVariables<StateVariables<TVariables, TInputVariables>>
+    >((update) => {
+      setStateVars((prev) => {
+        const updates =
+          typeof update === "function" ? update(prev as any) : update;
+
+        const filtered = Object.fromEntries(
+          Object.entries(updates).filter(([key]) => {
+            if (key in initialReactiveValues) {
+              if (__DEV__) {
+                console.warn(
+                  `Attempted to set reactive variable "${key}" via setVariables. ` +
+                    `Reactive variables are read-only and are ignored. `
+                );
+              }
+              return false;
+            }
+            return true;
+          })
+        );
+
+        if (Object.keys(filtered).length === 0) return prev;
+
+        return { ...prev, ...filtered };
+      });
+    }, []);
+
+    return [resolvedVariables, setVariables];
+  }
+
+  return { useHydratedVariables };
+}

package/src/openai/react/index.ts

@@ -2,4 +2,5 @@ export { useApp } from "./hooks/useApp.js";
 export { useToolName } from "./hooks/useToolName.js";
 export { useToolMetadata } from "./hooks/useToolMetadata.js";
 export { useToolInput } from "./hooks/useToolInput.js";
+export { createHydrationUtils } from "./hooks/createHydrationUtils.js";
 export { useWidgetState } from "./hooks/useWidgetState.js";

package/src/react/ApolloProvider.tsx

@@ -25,7 +25,7 @@ export function ApolloProvider({ children, client }: ApolloProvider.Props) {
     );
   }
 
-  use(client.waitForInitialization());
+  use(client.connect());
 
   return (
     <BaseApolloProvider client={client as BaseApolloClient}>

package/src/react/index.mcp.ts

@@ -2,6 +2,7 @@ export * from "./index.js";
 export {
   useApp,
   useToolInput,
+  createHydrationUtils,
   useToolMetadata,
   useToolName,
 } from "../mcp/react/index.js";

package/src/react/index.openai.ts

@@ -2,6 +2,7 @@ export * from "./index.js";
 export {
   useApp,
   useToolInput,
+  createHydrationUtils,
   useToolMetadata,
   useToolName,
 } from "../openai/react/index.js";

package/src/react/index.ts

@@ -1,6 +1,8 @@
 import { missingHook } from "./missingHook.js";
 
 export { ApolloProvider } from "./ApolloProvider.js";
+export { reactive } from "./reactive.js";
+export type { Reactive } from "./reactive.js";
 
 // Use `mcp` related types since these are the most common between the two
 // targets
@@ -17,3 +19,8 @@ export const useToolMetadata =
 
 export const useToolName =
   missingHook<typeof import("./index.mcp.js").useToolName>("useToolName");
+
+/** @experimental */
+export const createHydrationUtils = missingHook<
+  typeof import("./index.mcp.js").createHydrationUtils
+>("createHydrationUtils");

package/src/react/reactive.ts

@@ -0,0 +1,19 @@
+const REACTIVE = Symbol("apollo.reactive");
+
+export interface Reactive<T> {
+  readonly [REACTIVE]: true;
+  readonly value: T;
+}
+
+export function reactive<T>(value: T): Reactive<T> {
+  return { [REACTIVE]: true, value };
+}
+
+export function isReactive(value: unknown): value is Reactive<unknown> {
+  return (
+    typeof value === "object" &&
+    value !== null &&
+    REACTIVE in value &&
+    (value as Record<symbol, unknown>)[REACTIVE] === true
+  );
+}

package/src/testing/internal/mcp/graphqlToolResult.ts

@@ -2,20 +2,20 @@ import type { McpUiToolResultNotification } from "@modelcontextprotocol/ext-apps
 import type { FormattedExecutionResult } from "graphql";
 import type { ApolloMcpServerApps } from "../../../core/types";
 
-export function graphqlToolResult(
+export function graphqlToolResult<TData = Record<string, unknown>>(
   options:
-    | FormattedExecutionResult
+    | FormattedExecutionResult<TData>
     | {
         prefetch?: Record<string, FormattedExecutionResult>;
-        result: FormattedExecutionResult;
+        result: FormattedExecutionResult<TData>;
       }
 ): McpUiToolResultNotification["params"] {
   let structuredContent: ApolloMcpServerApps.StructuredContent;
 
   if ("data" in options) {
-    structuredContent = { result: options };
+    structuredContent = { result: options as FormattedExecutionResult };
   } else if ("result" in options) {
-    structuredContent = { result: options.result };
+    structuredContent = { result: options.result as FormattedExecutionResult };
 
     if (options.prefetch) {
       structuredContent.prefetch = options.prefetch;

package/src/utilities/getToolNamesFromDocument.ts

@@ -0,0 +1,15 @@
+import { getOperationDefinition } from "@apollo/client/utilities/internal";
+import { Kind, type DocumentNode } from "graphql";
+
+export function getToolNamesFromDocument(document: DocumentNode) {
+  const operationDef = getOperationDefinition(document);
+
+  return new Set(
+    operationDef?.directives
+      ?.filter((d) => d.name.value === "tool")
+      .flatMap((d) => {
+        const nameArg = d.arguments?.find((arg) => arg.name.value === "name");
+        return nameArg?.value.kind === Kind.STRING ? [nameArg.value.value] : [];
+      })
+  );
+}

package/src/utilities/getVariableNamesFromDocument.ts

@@ -0,0 +1,9 @@
+import { getOperationDefinition } from "@apollo/client/utilities/internal";
+import type { DocumentNode } from "graphql";
+
+export function getVariableNamesFromDocument(document: DocumentNode) {
+  const operationDef = getOperationDefinition(document);
+  return new Set(
+    operationDef?.variableDefinitions?.map((v) => v.variable.name.value) ?? []
+  );
+}

package/src/utilities/index.ts

@@ -1,5 +1,8 @@
 export { aiClientSymbol } from "./constants.js";
 export { getVariablesForOperationFromToolInput } from "./getVariablesForOperationFromToolInput.js";
+export { getToolNamesFromDocument } from "./getToolNamesFromDocument.js";
+export { getVariableNamesFromDocument } from "./getVariableNamesFromDocument.js";
 export { cacheAsync } from "./cacheAsync.js";
 export { invariant } from "./invariant.js";
 export { promiseWithResolvers } from "./promiseWithResolvers.js";
+export { warnOnVariableMismatch } from "./warnOnVariableMismatch.js";

package/src/utilities/warnOnVariableMismatch.ts

@@ -0,0 +1,20 @@
+import type { DocumentNode, OperationVariables } from "@apollo/client";
+import { getOperationName } from "@apollo/client/utilities/internal/internal.cjs";
+
+export function warnOnVariableMismatch(
+  document: DocumentNode,
+  toolInput: OperationVariables,
+  actualVariables: OperationVariables | undefined
+) {
+  const operationName = getOperationName(document, "(anonymous)");
+
+  console.warn(
+    `The operation "${operationName}" has a @tool directive matching the current ` +
+      "tool call, but the variables provided to the query don't match the tool " +
+      "input. Use the `useHydratedVariables` hook returned from `createHydrationUtils` " +
+      "to provide the hydrated variables to the query. " +
+      "\n\nExpected variables:\n%o\n\nReceived:\n%o",
+    toolInput,
+    actualVariables
+  );
+}

package/src/vite/__tests__/apolloClientAiApps.test.ts

@@ -1,7 +1,7 @@
 import { beforeEach, describe, expect, test, vi } from "vitest";
 import fs from "node:fs";
 import { gql, type DocumentNode } from "@apollo/client";
-import { print } from "@apollo/client/utilities";
+import { getMainDefinition, print } from "@apollo/client/utilities";
 import { getOperationName } from "@apollo/client/utilities/internal";
 import { vol } from "memfs";
 import { apolloClientAiApps } from "../apolloClientAiApps.js";
@@ -11,6 +11,8 @@ import type {
   ManifestWidgetSettings,
 } from "../../types/application-manifest.js";
 import { explorer } from "../utilities/config.js";
+import { invariant } from "@apollo/client/utilities/invariant";
+import { Kind } from "graphql";
 
 beforeEach(() => {
   explorer.clearCaches();
@@ -121,6 +123,173 @@ describe("operations", () => {
     `);
   });
 
+  test("handles operations with fragments in the same file", async () => {
+    vol.fromJSON({
+      "package.json": mockPackageJson(),
+      "src/my-component.tsx": declareOperation(gql`
+        query HelloWorldQuery($name: string!)
+        @tool(name: "hello-world", description: "This is an awesome tool!") {
+          greeting {
+            message
+            recipient {
+              ...RecipientFragment
+            }
+          }
+        }
+
+        fragment RecipientFragment on Recipient {
+          id
+          name
+        }
+      `),
+    });
+
+    await using server = await setupServer({
+      plugins: [apolloClientAiApps({ targets: ["mcp"] })],
+    });
+    await server.listen();
+
+    const manifest = readManifestFile();
+    expect(manifest).toMatchInlineSnapshot(`
+      {
+        "appVersion": "1.0.0",
+        "csp": {
+          "connectDomains": [],
+          "frameDomains": [],
+          "redirectDomains": [],
+          "resourceDomains": [],
+        },
+        "format": "apollo-ai-app-manifest",
+        "hash": "abc",
+        "operations": [
+          {
+            "body": "query HelloWorldQuery {
+        greeting {
+          message
+          recipient {
+            ...RecipientFragment
+            __typename
+          }
+          __typename
+        }
+      }
+
+      fragment RecipientFragment on Recipient {
+        id
+        name
+        __typename
+      }",
+            "id": "1646a86ae2ff5ad75457161be5cff80f3ba5172da573a0fc796b268870119020",
+            "name": "HelloWorldQuery",
+            "prefetch": false,
+            "tools": [
+              {
+                "description": "This is an awesome tool!",
+                "name": "hello-world",
+              },
+            ],
+            "type": "query",
+            "variables": {
+              "name": "string",
+            },
+          },
+        ],
+        "resource": "http://localhost:3333",
+        "version": "1",
+      }
+    `);
+  });
+
+  test("handles operations with fragments in other files", async () => {
+    vol.fromJSON({
+      "package.json": mockPackageJson(),
+      "src/first-recipient.tsx": declareFragment(gql`
+        fragment OtherRecipientFragment on Recipient {
+          name
+        }
+      `),
+      "src/my-component.tsx": declareOperation(gql`
+        query HelloWorldQuery($name: string!)
+        @tool(name: "hello-world", description: "This is an awesome tool!") {
+          greeting {
+            message
+            recipient {
+              ...RecipientFragment
+              ...OtherRecipientFragment
+            }
+          }
+        }
+      `),
+      "src/my-fragment.tsx": declareFragment(gql`
+        fragment RecipientFragment on Recipient {
+          id
+          name
+        }
+      `),
+    });
+
+    await using server = await setupServer({
+      plugins: [apolloClientAiApps({ targets: ["mcp"] })],
+    });
+    await server.listen();
+
+    const manifest = readManifestFile();
+    expect(manifest).toMatchInlineSnapshot(`
+      {
+        "appVersion": "1.0.0",
+        "csp": {
+          "connectDomains": [],
+          "frameDomains": [],
+          "redirectDomains": [],
+          "resourceDomains": [],
+        },
+        "format": "apollo-ai-app-manifest",
+        "hash": "abc",
+        "operations": [
+          {
+            "body": "query HelloWorldQuery {
+        greeting {
+          message
+          recipient {
+            ...RecipientFragment
+            ...OtherRecipientFragment
+            __typename
+          }
+          __typename
+        }
+      }
+
+      fragment OtherRecipientFragment on Recipient {
+        name
+        __typename
+      }
+
+      fragment RecipientFragment on Recipient {
+        id
+        name
+        __typename
+      }",
+            "id": "c65cb5ec2dc76bbfa992cd2a98c05ab3f909349f3f1478e608a7f16ae29bdd4a",
+            "name": "HelloWorldQuery",
+            "prefetch": false,
+            "tools": [
+              {
+                "description": "This is an awesome tool!",
+                "name": "hello-world",
+              },
+            ],
+            "type": "query",
+            "variables": {
+              "name": "string",
+            },
+          },
+        ],
+        "resource": "http://localhost:3333",
+        "version": "1",
+      }
+    `);
+  });
+
   test("does not write to dev application manifest file when using a build command", async () => {
     vol.fromJSON({
       "package.json": mockPackageJson(),
@@ -289,7 +458,7 @@ describe("operations", () => {
       });
       await server.listen();
     }).rejects.toThrowErrorMatchingInlineSnapshot(
-      `[Error: Found an unsupported operation type. Only Query and Mutation are supported.]`
+      `[Error: Found unsupported operation type 'subscription'. Only queries and mutations are supported.]`
     );
   });
 
@@ -1461,6 +1630,18 @@ function declareOperation(operation: DocumentNode) {
   return `const ${varName} = gql\`\n${print(operation)}\n\``;
 }
 
+function declareFragment(fragment: DocumentNode) {
+  const definition = getMainDefinition(fragment);
+  invariant(
+    definition.kind === Kind.FRAGMENT_DEFINITION,
+    "declareFragment must receive a fragment definition"
+  );
+
+  const name = definition.name.value;
+  const varName = name.replace(/([a-z])([A-Z])/g, "$1_$2").toUpperCase();
+  return `const ${varName} = gql\`\n${print(fragment)}\n\``;
+}
+
 function mockPackageJson(config?: Record<string, unknown>) {
   return JSON.stringify({ version: "1.0.0", ...config });
 }