@miragon/mcp-toolkit-ui 2.7.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

@@ -5,10 +5,9 @@ import {
   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", () => {

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

@@ -1,44 +1,55 @@
-import { useEffect, useRef, useState } from "react"
-import { parseToolResult } from "../lib/parse-tool-result.js"
+import { useEffect, useState, useSyncExternalStore } from "react"
+import { installToolResultObserver, type ToolResultObserver } from "../lib/tool-result-observer.js"
+import {
+  createRecoveryController,
+  type RecoveryPolicy,
+  type ToolResultRecoveryState,
+} from "./tool-result-recovery-model.js"
+
+export { decodeRecoveredResult } from "./tool-result-recovery-model.js"
+export type {
+  RecoveryBlockReason,
+  RecoveryPolicy,
+  ToolResultRecoveryState,
+  ToolSafety,
+} from "./tool-result-recovery-model.js"
 
 /**
  * Recovery for hosts that deliver the `ui/notifications/tool-result`
  * notification without `structuredContent` (observed on claude.ai and Claude
- * Desktop, which forward only the model-facing `content` blocks). Since
- * mcp-use 2.x such a result is not even observable: the `useToolContext`
- * latch ignores content-only successes, so the tool status stays `"pending"`
- * forever and the shell would idle on its loading skeleton.
+ * Desktop, which forward only the model-facing `content` blocks). mcp-use
+ * 2.x's `useToolContext` latch ignores such a content-only result, so its
+ * status stays `"pending"` forever and the shell would idle on its skeleton.
  *
- * Because the stripped notification is invisible, readiness is inferred two
- * ways: the hard `resultReady` signal (the latch fired but the payload is
- * invalid), or a grace timer — once the originating tool is known and no
- * result has latched within
- * {@link ToolResultRecoveryOptions.assumeStrippedAfterMs}, the result is
- * assumed delivered-but-stripped. Either way the hook re-executes the
- * originating tool once via the host bridge's official `callTool` — tool-call
- * RESPONSES carry `structuredContent` intact on those hosts — and exposes the
- * decoded payload. Spec-conforming hosts (e.g. the mcp-use inspector) deliver
- * a valid payload up front, so the recovery never fires there.
+ * The hook re-executes the originating tool once via the host bridge's
+ * `callTool` — tool-call RESPONSES carry `structuredContent` intact on those
+ * hosts — and exposes the decoded payload. When it does so (toolkit#176):
  *
- * A call the host cancelled is never recovered: since mcp-use 2.7.1 the
- * cancellation latches as `status: "error"` (a `ToolCancelledError`), which
- * would otherwise count as a ready-but-invalid result — re-executing it would
- * override the user's decision to stop the call. Neither is a call that
- * failed (`isError` or a transport error, {@link ToolResultRecoveryOptions.errored}):
- * re-running it only repeats the failure, so the shell renders the error and
- * offers an explicit retry instead (toolkit#176).
+ * - **Only when the host stripped the result.** A passive observer of the
+ *   host's lifecycle notifications (`installToolResultObserver`) sees the
+ *   content-only result itself, so the recovery fires as it arrives, and a
+ *   slow tool on a compliant host is never re-run (K13): no result has
+ *   arrived because none was sent. Where the observer cannot vouch for that
+ *   (no browser frame, or installed after the handshake), the 2.x grace timer
+ *   ({@link ToolResultRecoveryOptions.assumeStrippedAfterMs}) still infers a
+ *   stripped result, once the call's arguments arrived.
+ * - **With the call's complete arguments** (N21): the observer's complete
+ *   `tool-input`, never a partial stream; an equal re-delivery is the same run.
+ * - **Never for a cancelled or failed call:** re-running repeats the failure
+ *   or overrides the user's stop; the shell renders that state and offers an
+ *   explicit retry.
+ * - **Only for a tool the policy allows** ({@link ToolResultRecoveryOptions.policy}):
+ *   by default never one the host reports as writing (`readOnlyHint: false`
+ *   or `destructiveHint: true` in `hostContext.toolInfo.tool.annotations`).
+ *   A blocked recovery is reported (`blocked`) so the view can show its error.
  *
- * Assumes the originating tool is a read-only render pipeline (true for all
- * toolkit `render-view`/`show_*` tools) — the re-execution must be safe to
- * repeat, and may rarely run concurrently with the original invocation when
- * the grace timer beats a slow but conforming host.
+ * The decision and its timer live in `createRecoveryController`; this hook
+ * only binds it to React.
  */
 
 /**
- * Default grace period before a silent latch counts as a stripped result.
- * Long enough for a conforming host to latch a typical render tool first,
- * short enough that users on stripping hosts are not left staring at the
- * skeleton noticeably.
+ * Default grace period before a silent latch counts as a stripped result —
+ * used only where the transport observer cannot vouch that nothing arrived.
  */
 export const DEFAULT_ASSUME_STRIPPED_AFTER_MS = 2500
 
@@ -46,8 +57,7 @@ export interface ToolResultRecoveryOptions {
   /**
    * The latch fired (`useToolContext().status !== "pending"`) — a structured
    * result, a tool error, or a host cancellation arrived. On hosts that strip
-   * `structuredContent` this never turns true; the grace timer below covers
-   * that case.
+   * `structuredContent` this never turns true.
    */
   resultReady: boolean
   /**
@@ -66,10 +76,7 @@ export interface ToolResultRecoveryOptions {
   errored?: boolean
   /** The payload the host delivered (`useToolContext().toolOutput`). */
   props: unknown
-  /**
-   * Validates a candidate payload. MUST be referentially stable (module-level
-   * function or memoised) — it is an effect dependency.
-   */
+  /** Validates a candidate payload. */
   isValid: (value: unknown) => boolean
   /** The tool arguments from `ui/notifications/tool-input` (`useToolContext().toolInput`). */
   toolInput: Record<string, unknown> | null | undefined
@@ -79,23 +86,21 @@ export interface ToolResultRecoveryOptions {
   callTool: (name: string, args: Record<string, unknown>) => Promise<unknown>
   /**
    * Grace period in ms after which a silent latch counts as a stripped
-   * result. Armed once the originating tool is known while `resultReady` is
-   * still false; re-armed on every `toolInput` update, because streaming
-   * arguments mean the invocation may not have finished yet. Defaults to
+   * result. Applies only where the transport observer cannot vouch that no
+   * result arrived (see above); restarts for new arguments. Defaults to
    * {@link DEFAULT_ASSUME_STRIPPED_AFTER_MS}.
    */
   assumeStrippedAfterMs?: number
+  /** Which tools may be re-executed on their own. Default `"auto"`. */
+  policy?: RecoveryPolicy
+  /**
+   * The transport observer. Defaults to the document's
+   * (`installToolResultObserver()`); `null` disables it, leaving the grace
+   * timer to decide alone (the 2.x behaviour).
+   */
+  observer?: ToolResultObserver | null
 }
 
-export interface ToolResultRecoveryState<T> {
-  /** The recovered payload, `null` until the re-execution resolved (or when it failed). */
-  data: T | null
-  /** `true` when the single recovery attempt failed — the shell keeps its skeleton, no retry loop. */
-  failed: boolean
-}
-
-const IDLE: ToolResultRecoveryState<never> = { data: null, failed: false }
-
 /**
  * Resolves the originating tool name from the host context's `toolInfo`.
  * Reads the SEP-1865 shape first (`toolInfo.tool.name`), then the flattened
@@ -119,26 +124,10 @@ export function resolveRecoveryToolName(hostContext: unknown): string | undefine
 }
 
 /**
- * Decodes a re-executed tool result and validates it. Exposed for tests;
- * the hook uses it on the `callTool` response.
- */
-export function decodeRecoveredResult<T>(
-  result: unknown,
-  isValid: (value: unknown) => boolean,
-): T | null {
-  let decoded: unknown
-  try {
-    decoded = parseToolResult<unknown>(result)
-  } catch {
-    return null
-  }
-  return decoded !== null && isValid(decoded) ? (decoded as T) : null
-}
-
-/**
- * Whether the recovery may fire: the result latched (or is assumed stripped)
- * and the call was neither cancelled nor failed. Exported for tests — this is
- * the #176 rule "never recover on an error status".
+ * Whether a latched or assumed-stripped result may be recovered at all: the
+ * call was neither cancelled nor failed. Kept for existing imports; the full
+ * rule (observer, policy, arguments) is `decideRecovery`.
+ * @deprecated Use `decideRecovery` from the recovery model.
  */
 export function isRecoveryReady(state: {
   resultReady: boolean
@@ -163,81 +152,34 @@ export function useToolResultRecovery<T>(
     hostContext,
     callTool,
     assumeStrippedAfterMs = DEFAULT_ASSUME_STRIPPED_AFTER_MS,
+    policy = "auto",
+    observer,
   } = options
-  const [state, setState] = useState<ToolResultRecoveryState<T>>(IDLE)
-  // Single-flight: the promise ref survives StrictMode's dev-mode
-  // effect re-run (cleanup + setup on the same instance), so the re-run
-  // re-subscribes to the SAME in-flight call instead of firing a second one.
-  const flightRef = useRef<Promise<ToolResultRecoveryState<T>> | null>(null)
-  const flightInputRef = useRef<unknown>(null)
-
-  const toolName = resolveRecoveryToolName(hostContext)
-
-  // A stripped result is unobservable (the 2.x latch ignores content-only
-  // results), so infer it: once a recovery could fire but the latch has been
-  // silent for the grace period, treat the result as delivered-but-stripped.
-  // Re-armed on every toolInput identity change — streaming arguments mean
-  // the invocation may still be running, so the grace period starts over.
-  // The re-arm is a render-phase reset (not an effect) so the cascading-
-  // render lint stays satisfied; the timer effect below restarts anyway
-  // because toolInput is one of its dependencies.
-  const [assumedReady, setAssumedReady] = useState(false)
-  const [prevTimerInput, setPrevTimerInput] = useState<unknown>(toolInput)
-  if (prevTimerInput !== toolInput) {
-    setPrevTimerInput(toolInput)
-    if (assumedReady) setAssumedReady(false)
-  }
-  const suppressed = cancelled || errored
+  // One controller per mount; it survives StrictMode's effect re-run, so an
+  // in-flight call is never fired twice.
+  const [controller] = useState(() =>
+    createRecoveryController<T>({
+      observer: observer === undefined ? installToolResultObserver() : observer,
+    }),
+  )
   useEffect(() => {
-    if (resultReady || suppressed || toolName === undefined) return
-    const timer = setTimeout(() => setAssumedReady(true), assumeStrippedAfterMs)
-    return () => clearTimeout(timer)
-  }, [resultReady, suppressed, toolName, toolInput, assumeStrippedAfterMs])
-
-  // A cancelled or failed call never counts as ready, so the effect below
-  // invalidates any in-flight attempt the same way it does for a superseded run.
-  const ready = isRecoveryReady({ resultReady, assumedReady, cancelled, errored })
-
+    controller.start()
+    return () => controller.stop()
+  }, [controller])
   useEffect(() => {
-    // A new tool run (result no longer ready, or fresh tool-input identity)
-    // invalidates any previous recovery attempt and re-arms the single-flight.
-    if (flightRef.current && (!ready || flightInputRef.current !== toolInput)) {
-      flightRef.current = null
-      setState((prev) => (prev.data === null && !prev.failed ? prev : IDLE))
-    }
-    // Never fire on conforming hosts (valid payload) or before the
-    // originating tool is known.
-    if (!ready || isValid(props) || toolName === undefined) return
-
-    if (!flightRef.current) {
-      flightInputRef.current = toolInput
-      flightRef.current = (async (): Promise<ToolResultRecoveryState<T>> => {
-        try {
-          const result = await callTool(toolName, toolInput ?? {})
-          const data = decodeRecoveredResult<T>(result, isValid)
-          if (data !== null) return { data, failed: false }
-          console.warn(
-            `[mcp-toolkit] tool-result recovery: re-executing "${toolName}" returned no usable payload`,
-          )
-          return { data: null, failed: true }
-        } catch (err) {
-          console.warn(`[mcp-toolkit] tool-result recovery for "${toolName}" failed:`, err)
-          return { data: null, failed: true }
-        }
-      })()
-    }
-
-    const flight = flightRef.current
-    let stale = false
-    void flight.then((outcome) => {
-      // Drop the outcome when this subscription was superseded or the flight
-      // was invalidated by a newer tool run.
-      if (!stale && flightRef.current === flight) setState(outcome)
+    controller.update({
+      resultReady,
+      cancelled,
+      errored,
+      props,
+      isValid,
+      toolInput,
+      hostContext,
+      callTool,
+      policy,
+      graceMs: assumeStrippedAfterMs,
+      resolveToolName: resolveRecoveryToolName,
     })
-    return () => {
-      stale = true
-    }
-  }, [ready, props, isValid, toolName, toolInput, callTool])
-
-  return state
+  })
+  return useSyncExternalStore(controller.subscribe, controller.getState, controller.getState)
 }