@apollo/client-ai-apps 0.4.0 → 0.5.0

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

@@ -0,0 +1,86 @@
+import { afterEach, expect, test, vi } from "vitest";
+import {
+  mockApplicationManifest,
+  mockMcpHost,
+  spyOnConsole,
+  stubOpenAiGlobals,
+} from "../../../../testing/internal/index.js";
+import {
+  disableActEnvironment,
+  renderHookToSnapshotStream,
+} from "@testing-library/react-render-stream";
+import { useToolMetadata } from "../useToolMetadata.js";
+import { ApolloClient } from "../../../core/ApolloClient.js";
+import { InMemoryCache } from "@apollo/client";
+import { Suspense } from "react";
+import { ApolloProvider } from "../../../../react/ApolloProvider.js";
+
+afterEach(() => {
+  vi.unstubAllGlobals();
+});
+
+test("returns the tool metadata from window.openai", async () => {
+  using _ = spyOnConsole("debug");
+  stubOpenAiGlobals({ toolResponseMetadata: { foo: true } });
+  const client = new ApolloClient({
+    cache: new InMemoryCache(),
+    manifest: mockApplicationManifest(),
+  });
+
+  using host = await mockMcpHost();
+  host.onCleanup(() => client.stop());
+
+  host.sendToolInput({ arguments: {} });
+  host.sendToolResult({
+    content: [],
+    structuredContent: {},
+  });
+
+  using _disabledAct = disableActEnvironment();
+  const { takeSnapshot } = await renderHookToSnapshotStream(
+    () => useToolMetadata(),
+    {
+      wrapper: ({ children }) => (
+        <Suspense>
+          <ApolloProvider client={client}>{children}</ApolloProvider>
+        </Suspense>
+      ),
+    }
+  );
+
+  await expect(takeSnapshot()).resolves.toEqual({ foo: true });
+  await expect(takeSnapshot).not.toRerender();
+});
+
+test("returns null when not set", async () => {
+  using _ = spyOnConsole("debug");
+  stubOpenAiGlobals();
+  const client = new ApolloClient({
+    cache: new InMemoryCache(),
+    manifest: mockApplicationManifest(),
+  });
+
+  using host = await mockMcpHost();
+  host.onCleanup(() => client.stop());
+
+  host.sendToolInput({ arguments: {} });
+  host.sendToolResult({
+    content: [],
+    structuredContent: {},
+  });
+
+  using _disabledAct = disableActEnvironment();
+  const { takeSnapshot } = await renderHookToSnapshotStream(
+    () => useToolMetadata(),
+    {
+      wrapper: ({ children }) => (
+        <Suspense>
+          <ApolloProvider client={client}>{children}</ApolloProvider>
+        </Suspense>
+      ),
+    }
+  );
+
+  await expect(takeSnapshot()).resolves.toBeNull();
+  await expect(takeSnapshot).not.toRerender();
+});

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

@@ -0,0 +1,50 @@
+import { expect, test } from "vitest";
+import { useToolName } from "../useToolName.js";
+import {
+  graphqlToolResult,
+  minimalHostContextWithToolName,
+  mockApplicationManifest,
+  mockMcpHost,
+  spyOnConsole,
+  stubOpenAiGlobals,
+} from "../../../../testing/internal/index.js";
+import { ApolloClient } from "../../../core/ApolloClient.js";
+import { InMemoryCache } from "@apollo/client";
+import {
+  disableActEnvironment,
+  renderHookToSnapshotStream,
+} from "@testing-library/react-render-stream";
+import { Suspense } from "react";
+import { ApolloProvider } from "../../../../react/ApolloProvider.js";
+
+test("returns the tool name from the MCP host", async () => {
+  using _ = spyOnConsole("debug");
+  stubOpenAiGlobals({ toolResponseMetadata: {} });
+  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/{react → openai/react}/hooks/__tests__/useWidgetState.test.tsx

@@ -4,7 +4,7 @@ import {
   renderHookToSnapshotStream,
 } from "@testing-library/react-render-stream";
 import { useWidgetState } from "../useWidgetState.js";
-import { stubOpenAiGlobals } from "../../../testing/internal/index.js";
+import { stubOpenAiGlobals } from "../../../../testing/internal/index.js";
 
 afterEach(() => {
   vi.unstubAllGlobals();

package/src/openai/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/openai/react/hooks/useApp.ts

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

package/src/{react → openai/react}/hooks/useOpenAiGlobal.ts

@@ -1,6 +1,6 @@
 import { useSyncExternalStore, useCallback } from "react";
-import type { SetGlobalsEvent, OpenAiGlobals } from "../../types/openai.js";
-import { SET_GLOBALS_EVENT_TYPE } from "../../types/openai.js";
+import type { SetGlobalsEvent, OpenAiGlobals } from "../../types.js";
+import { SET_GLOBALS_EVENT_TYPE } from "../../types.js";
 
 export function useOpenAiGlobal<K extends keyof OpenAiGlobals>(
   key: K

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

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

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

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

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

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

package/src/{react → openai/react}/hooks/useWidgetState.ts

@@ -1,6 +1,6 @@
 import type { SetStateAction } from "react";
 import { useCallback, useState } from "react";
-import type { UnknownObject } from "../../types/openai.js";
+import type { UnknownObject } from "../../types.js";
 import { useOpenAiGlobal } from "./useOpenAiGlobal.js";
 
 export function useWidgetState<T extends UnknownObject>(

package/src/openai/react/index.ts

@@ -0,0 +1,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 { useWidgetState } from "./hooks/useWidgetState.js";

package/src/{types/openai.ts → openai/types.ts}

@@ -1,19 +1,11 @@
-export type UnknownObject = Record<string, unknown>;
-
-declare global {
-  interface Window {
-    openai: API<any> & OpenAiGlobals;
-  }
+import type { ApolloMcpServerApps } from "../core/types";
 
-  interface WindowEventMap {
-    [SET_GLOBALS_EVENT_TYPE]: SetGlobalsEvent;
-  }
-}
+export type UnknownObject = Record<string, unknown>;
 
 export type OpenAiGlobals<
   ToolInput extends UnknownObject = UnknownObject,
-  ToolOutput extends UnknownObject = UnknownObject,
-  ToolResponseMetadata extends UnknownObject = UnknownObject,
+  ToolOutput = ApolloMcpServerApps.StructuredContent,
+  ToolResponseMetadata = UnknownObject,
   WidgetState extends UnknownObject = UnknownObject,
 > = {
   theme: Theme;
@@ -65,7 +57,11 @@ export class SetGlobalsEvent extends CustomEvent<{
 export type CallTool = (
   name: string,
   args: Record<string, unknown>
-) => Promise<any>;
+) => Promise<CallToolResponse>;
+
+export type CallToolResponse = {
+  result: string;
+};
 
 export type DisplayMode = "pip" | "inline" | "fullscreen";
 

package/src/react/ApolloProvider.tsx

@@ -1,8 +1,14 @@
 import type { ReactNode } from "react";
-import { useEffect, useState } from "react";
+import { use } from "react";
 import { ApolloProvider as BaseApolloProvider } from "@apollo/client/react";
-import type { ApolloClient } from "../core/ApolloClient.js";
-import { SET_GLOBALS_EVENT_TYPE } from "../types/openai.js";
+import type { ApolloClient as BaseApolloClient } from "@apollo/client";
+import { ApolloClient as OpenAiApolloClient } from "../openai/core/ApolloClient.js";
+import { ApolloClient as McpApolloClient } from "../mcp/core/ApolloClient.js";
+import type { ApolloClient as FallbackApolloClient } from "../core/ApolloClient.js";
+import { __DEV__ } from "@apollo/client/utilities/environment";
+import { aiClientSymbol, invariant } from "../utilities/index.js";
+
+type ApolloClient = OpenAiApolloClient | McpApolloClient | FallbackApolloClient;
 
 export declare namespace ApolloProvider {
   export interface Props {
@@ -11,33 +17,19 @@ export declare namespace ApolloProvider {
   }
 }
 
-export const ApolloProvider = ({ children, client }: ApolloProvider.Props) => {
-  const [hasPreloaded, setHasPreloaded] = useState(false);
-
-  // This is to prevent against a race condition. We don't know if window.openai will be available when this loads or if it will become available shortly after.
-  // So... we create the event listener and whenever it is available, then we can process the prefetch/tool data.
-  // In practice, this should be pretty much instant
-  useEffect(() => {
-    const prefetchData = async () => {
-      await client.prefetchData();
-      setHasPreloaded(true);
-      window.removeEventListener(SET_GLOBALS_EVENT_TYPE, prefetchData);
-    };
-
-    window.addEventListener(SET_GLOBALS_EVENT_TYPE, prefetchData, {
-      passive: true,
-    });
-
-    if (window.openai?.toolOutput) {
-      window.dispatchEvent(new CustomEvent(SET_GLOBALS_EVENT_TYPE));
-    }
+export function ApolloProvider({ children, client }: ApolloProvider.Props) {
+  if (__DEV__) {
+    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 () => {
-      window.removeEventListener(SET_GLOBALS_EVENT_TYPE, prefetchData);
-    };
-  }, []);
+  use(client.waitForInitialization());
 
-  return hasPreloaded ?
-      <BaseApolloProvider client={client}>{children}</BaseApolloProvider>
-    : null;
-};
+  return (
+    <BaseApolloProvider client={client as BaseApolloClient}>
+      {children}
+    </BaseApolloProvider>
+  );
+}

package/src/react/__tests__/ApolloProvider/mcp.test.tsx

@@ -0,0 +1,74 @@
+import { expect, test } from "vitest";
+import { Suspense } from "react";
+import { ApolloProvider } from "../../ApolloProvider.js";
+import { waitFor } from "@testing-library/react";
+import { ApolloClient } from "../../../mcp/core/ApolloClient.js";
+import { gql, InMemoryCache } from "@apollo/client";
+import { print } from "@apollo/client/utilities";
+import {
+  minimalHostContextWithToolName,
+  mockApplicationManifest,
+  mockMcpHost,
+  renderAsync,
+  spyOnConsole,
+} from "../../../testing/internal/index.js";
+
+test("writes to the cache as soon as toolOutput is available", async () => {
+  using _ = spyOnConsole("debug");
+
+  const toolName = "GreetingQuery";
+  using host = await mockMcpHost({
+    hostContext: minimalHostContextWithToolName(toolName),
+  });
+  host.onCleanup(() => client.stop());
+
+  const query = gql`
+    query GreetingQuery {
+      greeting
+    }
+  `;
+  const data = {
+    greeting: "hello",
+  };
+
+  const client = new ApolloClient({
+    cache: new InMemoryCache(),
+    manifest: mockApplicationManifest({
+      operations: [
+        {
+          id: "1",
+          name: "GreetingQuery",
+          body: print(query),
+          prefetch: true,
+          prefetchID: "__anonymous",
+          type: "query",
+          variables: {},
+          tools: [{ name: toolName, description: "Fetches a greeting" }],
+        },
+      ],
+    }),
+  });
+
+  await renderAsync(<ApolloProvider client={client} />, {
+    wrapper: ({ children }) => <Suspense>{children}</Suspense>,
+  });
+
+  await expect(
+    waitFor(() => expect(client.extract()).not.toEqual({}))
+  ).rejects.toThrow();
+
+  host.sendToolInput({ arguments: {} });
+  host.sendToolResult({
+    content: [{ type: "text", text: JSON.stringify({ result: { data } }) }],
+    structuredContent: { result: { data } },
+  });
+
+  await waitFor(() => {
+    expect(client.extract()).toEqual({
+      ROOT_QUERY: {
+        __typename: "Query",
+        greeting: "hello",
+      },
+    });
+  });
+});

package/src/react/__tests__/ApolloProvider/openai.test.tsx

@@ -0,0 +1,146 @@
+import { expect, test } from "vitest";
+import { Suspense } from "react";
+import { ApolloProvider } from "../../ApolloProvider.js";
+import { waitFor } from "@testing-library/react";
+import { ApolloClient } from "../../../openai/core/ApolloClient.js";
+import { gql, InMemoryCache } from "@apollo/client";
+import { print } from "@apollo/client/utilities";
+import {
+  minimalHostContextWithToolName,
+  mockApplicationManifest,
+  mockMcpHost,
+  renderAsync,
+  spyOnConsole,
+  stubOpenAiGlobals,
+} from "../../../testing/internal/index.js";
+
+test("writes data to the cache when immediately available", async () => {
+  stubOpenAiGlobals();
+  using _ = spyOnConsole("debug");
+
+  const query = gql`
+    query GreetingQuery {
+      greeting
+    }
+  `;
+
+  const data = {
+    greeting: "hello",
+  };
+
+  const client = new ApolloClient({
+    cache: new InMemoryCache(),
+    manifest: mockApplicationManifest({
+      operations: [
+        {
+          id: "1",
+          name: "GreetingQuery",
+          body: print(query),
+          prefetch: true,
+          prefetchID: "__anonymous",
+          type: "query",
+          variables: {},
+          tools: [],
+        },
+      ],
+    }),
+  });
+
+  using host = await mockMcpHost({
+    hostContext: minimalHostContextWithToolName("GreetingQuery"),
+  });
+  host.onCleanup(() => client.stop());
+
+  host.sendToolInput({ arguments: {} });
+  host.sendToolResult({
+    content: [],
+    structuredContent: {
+      result: {
+        data: null,
+      },
+      prefetch: {
+        __anonymous: { data },
+      },
+    },
+  });
+
+  await renderAsync(<ApolloProvider client={client} />, {
+    wrapper: ({ children }) => <Suspense>{children}</Suspense>,
+  });
+
+  await waitFor(() => {
+    expect(client.extract()).toEqual({
+      ROOT_QUERY: {
+        __typename: "Query",
+        greeting: "hello",
+      },
+    });
+  });
+});
+
+test("writes to the cache as soon as toolOutput is available", async () => {
+  stubOpenAiGlobals();
+  using _ = spyOnConsole("debug");
+
+  stubOpenAiGlobals();
+
+  const query = gql`
+    query GreetingQuery {
+      greeting
+    }
+  `;
+  const data = {
+    greeting: "hello",
+  };
+
+  const client = new ApolloClient({
+    cache: new InMemoryCache(),
+    manifest: mockApplicationManifest({
+      operations: [
+        {
+          id: "1",
+          name: "GreetingQuery",
+          body: print(query),
+          prefetch: true,
+          prefetchID: "__anonymous",
+          type: "query",
+          variables: {},
+          tools: [],
+        },
+      ],
+    }),
+  });
+
+  using host = await mockMcpHost({
+    hostContext: minimalHostContextWithToolName("GreetingQuery"),
+  });
+
+  host.onCleanup(() => client.stop());
+  host.sendToolInput({ arguments: {} });
+
+  await renderAsync(<ApolloProvider client={client} />, {
+    wrapper: ({ children }) => <Suspense>{children}</Suspense>,
+  });
+
+  await expect(
+    waitFor(() => expect(client.extract()).not.toEqual({}))
+  ).rejects.toThrow();
+
+  host.sendToolResult({
+    content: [],
+    structuredContent: {
+      prefetch: {
+        __anonymous: { data },
+      },
+    },
+  });
+
+  await waitFor(() => {
+    expect(client.extract()).toEqual({
+      ROOT_QUERY: {
+        __typename: "Query",
+        greeting: "hello",
+      },
+    });
+  });
+});

package/src/react/index.mcp.ts

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

package/src/react/index.openai.ts

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

package/src/react/index.ts

@@ -0,0 +1,19 @@
+import { missingHook } from "./missingHook.js";
+
+export { ApolloProvider } from "./ApolloProvider.js";
+
+// Use `mcp` related types since these are the most common between the two
+// targets
+export const useApp =
+  missingHook<typeof import("./index.mcp.js").useApp>("useApp");
+
+export const useToolInput =
+  missingHook<typeof import("./index.mcp.js").useToolInput>("useToolInput");
+
+export const useToolMetadata =
+  missingHook<typeof import("./index.mcp.js").useToolMetadata>(
+    "useToolMetadata"
+  );
+
+export const useToolName =
+  missingHook<typeof import("./index.mcp.js").useToolName>("useToolName");

package/src/react/missingHook.ts

@@ -0,0 +1,9 @@
+export function missingHook<HookFn extends (...args: any[]) => any>(
+  name: string
+) {
+  return (() => {
+    throw new Error(
+      `Cannot use the '${name}' hook without export conditions. Please set export conditions or import from the \`/openai\` or \`/mcp\` subpath directly.`
+    );
+  }) as unknown as HookFn;
+}

package/src/testing/internal/index.ts

@@ -1,2 +1,10 @@
 export { dispatchStateChange } from "./openai/dispatchStateChange.js";
+export { graphqlToolResult } from "./mcp/graphqlToolResult.js";
+export { mockApplicationManifest } from "./utilities/mockApplicationManifest.js";
 export { stubOpenAiGlobals } from "./openai/stubOpenAiGlobals.js";
+export { minimalHostContextWithToolName } from "./mcp/minimalHostContextWithToolName.js";
+export { mockMcpHost } from "./mcp/mockMcpHost.js";
+export { ObservableStream } from "./utilities/ObservableStream.js";
+export { renderAsync } from "./react/renderAsync.js";
+export { spyOnConsole } from "./utilities/spyOnConsole.js";
+export { wait } from "./utilities/wait.js";

package/src/testing/internal/matchers/index.ts

@@ -1 +1,3 @@
+import "./toComplete.js";
+import "./toEmitValue.js";
 import "./toRerender.js";

package/src/testing/internal/matchers/toComplete.ts

@@ -0,0 +1,46 @@
+import { expect } from "vitest";
+
+import type {
+  ObservableStream,
+  TakeOptions,
+} from "../utilities/ObservableStream.js";
+import { EventMismatchError } from "../utilities/ObservableStream.js";
+
+expect.extend({
+  async toComplete(actual, options?: TakeOptions) {
+    const stream = actual as ObservableStream<any>;
+    const hint = this.utils.matcherHint("toComplete", "stream", "");
+
+    try {
+      await stream.takeComplete(options);
+
+      return {
+        pass: true,
+        message: () => {
+          return hint + "\n\nExpected stream not to complete but it did.";
+        },
+      };
+    } catch (error) {
+      if (
+        error instanceof Error &&
+        error.message === "Timeout waiting for next event"
+      ) {
+        return {
+          pass: false,
+          message: () =>
+            hint + "\n\nExpected stream to complete but it did not.",
+        };
+      } else if (EventMismatchError.is(error)) {
+        return {
+          pass: false,
+          message: () =>
+            hint +
+            "\n\n" +
+            this.utils.printDiffOrStringify(error.actual, error.expected),
+        };
+      } else {
+        throw error;
+      }
+    }
+  },
+});