@miragon/mcp-toolkit-ui 2.6.0 → 2.8.0

package/src/lib/tool-query.ts

@@ -0,0 +1,253 @@
+/**
+ * The plain logic behind `useToolQuery`, `useViewToolQuery`, `useViewData`
+ * and `useToolMutation` (toolkit#176 K18): the effective cache key and its
+ * collision guard, the opt-in retry policy for tool errors, the TanStack
+ * options a tool query runs with, and per-call mutation callbacks that fire
+ * for EVERY call. No React, no `mcp-use/react`, so it is unit-tested without
+ * a renderer.
+ */
+import {
+  hashKey,
+  type MutateOptions,
+  type MutationFunctionContext,
+  type QueryClient,
+  type UseQueryOptions,
+} from "@tanstack/react-query"
+import { isToolError, parseToolResult } from "./parse-tool-result.js"
+
+/** A `(name, args)` tool caller — the shape of `useCallTool()` and `HostBridge.callTool`. */
+export type ToolCaller = (name: string, args: object) => Promise<unknown>
+
+/**
+ * TanStack query options a tool query passes through: everything except the
+ * key and the fetcher, which the hooks build from the tool and its args
+ * (`retry`, `retryDelay`, `staleTime`, `placeholderData` / `keepPreviousData`,
+ * `initialData` + `initialDataUpdatedAt`, `enabled`, `select`, …).
+ */
+export type UseToolQueryOptions<TData = unknown, TSelected = TData> = Omit<
+  UseQueryOptions<TData, Error, TSelected>,
+  "queryKey" | "queryFn"
+>
+
+/**
+ * Neither an `AppQueryProvider` `callTool` nor a `HostBridgeProvider` is in
+ * reach, so the query cannot call any tool. Never retried: nothing changes
+ * until the mount does.
+ */
+export class NoToolTransportError extends Error {
+  constructor(toolName: string) {
+    super(
+      `[mcp-toolkit] cannot call "${toolName}": no tool transport. Mount the widget under ` +
+        "AppQueryProvider with callTool, or a HostBridgeProvider.",
+    )
+    this.name = "NoToolTransportError"
+  }
+}
+
+/** Retries {@link shouldRetryToolCall} grants a failed transport (TanStack's own default count). */
+export const TOOL_QUERY_RETRIES = 3
+
+/**
+ * A retry policy for tool queries: up to {@link TOOL_QUERY_RETRIES} retries
+ * for a failed transport, none for a tool's own `isError` answer
+ * (`isToolError`) or a missing transport — repeating a deterministic answer
+ * only repeats it and holds the error back about 7 s.
+ *
+ * Opt-in in 2.x (the hooks keep the client's `retry`, TanStack's 3 by
+ * default, because a tool may also answer a transient failure with
+ * `isError`): pass `retry: shouldRetryToolCall` to a hook, or set it as the
+ * client's default (`queryClient.setDefaultOptions`, or an `AppQueryProvider`
+ * `client` of your own).
+ */
+export function shouldRetryToolCall(failureCount: number, error: unknown): boolean {
+  if (isToolError(error) || error instanceof NoToolTransportError) return false
+  return failureCount < TOOL_QUERY_RETRIES
+}
+
+/**
+ * The effective cache key of a tool query (toolkit#176 K18).
+ *
+ * - 2.x default: the caller's key, then the args (`[...queryKey, args]`), so
+ *   two callers sharing a key with different args get separate entries.
+ * - `toolNameInKey`: the caller's key, the tool name, then the args
+ *   (`[...queryKey, toolName, args]`), so two tools never share an entry,
+ *   whatever key their callers pass. Opt-in in 2.x (`AppQueryProvider` /
+ *   `McpAppView` `toolNameInQueryKey`), the default from the next major.
+ *
+ * Build exact keys (`getQueryData`, `setQueryData`, `cancelQueries` with
+ * `exact: true`) with this helper, or with `useToolQueryKey`, which reads the
+ * provider's setting, rather than by hand: code that goes through them keeps
+ * working when the default changes. Prefix invalidation by `queryKey`
+ * matches either way.
+ */
+export function toolQueryKey(
+  queryKey: readonly unknown[],
+  toolName: string,
+  args: object,
+  toolNameInKey = false,
+): unknown[] {
+  return toolNameInKey ? [...queryKey, toolName, args] : [...queryKey, args]
+}
+
+/**
+ * The tool that last claimed each cache key, per client — the guard against
+ * two tools sharing one entry (their answers would overwrite each other).
+ */
+const keyOwners = new WeakMap<QueryClient, Map<string, string>>()
+
+/**
+ * Claim the cache key `keyId` (`hashKey` of a {@link toolQueryKey}) for
+ * `toolName` on `client`. Returns the tool that held it
+ * before when that is ANOTHER tool — a collision the hooks report — else
+ * `null`. The latest claim wins, so each collision is reported once per
+ * switch.
+ */
+export function claimToolQueryKey(
+  client: QueryClient,
+  keyId: string,
+  toolName: string,
+): string | null {
+  const owners = keyOwners.get(client) ?? new Map<string, string>()
+  keyOwners.set(client, owners)
+  const previous = owners.get(keyId)
+  owners.set(keyId, toolName)
+  return previous !== undefined && previous !== toolName ? previous : null
+}
+
+/** The console line for a {@link claimToolQueryKey} collision. */
+export function toolKeyCollisionMessage(keyId: string, toolName: string, previous: string): string {
+  return (
+    `[mcp-toolkit] the query key ${keyId} is used by "${previous}" and "${toolName}": ` +
+    "both share one cache entry and overwrite each other's answer. Give each tool its own queryKey."
+  )
+}
+
+/** How {@link toolQueryOptions} decodes a result and builds the key. */
+export interface ToolQueryBuild<TData> {
+  /** Decodes the tool result. Default: `parseToolResult` (structured-first). */
+  decode?: (result: unknown) => TData
+  /** Put the tool name into the key (see {@link toolQueryKey}). Default `false`. */
+  toolNameInKey?: boolean
+}
+
+/**
+ * The TanStack options of a tool query: the caller's options, the
+ * {@link toolQueryKey} and a fetcher that calls `toolName` and decodes the
+ * result. `retry` stays the caller's, else the client's default; only a
+ * missing transport is never retried. Without a transport the fetcher fails
+ * with {@link NoToolTransportError} instead of leaving the query idle.
+ */
+export function toolQueryOptions<TData, TSelected = TData>(
+  queryKey: readonly unknown[],
+  toolName: string,
+  args: object,
+  callTool: ToolCaller | undefined,
+  opts?: UseToolQueryOptions<TData, TSelected>,
+  build: ToolQueryBuild<TData> = {},
+): UseQueryOptions<TData, Error, TSelected> {
+  const decode = build.decode ?? ((result: unknown) => parseToolResult<TData>(result))
+  return {
+    ...opts,
+    // An omitted `retry` must stay omitted: an explicit `undefined` would
+    // shadow the client's default.
+    ...(callTool ? null : { retry: false }),
+    queryKey: toolQueryKey(queryKey, toolName, args, build.toolNameInKey),
+    queryFn: async () => {
+      if (!callTool) throw new NoToolTransportError(toolName)
+      return decode(await callTool(toolName, args))
+    },
+  }
+}
+
+/**
+ * The content of the payload last seeded per query key, per client: a
+ * remounted or re-rendered OLD payload must not roll a refetched answer
+ * back, while a payload not seeded before is a new delivery and wins.
+ */
+const lastSeeded = new WeakMap<QueryClient, Map<string, string>>()
+
+function contentOf(value: unknown): string | undefined {
+  try {
+    return JSON.stringify(value)
+  } catch {
+    return undefined
+  }
+}
+
+/**
+ * Seed `queryKey` with host-delivered `seed`: written when the entry is
+ * empty, or when the seed is a payload this client has not seeded for the
+ * key before (a new delivery). Returns whether the cache was written.
+ */
+export function seedToolQuery(
+  client: QueryClient,
+  queryKey: readonly unknown[],
+  seed: unknown,
+  updatedAt?: number,
+): boolean {
+  const keyId = hashKey(queryKey)
+  const seeded = lastSeeded.get(client) ?? new Map<string, string>()
+  lastSeeded.set(client, seeded)
+  const content = contentOf(seed)
+  const delivered = content === undefined || seeded.get(keyId) !== content
+  if (content !== undefined) seeded.set(keyId, content)
+  if (!delivered && client.getQueryData(queryKey) !== undefined) return false
+  client.setQueryData(queryKey, seed, updatedAt === undefined ? undefined : { updatedAt })
+  return true
+}
+
+/** Callbacks of one `mutate` call, typed as TanStack's per-call options. */
+export type PerCallCallbacks<TData, TVariables> = MutateOptions<TData, Error, TVariables, unknown>
+
+function settle(run: () => void): void {
+  try {
+    run()
+  } catch (err) {
+    // Surface a throwing callback as an unhandled rejection, as TanStack
+    // does, without skipping the callbacks after it.
+    void Promise.reject(err instanceof Error ? err : new Error(String(err)))
+  }
+}
+
+/**
+ * Fire one call's per-call `onSuccess` / `onError` / `onSettled` once `call`
+ * settles — for EVERY call. TanStack's own `mutate` keeps only the latest
+ * call's callbacks, so two quick calls (retrying two incidents in a row)
+ * would confirm only the second. `context()` is read when the call settles:
+ * `null` (the component unmounted) skips the callbacks, as TanStack does.
+ * A rejected `call` is handled here, never left unhandled.
+ */
+export function settleEachCall<TData, TVariables>(
+  call: Promise<TData>,
+  variables: TVariables,
+  callbacks: PerCallCallbacks<TData, TVariables> | undefined,
+  context: () => MutationFunctionContext | null,
+): void {
+  call.then(
+    (data) => {
+      const ctx = callbacks ? context() : null
+      if (!callbacks || !ctx) return
+      settle(() => callbacks.onSuccess?.(data, variables, undefined, ctx))
+      settle(() => callbacks.onSettled?.(data, null, variables, undefined, ctx))
+    },
+    (error: Error) => {
+      const ctx = callbacks ? context() : null
+      if (!callbacks || !ctx) return
+      settle(() => callbacks.onError?.(error, variables, undefined, ctx))
+      settle(() => callbacks.onSettled?.(undefined, error, variables, undefined, ctx))
+    },
+  )
+}
+
+/**
+ * Invalidate the query keys (prefixes) a mutation names: matching ACTIVE
+ * queries refetch now, inactive ones are only marked stale and refetch when
+ * they are next used — never a refetch of every inactive cached page.
+ */
+export async function invalidateToolQueries(
+  client: QueryClient,
+  keys: readonly (readonly unknown[])[] | undefined,
+): Promise<void> {
+  if (!keys) return
+  await Promise.all(keys.map((queryKey) => client.invalidateQueries({ queryKey })))
+}

package/src/lib/tool-result-observer.test.ts

@@ -0,0 +1,296 @@
+import { describe, expect, it, vi } from "vitest"
+import {
+  classifyHostMessage,
+  classifyToolResult,
+  createToolResultObserver,
+  EMPTY_TOOL_OBSERVATION,
+  installToolResultObserver,
+  reduceToolObservation,
+  type HostLifecycleMessage,
+  type ToolObservation,
+} from "./tool-result-observer.js"
+
+/**
+ * The passive observer behind the stripped-result recovery (toolkit#176 K13,
+ * N21): it must tell a stripped result from a structured one and from "not
+ * yet", read only the COMPLETE tool input, and never act on foreign traffic.
+ */
+
+const rpc = (method: string, params?: Record<string, unknown>) => ({
+  jsonrpc: "2.0",
+  method,
+  ...(params ? { params } : {}),
+})
+
+/** The host's answer to `ui/initialize`, as ext-apps sends it. */
+const HANDSHAKE = {
+  jsonrpc: "2.0",
+  id: 1,
+  result: {
+    protocolVersion: "2026-01-26",
+    hostInfo: { name: "host", version: "1" },
+    hostCapabilities: { serverTools: {} },
+    hostContext: { toolInfo: { tool: { name: "render-view" } } },
+  },
+}
+
+describe("classifyToolResult (mirrors the mcp-use 2.x latch)", () => {
+  it("an isError result is an error, whatever else it carries", () => {
+    expect(classifyToolResult({ isError: true })).toBe("error")
+    expect(classifyToolResult({ isError: true, structuredContent: { a: 1 } })).toBe("error")
+  })
+
+  it("a result with structuredContent — even null — is structured", () => {
+    expect(classifyToolResult({ structuredContent: { a: 1 } })).toBe("structured")
+    expect(classifyToolResult({ structuredContent: null })).toBe("structured")
+  })
+
+  it("a content-only result is the stripped one claude.ai sends", () => {
+    expect(classifyToolResult({ content: [{ type: "text", text: "Rendered." }] })).toBe("stripped")
+    expect(classifyToolResult({})).toBe("stripped")
+    expect(classifyToolResult({ isError: false })).toBe("stripped")
+  })
+})
+
+describe("classifyHostMessage", () => {
+  it("recognises the host's ui/initialize answer by its capabilities or protocol version", () => {
+    expect(classifyHostMessage(HANDSHAKE)).toEqual({ kind: "handshake" })
+    expect(
+      classifyHostMessage({ jsonrpc: "2.0", id: 7, result: { protocolVersion: "x" } }),
+    ).toEqual({ kind: "handshake" })
+    expect(
+      classifyHostMessage({ jsonrpc: "2.0", id: 7, result: { hostCapabilities: {} } }),
+    ).toEqual({ kind: "handshake" })
+  })
+
+  it("ignores every other response (a tools/call answer, an error, a bare result)", () => {
+    expect(
+      classifyHostMessage({ jsonrpc: "2.0", id: 2, result: { structuredContent: {} } }),
+    ).toBeNull()
+    expect(classifyHostMessage({ jsonrpc: "2.0", id: 2, error: { code: 1 } })).toBeNull()
+    expect(classifyHostMessage({ jsonrpc: "2.0", id: 2, result: "x" })).toBeNull()
+    expect(
+      classifyHostMessage({ jsonrpc: "2.0", result: { protocolVersion: "x" } }),
+      "a response needs an id",
+    ).toBeNull()
+    expect(
+      classifyHostMessage({ jsonrpc: "2.0", id: 3, result: { protocolVersion: 1 } }),
+    ).toBeNull()
+  })
+
+  it("reads the complete tool input, never the partial stream", () => {
+    expect(
+      classifyHostMessage(rpc("ui/notifications/tool-input", { arguments: { id: 1 } })),
+    ).toEqual({ kind: "tool-input", arguments: { id: 1 } })
+    expect(
+      classifyHostMessage(rpc("ui/notifications/tool-input-partial", { arguments: { id: 1 } })),
+    ).toBeNull()
+  })
+
+  it("a tool input without usable arguments is an empty argument set", () => {
+    expect(classifyHostMessage(rpc("ui/notifications/tool-input"))).toEqual({
+      kind: "tool-input",
+      arguments: {},
+    })
+    expect(
+      classifyHostMessage(rpc("ui/notifications/tool-input", { arguments: ["not", "a", "map"] })),
+    ).toEqual({ kind: "tool-input", arguments: {} })
+  })
+
+  it("classifies the tool result and the cancellation", () => {
+    expect(classifyHostMessage(rpc("ui/notifications/tool-result", { content: [] }))).toEqual({
+      kind: "tool-result",
+      result: "stripped",
+    })
+    expect(
+      classifyHostMessage(rpc("ui/notifications/tool-result", { structuredContent: {} })),
+    ).toEqual({ kind: "tool-result", result: "structured" })
+    expect(classifyHostMessage(rpc("ui/notifications/tool-result"))).toEqual({
+      kind: "tool-result",
+      result: "stripped",
+    })
+    expect(classifyHostMessage(rpc("ui/notifications/tool-cancelled"))).toEqual({
+      kind: "tool-cancelled",
+    })
+  })
+
+  it("ignores foreign traffic and other methods", () => {
+    expect(classifyHostMessage(null)).toBeNull()
+    expect(classifyHostMessage("ui/notifications/tool-result")).toBeNull()
+    expect(classifyHostMessage([HANDSHAKE])).toBeNull()
+    expect(classifyHostMessage({ method: "ui/notifications/tool-result" })).toBeNull()
+    expect(classifyHostMessage({ ...HANDSHAKE, jsonrpc: "1.0" })).toBeNull()
+    expect(classifyHostMessage(rpc("ui/notifications/host-context-changed", {}))).toBeNull()
+    expect(classifyHostMessage(rpc("ui/resource-teardown", {}))).toBeNull()
+  })
+})
+
+describe("reduceToolObservation", () => {
+  const fold = (...messages: HostLifecycleMessage[]): ToolObservation =>
+    messages.reduce(reduceToolObservation, EMPTY_TOOL_OBSERVATION)
+
+  it("starts empty", () => {
+    expect(EMPTY_TOOL_OBSERVATION).toEqual({
+      sawHandshake: false,
+      completeInput: undefined,
+      result: "none",
+      cancelled: false,
+    })
+  })
+
+  it("records the handshake, the latest complete input and the result", () => {
+    expect(
+      fold(
+        { kind: "handshake" },
+        { kind: "tool-input", arguments: { a: 1 } },
+        { kind: "tool-input", arguments: { a: 2 } },
+        { kind: "tool-result", result: "stripped" },
+      ),
+    ).toEqual({ sawHandshake: true, completeInput: { a: 2 }, result: "stripped", cancelled: false })
+  })
+
+  it("settles once, like the mcp-use latch: nothing changes a delivered call", () => {
+    const settled = fold(
+      { kind: "tool-input", arguments: { a: 1 } },
+      { kind: "tool-result", result: "structured" },
+    )
+    expect(reduceToolObservation(settled, { kind: "tool-input", arguments: { a: 9 } })).toBe(
+      settled,
+    )
+    expect(reduceToolObservation(settled, { kind: "tool-result", result: "stripped" })).toBe(
+      settled,
+    )
+    expect(reduceToolObservation(settled, { kind: "tool-cancelled" })).toBe(settled)
+  })
+
+  it("a cancellation settles the call too", () => {
+    const cancelled = fold({ kind: "tool-cancelled" })
+    expect(cancelled.cancelled).toBe(true)
+    expect(reduceToolObservation(cancelled, { kind: "tool-result", result: "stripped" })).toBe(
+      cancelled,
+    )
+  })
+
+  it("a repeated handshake is the same state; a late one still counts after a result", () => {
+    const shaken = fold({ kind: "handshake" })
+    expect(reduceToolObservation(shaken, { kind: "handshake" })).toBe(shaken)
+    const late = fold({ kind: "tool-result", result: "stripped" }, { kind: "handshake" })
+    expect(late).toMatchObject({ sawHandshake: true, result: "stripped" })
+  })
+})
+
+describe("createToolResultObserver", () => {
+  const post = (target: EventTarget, data: unknown, source: unknown = null) => {
+    const event = new MessageEvent("message", { data })
+    Object.defineProperty(event, "source", { value: source })
+    target.dispatchEvent(event)
+  }
+
+  it("folds the messages it accepts and notifies only on a change", () => {
+    const target = new EventTarget()
+    const observer = createToolResultObserver(target)
+    const listener = vi.fn()
+    observer.subscribe(listener)
+    post(target, HANDSHAKE)
+    post(target, rpc("ui/notifications/tool-input", { arguments: { id: 7 } }))
+    post(target, rpc("ui/notifications/tool-input-partial", { arguments: { id: 8 } }))
+    post(target, { unrelated: true })
+    post(target, rpc("ui/notifications/tool-result", { content: [] }))
+    post(target, rpc("ui/notifications/tool-result", { structuredContent: {} }))
+    expect(observer.getSnapshot()).toEqual({
+      sawHandshake: true,
+      completeInput: { id: 7 },
+      result: "stripped",
+      cancelled: false,
+    })
+    expect(listener).toHaveBeenCalledTimes(3)
+  })
+
+  it("drops what the filter refuses (another frame's traffic)", () => {
+    const target = new EventTarget()
+    const host = { name: "host" }
+    const observer = createToolResultObserver(target, (event) => event.source === host)
+    post(target, HANDSHAKE, { name: "intruder" })
+    post(target, rpc("ui/notifications/tool-result", { content: [] }), null)
+    expect(observer.getSnapshot()).toBe(EMPTY_TOOL_OBSERVATION)
+    post(target, HANDSHAKE, host)
+    expect(observer.getSnapshot().sawHandshake).toBe(true)
+  })
+
+  it("unsubscribes a listener, and dispose stops listening but keeps the last snapshot", () => {
+    const target = new EventTarget()
+    const observer = createToolResultObserver(target)
+    const kept = vi.fn()
+    const dropped = vi.fn()
+    observer.subscribe(kept)
+    const unsubscribe = observer.subscribe(dropped)
+    unsubscribe()
+    post(target, HANDSHAKE)
+    expect(kept).toHaveBeenCalledTimes(1)
+    expect(dropped).not.toHaveBeenCalled()
+    observer.dispose()
+    post(target, rpc("ui/notifications/tool-result", { content: [] }))
+    expect(observer.getSnapshot()).toMatchObject({ sawHandshake: true, result: "none" })
+    expect(kept).toHaveBeenCalledTimes(1)
+  })
+})
+
+describe("installToolResultObserver", () => {
+  it("is null outside a framed browser document (no host to observe)", () => {
+    expect(installToolResultObserver()).toBeNull()
+  })
+
+  it("is null in a top-level window (no host frames the view)", async () => {
+    vi.resetModules()
+    const target = new EventTarget()
+    vi.stubGlobal("window", Object.assign(target, { parent: target }))
+    try {
+      const fresh = await import("./tool-result-observer.js")
+      expect(fresh.installToolResultObserver()).toBeNull()
+    } finally {
+      vi.unstubAllGlobals()
+      vi.resetModules()
+    }
+  })
+
+  it("listens from the moment the bundle evaluates, before any explicit install", async () => {
+    vi.resetModules()
+    const target = new EventTarget()
+    const parent = { name: "host" }
+    vi.stubGlobal("window", Object.assign(target, { parent }))
+    try {
+      const fresh = await import("./tool-result-observer.js")
+      post(target, HANDSHAKE, parent)
+      expect(fresh.installToolResultObserver()?.getSnapshot().sawHandshake).toBe(true)
+    } finally {
+      vi.unstubAllGlobals()
+      vi.resetModules()
+    }
+  })
+
+  it("installs once in a framed document and listens to the parent only", async () => {
+    vi.resetModules()
+    const target = new EventTarget()
+    const parent = { name: "host" }
+    vi.stubGlobal("window", Object.assign(target, { parent }))
+    try {
+      const fresh = await import("./tool-result-observer.js")
+      const observer = fresh.installToolResultObserver()
+      expect(observer).not.toBeNull()
+      expect(fresh.installToolResultObserver(), "idempotent").toBe(observer)
+      post(target, HANDSHAKE, { name: "other" })
+      expect(observer?.getSnapshot().sawHandshake).toBe(false)
+      post(target, HANDSHAKE, parent)
+      expect(observer?.getSnapshot().sawHandshake).toBe(true)
+    } finally {
+      vi.unstubAllGlobals()
+      vi.resetModules()
+    }
+  })
+
+  function post(target: EventTarget, data: unknown, source: unknown) {
+    const event = new MessageEvent("message", { data })
+    Object.defineProperty(event, "source", { value: source })
+    target.dispatchEvent(event)
+  }
+})