Preserve generated UI mutation pending state

RhysSullivan · RhysSullivan · commit a2023cad247d · 2026-05-13T12:05:05.000-07:00
diff --git a/packages/core/execution/src/description.ts b/packages/core/execution/src/description.ts
@@ -72,7 +72,6 @@ const formatDescription = (sources: readonly Source[]): string => {
     "- Use the discovered output shape exactly. Do not invent wrapper fields like `data.domain` or `data.items` unless the schema/sample shows them.",
     "- For toggles and switches, mutate with the checked value from the event instead of inverting possibly stale query data.",
     "- For optimistic writes, use TanStack `onMutate` / `onError` / `onSettled`: cancel the query, snapshot old data, `setQueryData`, roll back on error, then invalidate.",
-    "- Success messages must use `mutation.variables` or neutral wording, not query-derived state that may still be stale while invalidation refetches.",
     "- Only hardcode small display constants like labels, colors, tab names, and chart configuration. Never embed tool response rows, API results, summaries, or dashboard data as literals in the component.",
     "- Always render loading and error states from `useQuery` / `useMutation`; do not replace them with hardcoded fallback data.",
     "- Tools: `tools.<namespace>.<tool>(args)` — call any configured API tool (never use raw `fetch`). Tool helpers: `.queryOptions(args, options)`, `.mutationOptions(options)`, `.queryKey(args)`, `.pathKey()`, and `.mutationKey()`.",
diff --git a/packages/hosts/mcp/src/server.test.ts b/packages/hosts/mcp/src/server.test.ts
@@ -254,49 +254,6 @@ describe("MCP host server — client with elicitation", () => {
     );
   });
 
-  it("render-ui rejects stale toggle mutation state patterns", async () => {
-    await withClient(
-      makeStubEngine({}),
-      APPS_ELICITATION_CAPS,
-      async (client) => {
-        const invertedToggle = await client.callTool({
-          name: "render-ui",
-          arguments: {
-            code: [
-              "function App() {",
-              "  const autoRenew = false;",
-              "  const mutation = useMutation(tools.vercel_api.domainsRegistrar.updateDomainAutoRenew.mutationOptions({}));",
-              "  return <Switch onCheckedChange={() => mutation.mutate({ body: { autoRenew: !autoRenew } })} />;",
-              "}",
-            ].join("\n"),
-          },
-        });
-        expect(invertedToggle.isError).toBe(true);
-        expect(textOf(invertedToggle)).toContain(
-          'Mutation payloads must not invert query-backed state like "!autoRenew"',
-        );
-
-        const staleSuccessLabel = await client.callTool({
-          name: "render-ui",
-          arguments: {
-            code: [
-              "function App() {",
-              "  const autoRenew = false;",
-              "  const mutation = { isSuccess: true };",
-              '  return <div>{mutation.isSuccess && <span>Auto-renew {autoRenew ? "enabled" : "disabled"} successfully</span>}</div>;',
-              "}",
-            ].join("\n"),
-          },
-        });
-        expect(staleSuccessLabel.isError).toBe(true);
-        expect(textOf(staleSuccessLabel)).toContain(
-          'Mutation success text must not read "autoRenew"',
-        );
-      },
-      { plugins: [DYNAMIC_UI_PLUGIN] },
-    );
-  });
-
   it("splits execution and UI rendering into separate model-facing tool descriptions", async () => {
     const description = [
       "Execute TypeScript in a sandboxed runtime.",
@@ -337,12 +294,6 @@ describe("MCP host server — client with elicitation", () => {
         expect(renderUi?.description).toContain("For optimistic UI, use `onMutate`");
         expect(renderUi?.description).toContain("Do not call API tools first");
         expect(renderUi?.description).toContain("Do not redeclare or destructure provided globals");
-        expect(renderUi?.description).toContain(
-          "server rejects toggle mutations that invert query-backed booleans",
-        );
-        expect(renderUi?.description).toContain(
-          "server rejects mutation success labels that read enabled/disabled wording",
-        );
         expect(renderUi?.description).toContain("server rejects obvious hardcoded live-data");
         expect(renderUi?.description).toContain("server rejects redeclarations");
         expect(renderUi?.description).toContain("- `axiom_mcp`");
diff --git a/packages/hosts/mcp/src/shell/hooks.ts b/packages/hosts/mcp/src/shell/hooks.ts
@@ -7,7 +7,7 @@ import {
   skipToken,
   useMutation as useTanStackMutation,
   useQuery as useTanStackQuery,
-  useQueryClient,
+  useQueryClient as useTanStackQueryClient,
   type MutationKey,
   type QueryClient as QueryClientType,
   type QueryKey,
@@ -17,17 +17,11 @@ import {
   type UseQueryResult,
 } from "@tanstack/react-query";
 
-export {
-  QueryClient,
-  QueryClientProvider,
-  mutationOptions,
-  queryOptions,
-  skipToken,
-  useQueryClient,
-};
+export { QueryClient, QueryClientProvider, mutationOptions, queryOptions, skipToken };
 
 let legacyQueryId = 0;
 let legacyMutationId = 0;
+const invalidationScopes: Array<Array<Promise<unknown>>> = [];
 
 type RefetchableQuery = {
   refetch: () => unknown;
@@ -49,6 +43,68 @@ const nextLegacyMutationKey = (): MutationKey => [
   ++legacyMutationId,
 ];
 
+const trackInvalidation = (promise: Promise<unknown>) => {
+  for (const scope of invalidationScopes) {
+    scope.push(promise);
+  }
+  return promise;
+};
+
+const trackMutationCallback = async <T>(callback: () => T | Promise<T>): Promise<T> => {
+  const scope: Array<Promise<unknown>> = [];
+  invalidationScopes.push(scope);
+  try {
+    const result = await callback();
+    await Promise.allSettled(scope);
+    return result;
+  } finally {
+    invalidationScopes.pop();
+  }
+};
+
+const wrapQueryClient = (client: QueryClientType): QueryClientType =>
+  new Proxy(client, {
+    get(target, prop, receiver) {
+      const value = Reflect.get(target, prop, receiver) as unknown;
+      if (prop === "invalidateQueries" && typeof value === "function") {
+        return (...args: unknown[]) =>
+          trackInvalidation(
+            (value as (...args: unknown[]) => Promise<unknown>).apply(target, args),
+          );
+      }
+      return typeof value === "function" ? value.bind(target) : value;
+    },
+  }) as QueryClientType;
+
+export const useQueryClient = (queryClient?: QueryClientType): QueryClientType => {
+  const client = useTanStackQueryClient(queryClient);
+  const wrappedRef = useRef<{ client: QueryClientType; wrapped: QueryClientType } | null>(null);
+
+  if (wrappedRef.current?.client !== client) {
+    wrappedRef.current = { client, wrapped: wrapQueryClient(client) };
+  }
+
+  return wrappedRef.current.wrapped;
+};
+
+const wrapMutationOptions = <TData, TError, TVariables, TContext>(
+  options: UseMutationOptions<TData, TError, TVariables, TContext>,
+): UseMutationOptions<TData, TError, TVariables, TContext> => ({
+  ...options,
+  onSuccess: options.onSuccess
+    ? (...args: Parameters<NonNullable<typeof options.onSuccess>>) =>
+        trackMutationCallback(() => options.onSuccess?.(...args))
+    : undefined,
+  onError: options.onError
+    ? (...args: Parameters<NonNullable<typeof options.onError>>) =>
+        trackMutationCallback(() => options.onError?.(...args))
+    : undefined,
+  onSettled: options.onSettled
+    ? (...args: Parameters<NonNullable<typeof options.onSettled>>) =>
+        trackMutationCallback(() => options.onSettled?.(...args))
+    : undefined,
+});
+
 /**
  * TanStack Query's `useQuery`, plus compatibility for the original generated
  * UI shorthand: `useQuery(() => tools.namespace.tool(args))`.
@@ -123,14 +179,14 @@ export function useMutation(
       ...tanstackOptions,
       mutationFn: optionsOrFn as (input: unknown) => Promise<unknown>,
       onSuccess: async (data, variables, context, mutationContext) => {
-        await onSuccess?.(data, variables, context, mutationContext);
+        await trackMutationCallback(() => onSuccess?.(data, variables, context, mutationContext));
         await Promise.all(invalidates?.map((query) => query.refetch()) ?? []);
       },
     });
   }
 
   return useTanStackMutation(
-    optionsOrFn as UseMutationOptions<unknown, Error, unknown, unknown>,
+    wrapMutationOptions(optionsOrFn as UseMutationOptions<unknown, Error, unknown, unknown>),
     optionsOrQueryClient as QueryClientType | undefined,
   );
 }
diff --git a/packages/hosts/mcp/src/shell/mcp-app.browser.test.ts b/packages/hosts/mcp/src/shell/mcp-app.browser.test.ts
@@ -169,10 +169,11 @@ function App() {
   const domainQuery = useQuery(tools.inventory.domains.getDomain.queryOptions(domainArgs));
   const updateAutoRenew = useMutation(
     tools.inventory.domains.updateDomainAutoRenew.mutationOptions({
-      onSuccess: () =>
+      onSuccess: () => {
         queryClient.invalidateQueries({
           queryKey: tools.inventory.domains.getDomain.queryKey(domainArgs),
-        }),
+        });
+      },
     })
   );
 
@@ -184,6 +185,12 @@ function App() {
         <div id="auto-renew-state">
           {domainQuery.isLoading ? "loading" : String(autoRenew)}
         </div>
+        <div id="auto-renew-pending">{String(updateAutoRenew.isPending)}</div>
+        <div id="auto-renew-success">
+          {updateAutoRenew.isSuccess
+            ? "Auto-renew " + (autoRenew ? "enabled" : "disabled") + " successfully"
+            : ""}
+        </div>
         <Button
           id="auto-renew-toggle"
           disabled={domainQuery.isLoading || updateAutoRenew.isPending}
@@ -337,6 +344,7 @@ const startOpenApiServer = (): Promise<OpenApiServer> =>
   new Promise((resolveServer, rejectServer) => {
     let baseUrl = "";
     let domainAutoRenew = false;
+    let nextDomainGetDelayMs = 0;
     const postRequests: string[] = [];
 
     const server: Server = createServer(async (request, response) => {
@@ -501,6 +509,11 @@ const startOpenApiServer = (): Promise<OpenApiServer> =>
       if (request.method === "GET" && request.url?.startsWith("/domains/")) {
         const domain = decodeURIComponent(request.url.slice("/domains/".length));
         if (!domain.includes("/")) {
+          const delayMs = nextDomainGetDelayMs;
+          nextDomainGetDelayMs = 0;
+          if (delayMs > 0) {
+            await new Promise((resolveDelay) => setTimeout(resolveDelay, delayMs));
+          }
           response.statusCode = 200;
           response.setHeader("content-type", "application/json");
           response.end(JSON.stringify({ domain, renew: domainAutoRenew }));
@@ -533,6 +546,7 @@ const startOpenApiServer = (): Promise<OpenApiServer> =>
         postRequests.push(body);
         const parsed = JSON.parse(body) as { autoRenew?: unknown };
         domainAutoRenew = parsed.autoRenew === true;
+        nextDomainGetDelayMs = 300;
         const domainPath = request.url.slice(
           "/domains/".length,
           request.url.length - "/auto-renew".length,
@@ -1030,11 +1044,22 @@ describe("MCP app generated UI browser isolation", () => {
       expect(openApiServer.postRequests).toHaveLength(initialPostCount);
 
       await shellFrame.getByRole("button", { name: "Approve" }).click({ timeout: 10_000 });
+      await innerFrame.waitForTimeout(100);
+      expect(await innerFrame.locator("#auto-renew-pending").textContent()).toBe("true");
+      expect(await innerFrame.locator("#auto-renew-success").textContent()).toBe("");
+
       await innerFrame.waitForFunction(
         () => document.querySelector("#auto-renew-state")?.textContent === "true",
         undefined,
         { timeout: 10_000 },
       );
+      await innerFrame.waitForFunction(
+        () =>
+          document.querySelector("#auto-renew-success")?.textContent ===
+          "Auto-renew enabled successfully",
+        undefined,
+        { timeout: 10_000 },
+      );
 
       expect(openApiServer.postRequests).toHaveLength(initialPostCount + 1);
       expect(openApiServer.postRequests.at(-1)).toBe(JSON.stringify({ autoRenew: true }));
diff --git a/packages/plugins/dynamic-ui/src/mcp.ts b/packages/plugins/dynamic-ui/src/mcp.ts
@@ -338,7 +338,6 @@ export const buildRenderUiDescription = (executeDescription: string): string =>
       "- Use the discovered output shape exactly. Do not invent wrapper fields like `data.domain` or `data.items` unless the schema/sample shows them.",
       "- For toggles and switches, mutate with the checked value from the event instead of inverting possibly stale query data.",
       "- For optimistic writes, use TanStack `onMutate` / `onError` / `onSettled`: cancel the query, snapshot old data, `setQueryData`, roll back on error, then invalidate.",
-      "- Success messages must use `mutation.variables` or neutral wording, not query-derived state that may still be stale while invalidation refetches.",
       "- Only hardcode small display constants like labels, colors, tab names, and chart configuration. Never embed tool response rows, API results, summaries, or dashboard data as literals in the component.",
       "- Always render loading and error states from `useQuery` / `useMutation`; do not replace them with hardcoded fallback data.",
       `- shadcn/ui components available by name: ${SHADCN_COMPONENTS}`,
@@ -373,9 +372,6 @@ export const buildRenderUiDescription = (executeDescription: string): string =>
     "- For simple writes, invalidate with `queryClient.invalidateQueries(tools.<namespace>.<queryTool>.queryFilter(args))` in `onSuccess` or `onSettled`.",
     "- For toggles and switches, pass the new checked value into `mutate`: `onCheckedChange={(checked) => mutation.mutate({ body: { enabled: checked } })}`.",
     "- For optimistic UI, use `onMutate` to `cancelQueries`, snapshot `getQueryData`, and `setQueryData`; return the snapshot, restore it in `onError`, and invalidate in `onSettled`.",
-    "- Do not send mutation payloads by inverting query state like `autoRenew: !autoRenew`; the cached value can be stale after an approval flow.",
-    '- Do not show success text by combining `mutation.isSuccess` with query-derived state like `autoRenew ? "enabled" : "disabled"`. `isSuccess` becomes true before invalidated reads finish refetching.',
-    '- For action-specific success text, read the submitted value from `mutation.variables` such as `mutation.variables?.body?.autoRenew ? "enabled" : "disabled"`, or show neutral text like `Saved successfully`.',
     "",
     "## Available UI Components",
     "",
@@ -391,8 +387,6 @@ export const buildRenderUiDescription = (executeDescription: string): string =>
     "- Keep data live by routing every API read/write through the provided `tools` proxy from TanStack Query or `run(code)`.",
     "- Do not redeclare or destructure provided globals. Hooks, components, icons, `tools`, and `run` are already in scope; use them directly.",
     "- Tool proxy helpers are TanStack-native: `.queryOptions(args, options)`, `.mutationOptions(options)`, `.queryKey(args)`, `.queryFilter(args, filters)`, `.pathKey()`, `.pathFilter(filters)`, and `.mutationKey()`.",
-    "- The server rejects toggle mutations that invert query-backed booleans such as `autoRenew: !autoRenew`; use the checked value from the UI event.",
-    "- The server rejects mutation success labels that read enabled/disabled wording from stale query state; use `mutation.variables` or neutral text.",
     "- The server rejects obvious hardcoded live-data snapshots such as `const rows = [{...}, {...}]`; regenerate with `useQuery` instead.",
     "- The server rejects redeclarations of provided globals such as `const { useState } = React` or `const Card = ...` before the UI reaches the iframe.",
     "",
@@ -416,19 +410,6 @@ const OBJECT_DESTRUCTURING_DECLARATION = /\b(?:const|let|var)\s*\{([^{}]*)\}\s*=
 const PROVIDED_GLOBAL_DECLARATION =
   /\b(?:const|let|var)\s+([A-Za-z_$][\w$]*)\b|\bfunction\s+([A-Za-z_$][\w$]*)\s*\(|\bclass\s+([A-Za-z_$][\w$]*)\b/g;
 
-const QUERY_BACKED_BOOLEAN_NAME =
-  "(?:autoRenew|enabled|isEnabled|checked|isChecked|active|isActive|disabled|isDisabled)";
-
-const MUTATION_INVERTS_QUERY_STATE = new RegExp(
-  String.raw`\bmutate\s*\(\s*\{[\s\S]{0,1000}\b[A-Za-z_$][\w$]*\s*:\s*!\s*(${QUERY_BACKED_BOOLEAN_NAME})\b`,
-  "i",
-);
-
-const MUTATION_SUCCESS_USES_QUERY_STATE_LABEL = new RegExp(
-  String.raw`\b\w+\.isSuccess\b[\s\S]{0,1200}\{\s*(${QUERY_BACKED_BOOLEAN_NAME})\s*\?\s*["'\`](?:enabled|disabled|on|off|activated|deactivated)\b`,
-  "i",
-);
-
 const firstDefined = (...values: Array<string | undefined>): string | undefined =>
   values.find((value): value is string => value !== undefined);
 
@@ -473,22 +454,6 @@ export const validateRenderUiCode = (code: string): string | null => {
     }
   }
 
-  const invertedState = MUTATION_INVERTS_QUERY_STATE.exec(code)?.[1];
-  if (invertedState) {
-    return [
-      `Mutation payloads must not invert query-backed state like "!${invertedState}".`,
-      "Switch and checkbox handlers receive the next checked value; pass that value into mutate instead.",
-    ].join(" ");
-  }
-
-  const staleSuccessState = MUTATION_SUCCESS_USES_QUERY_STATE_LABEL.exec(code)?.[1];
-  if (staleSuccessState) {
-    return [
-      `Mutation success text must not read "${staleSuccessState}" from query state.`,
-      "`mutation.isSuccess` becomes true before invalidated queries finish refetching; use mutation.variables for action-specific text or show neutral saved text.",
-    ].join(" ");
-  }
-
   for (const match of code.matchAll(OBJECT_ARRAY_LITERAL)) {
     const name = match[1];
     const body = match[2] ?? "";
