@miragon/mcp-toolkit-ui 2.6.0 → 2.8.0

package/src/hooks/use-tool-query.ts

@@ -1,63 +1,186 @@
-import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query"
-import { useCallTool } from "../providers/query-provider.js"
+import { useCallback, useEffect, useRef } from "react"
+import {
+  hashKey,
+  useMutation,
+  useQuery,
+  useQueryClient,
+  type MutationFunctionContext,
+  type UseMutationOptions,
+  type UseMutationResult,
+  type UseQueryResult,
+} from "@tanstack/react-query"
+import { useCallTool, useToolNameInQueryKey } from "../providers/query-provider.js"
 import { parseToolResult } from "../lib/parse-tool-result.js"
+import {
+  claimToolQueryKey,
+  invalidateToolQueries,
+  NoToolTransportError,
+  settleEachCall,
+  toolKeyCollisionMessage,
+  toolQueryKey,
+  toolQueryOptions,
+  type PerCallCallbacks,
+  type UseToolQueryOptions,
+} from "../lib/tool-query.js"
 
-export interface UseToolQueryOptions<TData = unknown, TSelected = TData> {
-  enabled?: boolean
-  select?: (data: TData) => TSelected
+export type { UseToolQueryOptions } from "../lib/tool-query.js"
+
+/**
+ * The exact cache key `useToolQuery`, `useViewToolQuery` and `useViewData`
+ * use for this tool and args under the nearest `AppQueryProvider`:
+ * `toolQueryKey` with the provider's `toolNameInQueryKey` setting. For
+ * `getQueryData`, `setQueryData` or `cancelQueries({ exact: true })`; prefix
+ * invalidation needs only the caller's `queryKey`. A new array per render.
+ */
+export function useToolQueryKey(
+  queryKey: readonly unknown[],
+  toolName: string,
+  args: object,
+): unknown[] {
+  return toolQueryKey(queryKey, toolName, args, useToolNameInQueryKey())
+}
+
+/**
+ * Warns when two tools use one cache key (toolkit#176 K18): their answers
+ * would overwrite each other. Only the 2.x key can collide: with the tool
+ * name in the key (`toolNameInQueryKey`), two tools never share one.
+ */
+export function useToolKeyGuard(
+  queryKey: readonly unknown[],
+  toolName: string,
+  args: object,
+): void {
+  const client = useQueryClient()
+  const keyId = hashKey(useToolQueryKey(queryKey, toolName, args))
+  useEffect(() => {
+    const previous = claimToolQueryKey(client, keyId, toolName)
+    if (previous !== null) console.warn(toolKeyCollisionMessage(keyId, toolName, previous))
+  }, [client, keyId, toolName])
+}
+
+/**
+ * Fetches a query again once a tool transport arrives, when it failed for
+ * want of one ({@link NoToolTransportError}): a `callTool` that resolves
+ * after the first render. TanStack refetches on a new key or `enabled`,
+ * never on a new fetcher, so the query would keep the error until a remount.
+ */
+export function useRefetchOnTransport(
+  callTool: unknown,
+  query: Pick<UseQueryResult<unknown, Error>, "error" | "refetch">,
+): void {
+  const { error, refetch } = query
+  useEffect(() => {
+    if (callTool && error instanceof NoToolTransportError) void refetch()
+  }, [callTool, error, refetch])
 }
 
 /**
  * Loads data from a widget tool via TanStack Query.
  *
- * `args` is appended to the effective query key so two callers that share
- * the supplied `queryKey` but pass different `args` get separate cache
- * entries (otherwise the second caller silently reads stale data from the
- * first — the classic queryKey-vs-args trap). `select` transforms data on
- * each read without mutating the cache.
+ * The effective cache key is `[...queryKey, args]` (`toolQueryKey`): two
+ * callers that share `queryKey` but pass different `args` get separate
+ * entries. Two tools must not share a key and args; the hook warns when they
+ * do. Under an `AppQueryProvider` with `toolNameInQueryKey` (opt-in in 2.x,
+ * the default from the next major) the key is `[...queryKey, toolName,
+ * args]` and two tools never collide; `useToolQueryKey` returns the key in
+ * force. `opts` passes TanStack's query options through (`retry`,
+ * `retryDelay`, `staleTime`, `placeholderData` / `keepPreviousData`,
+ * `initialData` + `initialDataUpdatedAt`, `enabled`, `select`, …), and the
+ * result exposes `refetch`, `failureCount` and `failureReason`. `retry`
+ * defaults to the client's; `retry: shouldRetryToolCall` stops retrying a
+ * tool's own `isError` answer.
+ *
+ * The tool goes through `AppQueryProvider`'s `callTool`, else the nearest
+ * `HostBridgeProvider`'s bridge; with neither, the query fails with a
+ * "no tool transport" error instead of idling, and fetches once one arrives.
  */
 export function useToolQuery<TData = unknown, TSelected = TData>(
-  queryKey: unknown[],
+  queryKey: readonly unknown[],
   toolName: string,
   args: object,
   opts?: UseToolQueryOptions<TData, TSelected>,
-) {
+): UseQueryResult<TSelected, Error> {
   const callTool = useCallTool()
+  const toolNameInKey = useToolNameInQueryKey()
+  useToolKeyGuard(queryKey, toolName, args)
+  const query = useQuery<TData, Error, TSelected>(
+    toolQueryOptions<TData, TSelected>(queryKey, toolName, args, callTool, opts, {
+      toolNameInKey,
+    }),
+  )
+  useRefetchOnTransport(callTool, query)
+  return query
+}
 
-  return useQuery<TData, Error, TSelected>({
-    queryKey: [...queryKey, args],
-    queryFn: async () => {
-      if (!callTool) throw new Error("callTool not available")
-      const result = await callTool(toolName, args)
-      return parseToolResult<TData>(result)
-    },
-    select: opts?.select,
-    enabled: !!callTool && (opts?.enabled ?? true),
-  })
+/** Options of {@link useToolMutation}: TanStack's mutation options plus `invalidateKeys`. */
+export type UseToolMutationOptions<TData = unknown> = Omit<
+  UseMutationOptions<TData, Error, Record<string, unknown>, unknown>,
+  "mutationFn" | "onMutate"
+> & {
+  /** Query keys (prefixes) to invalidate after a successful call. */
+  invalidateKeys?: unknown[][]
 }
 
 /**
- * Executes a mutation via a widget tool and invalidates configured query keys.
+ * Executes a mutation via a widget tool and invalidates `invalidateKeys`
+ * after each successful call (`invalidateQueries`: active queries refetch,
+ * inactive ones are marked stale and refetch when they are next used).
+ *
+ * `mutate(args, { onSuccess, onError, onSettled })` fires its per-call
+ * callbacks for EVERY call while the component is mounted, not only for the
+ * latest one (TanStack's own rule), so two quick calls each get their
+ * confirmation. The hook-level callbacks in `opts` run first; the
+ * invalidation runs before the hook-level `onSuccess`.
  */
 export function useToolMutation<TData = unknown>(
   toolName: string,
-  opts?: { invalidateKeys?: unknown[][] },
-) {
+  opts?: UseToolMutationOptions<TData>,
+): UseMutationResult<TData, Error, Record<string, unknown>> {
   const callTool = useCallTool()
   const queryClient = useQueryClient()
+  const { invalidateKeys, onSuccess, ...options } = opts ?? {}
 
-  return useMutation<TData, Error, Record<string, unknown>>({
+  const mutation = useMutation<TData, Error, Record<string, unknown>>({
+    ...options,
     mutationFn: async (args: Record<string, unknown>) => {
-      if (!callTool) throw new Error("callTool not available")
+      if (!callTool) throw new NoToolTransportError(toolName)
       const result = await callTool(toolName, args)
       return parseToolResult<TData>(result)
     },
-    onSuccess: async () => {
-      if (opts?.invalidateKeys) {
-        await Promise.all(
-          opts.invalidateKeys.map((key) => queryClient.refetchQueries({ queryKey: key })),
-        )
-      }
+    onSuccess: async (data, variables, onMutateResult, context) => {
+      await invalidateToolQueries(queryClient, invalidateKeys)
+      await onSuccess?.(data, variables, onMutateResult, context)
     },
   })
+
+  // Per-call callbacks fire only while mounted (TanStack's rule, kept), with
+  // the latest meta/mutationKey. Both refs are written in effects and read
+  // only when a call settles, so `mutate` stays stable across renders.
+  const mounted = useRef(false)
+  const scope = useRef<Pick<MutationFunctionContext, "meta" | "mutationKey">>({
+    meta: undefined,
+    mutationKey: undefined,
+  })
+  useEffect(() => {
+    scope.current = { meta: options.meta, mutationKey: options.mutationKey }
+  })
+  useEffect(() => {
+    mounted.current = true
+    return () => {
+      mounted.current = false
+    }
+  }, [])
+  const { mutateAsync } = mutation
+  const mutate = useCallback(
+    (
+      variables: Record<string, unknown>,
+      callbacks?: PerCallCallbacks<TData, Record<string, unknown>>,
+    ) => {
+      settleEachCall(mutateAsync(variables), variables, callbacks, () =>
+        mounted.current ? { client: queryClient, ...scope.current } : null,
+      )
+    },
+    [mutateAsync, queryClient],
+  )
+  return { ...mutation, mutate }
 }

package/src/hooks/use-tool-result-recovery.effects.test.tsx

@@ -0,0 +1,141 @@
+// @vitest-environment happy-dom
+import { afterEach, beforeEach, describe, expect, it, vi, type Mock } from "vitest"
+import type { ToolResultObserver } from "../lib/tool-result-observer.js"
+import { createToolResultObserver } from "../lib/tool-result-observer.js"
+import { advanceTimers, hookProbe, inAct, renderWithEffects } from "../render-effects.js"
+import {
+  DEFAULT_ASSUME_STRIPPED_AFTER_MS,
+  useToolResultRecovery,
+  type ToolResultRecoveryOptions,
+} from "./use-tool-result-recovery.js"
+
+/**
+ * The hook's wiring with its effects running (toolkit#176 K13): the default
+ * observer is the document's (`installToolResultObserver`), the controller
+ * subscribes to it on mount, and a stripped result re-runs the tool from the
+ * observer's notification alone, while a slow tool on a compliant host is
+ * never re-run. The decision itself is pinned in the recovery model's tests.
+ */
+
+const installed = vi.hoisted(() => ({ observer: null as ToolResultObserver | null }))
+
+// The document's observer, as a framed view would install it; the tests drive
+// a fake host channel instead of window.parent.
+vi.mock(import("../lib/tool-result-observer.js"), async (importOriginal) => {
+  const actual = await importOriginal()
+  return { ...actual, installToolResultObserver: () => installed.observer }
+})
+
+/** A host channel the observer listens on, with the SEP-1865 messages a host sends. */
+function fakeHost() {
+  const channel = new EventTarget()
+  const observer = createToolResultObserver(channel)
+  const send = (data: unknown) =>
+    inAct(() => channel.dispatchEvent(new MessageEvent("message", { data })))
+  const notify = (method: string, params: Record<string, unknown>) =>
+    send({ jsonrpc: "2.0", method, params })
+  return {
+    observer,
+    handshake: () =>
+      send({
+        jsonrpc: "2.0",
+        id: 1,
+        result: { protocolVersion: "2026-01-26", hostCapabilities: {} },
+      }),
+    toolInput: (args: Record<string, unknown>) =>
+      notify("ui/notifications/tool-input", { arguments: args }),
+    strippedResult: () =>
+      notify("ui/notifications/tool-result", { content: [{ type: "text", text: "Rendered." }] }),
+    structuredResult: (structuredContent: unknown) =>
+      notify("ui/notifications/tool-result", { content: [], structuredContent }),
+  }
+}
+
+type CallTool = ToolResultRecoveryOptions["callTool"]
+
+const PAYLOAD = { ok: true, rows: 3 }
+const isValid = (value: unknown) => (value as { ok?: unknown } | null)?.ok === true
+
+let callTool: Mock<CallTool>
+
+/** Mount the hook with the view's inputs; the probe keeps its latest state. */
+async function mount(over: Partial<ToolResultRecoveryOptions> = {}) {
+  const options: ToolResultRecoveryOptions = {
+    resultReady: false,
+    props: undefined,
+    isValid,
+    // The view runtime's input: a partial stream may have merged into it.
+    toolInput: { id: 7 },
+    hostContext: { toolInfo: { tool: { name: "show_list" } } },
+    callTool,
+    ...over,
+  }
+  const probe = hookProbe(() => useToolResultRecovery<typeof PAYLOAD>(options))
+  const view = await renderWithEffects(<probe.Probe />)
+  return { ...view, state: probe.current }
+}
+
+beforeEach(() => {
+  vi.useFakeTimers()
+  callTool = vi.fn<CallTool>(() => Promise.resolve({ structuredContent: PAYLOAD }))
+})
+
+afterEach(() => {
+  vi.useRealTimers()
+  installed.observer = null
+})
+
+describe("useToolResultRecovery with its effects (K13)", () => {
+  it("keep: a structured result the latch delivered is never re-run", async () => {
+    const host = fakeHost()
+    installed.observer = host.observer
+    await host.handshake()
+    await host.toolInput({ id: 7 })
+    await host.structuredResult(PAYLOAD)
+    const view = await mount({ resultReady: true, props: PAYLOAD })
+    await advanceTimers(vi.advanceTimersByTime, DEFAULT_ASSUME_STRIPPED_AFTER_MS * 4)
+    expect(callTool).not.toHaveBeenCalled()
+    expect(view.state().data).toBeNull()
+    await view.unmount()
+  })
+
+  it("slow: a compliant host still answering is never re-run, however long it takes", async () => {
+    const host = fakeHost()
+    installed.observer = host.observer
+    await host.handshake()
+    await host.toolInput({ id: 7 })
+    const view = await mount()
+    await advanceTimers(vi.advanceTimersByTime, DEFAULT_ASSUME_STRIPPED_AFTER_MS + 1)
+    await advanceTimers(vi.advanceTimersByTime, 10_000)
+    expect(callTool).not.toHaveBeenCalled()
+    expect(view.state()).toMatchObject({ data: null, failed: false, blocked: null })
+    await view.unmount()
+  })
+
+  it("strip: the observer's notification alone re-runs the tool once, with the complete input", async () => {
+    const host = fakeHost()
+    installed.observer = host.observer
+    await host.handshake()
+    const view = await mount()
+    await host.toolInput({ id: 7, page: 2 })
+    expect(callTool).not.toHaveBeenCalled()
+    // Nothing re-renders the view here: only the subscription can react.
+    await host.strippedResult()
+    expect(callTool).toHaveBeenCalledTimes(1)
+    expect(callTool).toHaveBeenCalledWith("show_list", { id: 7, page: 2 })
+    expect(view.state()).toMatchObject({ data: PAYLOAD, failed: false, args: { id: 7, page: 2 } })
+    await advanceTimers(vi.advanceTimersByTime, 10_000)
+    expect(callTool, "a single flight").toHaveBeenCalledTimes(1)
+    await view.unmount()
+  })
+
+  it("without an observer (observer: null) the 2.x grace timer re-runs a silent call", async () => {
+    const view = await mount({ observer: null })
+    await advanceTimers(vi.advanceTimersByTime, DEFAULT_ASSUME_STRIPPED_AFTER_MS - 1)
+    expect(callTool).not.toHaveBeenCalled()
+    await advanceTimers(vi.advanceTimersByTime, 1)
+    expect(callTool).toHaveBeenCalledTimes(1)
+    expect(callTool).toHaveBeenCalledWith("show_list", { id: 7 })
+    await view.unmount()
+  })
+})

package/src/hooks/use-tool-result-recovery.test.ts

@@ -1,10 +1,13 @@
 import { describe, expect, it } from "vitest"
-import { decodeRecoveredResult, resolveRecoveryToolName } from "./use-tool-result-recovery.js"
+import {
+  decodeRecoveredResult,
+  isRecoveryReady,
+  resolveRecoveryToolName,
+} from "./use-tool-result-recovery.js"
 
-// The hook's effect wiring (single-flight, StrictMode re-run, host-payload
-// precedence) is exercised end-to-end by the consumer's Playwright host
-// simulation (structuredContent keep/strip scenarios); these tests pin the
-// pure decision logic.
+// The hook's effect wiring (default observer, subscription, keep/strip/slow)
+// runs in use-tool-result-recovery.effects.test.tsx and end-to-end in the
+// consumer's Playwright host simulation; these tests pin the pure helpers.
 
 describe("resolveRecoveryToolName", () => {
   it("reads the SEP-1865 shape toolInfo.tool.name first", () => {
@@ -63,3 +66,23 @@ describe("decodeRecoveredResult", () => {
     expect(decodeRecoveredResult({ content: [] }, isView)).toBeNull()
   })
 })
+
+describe("isRecoveryReady (toolkit#176)", () => {
+  const base = { resultReady: false, assumedReady: false, cancelled: false, errored: false }
+
+  it("fires once a result latched or the stripped-result grace period passed", () => {
+    expect(isRecoveryReady(base)).toBe(false)
+    expect(isRecoveryReady({ ...base, resultReady: true })).toBe(true)
+    expect(isRecoveryReady({ ...base, assumedReady: true })).toBe(true)
+  })
+
+  it("never re-executes a call that failed (K14: isError renders the error instead)", () => {
+    expect(isRecoveryReady({ ...base, resultReady: true, errored: true })).toBe(false)
+    expect(isRecoveryReady({ ...base, assumedReady: true, errored: true })).toBe(false)
+  })
+
+  it("never re-executes a call the host cancelled", () => {
+    expect(isRecoveryReady({ ...base, resultReady: true, cancelled: true })).toBe(false)
+    expect(isRecoveryReady({ ...base, assumedReady: true, cancelled: true })).toBe(false)
+  })
+})