fix assistant custom widget creation

apps/nextjs/src/components/assistant/assistant-auto-approval.spec.ts

@@ -2,7 +2,7 @@
 
 import { act, createElement, Fragment } from "react";
 import { createRoot } from "react-dom/client";
-import { afterEach, describe, expect, test } from "vitest";
+import { afterEach, describe, expect, test, vi } from "vitest";
 
 import {
   AssistantAutoApprovalProvider,
@@ -67,6 +67,19 @@ describe("assistant automatic approval tracker", () => {
     expect(tracker.claim("app-create-1")).toBe(true);
   });
 
+  test("bounds automatic retries and permits a fresh claim after completion", () => {
+    const tracker = createAssistantAutoApprovalTracker(2);
+
+    expect(tracker.claim("custom-widget-create-1")).toBe(true);
+    tracker.release("custom-widget-create-1");
+    expect(tracker.claim("custom-widget-create-1")).toBe(true);
+    tracker.release("custom-widget-create-1");
+    expect(tracker.claim("custom-widget-create-1")).toBe(false);
+
+    tracker.complete("custom-widget-create-1");
+    expect(tracker.claim("custom-widget-create-1")).toBe(true);
+  });
+
   test("keeps the setting for the same local conversation and resets it on a real switch", async () => {
     const container = document.createElement("div");
     containers.push(container);
@@ -185,4 +198,48 @@ describe("assistant automatic approval tracker", () => {
 
     await act(async () => root.unmount());
   });
+
+  test("retries an approval response that never settles and restores manual controls after two attempts", async () => {
+    vi.useFakeTimers();
+    const container = document.createElement("div");
+    containers.push(container);
+    document.body.append(container);
+    const root = createRoot(container);
+    let confirmationCount = 0;
+
+    await act(async () => {
+      root.render(
+        createElement(
+          AssistantAutoApprovalProvider,
+          { conversationId: "local-thread-1" },
+          createElement(
+            Fragment,
+            null,
+            createElement(AutoApprovalProbe),
+            createElement(AutomaticActionProbe, {
+              ready: true,
+              completed: false,
+              confirm: () => {
+                confirmationCount += 1;
+              },
+            }),
+          ),
+        ),
+      );
+    });
+    act(() => currentAutoApproval?.setEnabled(true));
+    expect(confirmationCount).toBe(1);
+    expect(automaticActionInProgress).toBe(true);
+
+    await act(async () => vi.advanceTimersByTime(5_000));
+    expect(confirmationCount).toBe(2);
+    expect(automaticActionInProgress).toBe(true);
+
+    await act(async () => vi.advanceTimersByTime(5_000));
+    expect(confirmationCount).toBe(2);
+    expect(automaticActionInProgress).toBe(false);
+
+    await act(async () => root.unmount());
+    vi.useRealTimers();
+  });
 });

apps/nextjs/src/components/assistant/assistant-auto-approval.tsx

@@ -3,20 +3,32 @@
 import type { PropsWithChildren } from "react";
 import { createContext, useCallback, useContext, useEffect, useMemo, useRef, useState } from "react";
 
-export const createAssistantAutoApprovalTracker = () => {
+const AUTO_APPROVAL_RETRY_TIMEOUT_MS = 5_000;
+const AUTO_APPROVAL_MAX_ATTEMPTS = 2;
+
+export const createAssistantAutoApprovalTracker = (maximumAttempts = Number.POSITIVE_INFINITY) => {
   const claimedToolCalls = new Set<string>();
+  const attemptsByToolCall = new Map<string, number>();
 
   return {
     claim(toolCallId: string) {
       if (claimedToolCalls.has(toolCallId)) return false;
+      const attempts = attemptsByToolCall.get(toolCallId) ?? 0;
+      if (attempts >= maximumAttempts) return false;
       claimedToolCalls.add(toolCallId);
+      attemptsByToolCall.set(toolCallId, attempts + 1);
       return true;
     },
     release(toolCallId: string) {
       claimedToolCalls.delete(toolCallId);
     },
+    complete(toolCallId: string) {
+      claimedToolCalls.delete(toolCallId);
+      attemptsByToolCall.delete(toolCallId);
+    },
     clear() {
       claimedToolCalls.clear();
+      attemptsByToolCall.clear();
     },
   };
 };
@@ -25,6 +37,8 @@ interface AssistantAutoApprovalContextValue {
   enabled: boolean;
   setEnabled: (enabled: boolean) => void;
   requestAction: (toolCallId: string, confirm: () => void) => boolean;
+  completeAction: (toolCallId: string) => void;
+  retryRevision: number;
 }
 
 const AssistantAutoApprovalContext = createContext<AssistantAutoApprovalContextValue | null>(null);
@@ -34,22 +48,41 @@ export const AssistantAutoApprovalProvider = ({
   conversationId,
 }: PropsWithChildren<{ conversationId: string }>) => {
   const [enabled, setEnabledState] = useState(false);
-  const trackerRef = useRef(createAssistantAutoApprovalTracker());
+  const [retryRevision, setRetryRevision] = useState(0);
+  const trackerRef = useRef(createAssistantAutoApprovalTracker(AUTO_APPROVAL_MAX_ATTEMPTS));
+  const retryTimersRef = useRef(new Map<string, number>());
   const previousConversationIdRef = useRef(conversationId);
 
-  const setEnabled = useCallback((nextEnabled: boolean) => {
-    if (!nextEnabled) trackerRef.current.clear();
-    setEnabledState(nextEnabled);
+  const clearRetryTimer = useCallback((toolCallId: string) => {
+    const timer = retryTimersRef.current.get(toolCallId);
+    if (timer !== undefined) window.clearTimeout(timer);
+    retryTimersRef.current.delete(toolCallId);
   }, []);
 
+  const clearAutomaticActions = useCallback(() => {
+    for (const timer of retryTimersRef.current.values()) window.clearTimeout(timer);
+    retryTimersRef.current.clear();
+    trackerRef.current.clear();
+  }, []);
+
+  const setEnabled = useCallback(
+    (nextEnabled: boolean) => {
+      if (!nextEnabled) clearAutomaticActions();
+      setEnabledState(nextEnabled);
+    },
+    [clearAutomaticActions],
+  );
+
   useEffect(() => {
     const previousConversationId = previousConversationIdRef.current;
     previousConversationIdRef.current = conversationId;
 
     if (previousConversationId === conversationId) return;
-    trackerRef.current.clear();
+    clearAutomaticActions();
     setEnabledState(false);
-  }, [conversationId]);
+  }, [clearAutomaticActions, conversationId]);
+
+  useEffect(() => clearAutomaticActions, [clearAutomaticActions]);
 
   const requestAction = useCallback(
     (toolCallId: string, confirm: () => void) => {
@@ -57,16 +90,36 @@ export const AssistantAutoApprovalProvider = ({
 
       try {
         confirm();
+        clearRetryTimer(toolCallId);
+        retryTimersRef.current.set(
+          toolCallId,
+          window.setTimeout(() => {
+            retryTimersRef.current.delete(toolCallId);
+            trackerRef.current.release(toolCallId);
+            setRetryRevision((revision) => revision + 1);
+          }, AUTO_APPROVAL_RETRY_TIMEOUT_MS),
+        );
         return true;
       } catch {
         trackerRef.current.release(toolCallId);
         return false;
       }
     },
-    [enabled],
+    [clearRetryTimer, enabled],
   );
 
-  const value = useMemo(() => ({ enabled, setEnabled, requestAction }), [enabled, requestAction, setEnabled]);
+  const completeAction = useCallback(
+    (toolCallId: string) => {
+      clearRetryTimer(toolCallId);
+      trackerRef.current.complete(toolCallId);
+    },
+    [clearRetryTimer],
+  );
+
+  const value = useMemo(
+    () => ({ enabled, setEnabled, requestAction, completeAction, retryRevision }),
+    [completeAction, enabled, requestAction, retryRevision, setEnabled],
+  );
 
   return <AssistantAutoApprovalContext.Provider value={value}>{children}</AssistantAutoApprovalContext.Provider>;
 };
@@ -88,20 +141,21 @@ export const useAssistantAutomaticAction = ({
   completed: boolean;
   confirm: () => void;
 }) => {
-  const { enabled, requestAction } = useAssistantAutoApproval();
+  const { completeAction, enabled, requestAction, retryRevision } = useAssistantAutoApproval();
   const [inProgress, setInProgress] = useState(false);
   const confirmRef = useRef(confirm);
   confirmRef.current = confirm;
 
   useEffect(() => {
     if (!ready || completed || !enabled) {
+      if (completed) completeAction(toolCallId);
       setInProgress(false);
       return;
     }
 
     const requested = requestAction(toolCallId, () => confirmRef.current());
     setInProgress(requested);
-  }, [completed, enabled, ready, requestAction, toolCallId]);
+  }, [completeAction, completed, enabled, ready, requestAction, retryRevision, toolCallId]);
 
   return inProgress;
 };

apps/nextjs/src/components/assistant/assistant-auto-submit.spec.ts

@@ -42,6 +42,38 @@ describe("shouldAutomaticallyContinueAssistant", () => {
     ).toBe(true);
   });
 
+  test("continues an approved mutation when the same step also reports a provider-executed server tool", () => {
+    expect(
+      shouldAutomaticallyContinueAssistant({
+        messages: [
+          {
+            id: "assistant-message",
+            role: "assistant",
+            parts: [
+              { type: "step-start" },
+              {
+                type: "dynamic-tool",
+                toolName: "openrouter_web_search",
+                toolCallId: "search-1",
+                input: { query: "PokéAPI" },
+                state: "input-available",
+                providerExecuted: true,
+              },
+              {
+                type: "dynamic-tool",
+                toolName: "customWidget_createFromPreview",
+                toolCallId: "create-1",
+                input: { previewSessionId: "preview-1" },
+                state: "approval-responded",
+                approval: { id: "approval-1", approved: true },
+              },
+            ],
+          },
+        ],
+      }),
+    ).toBe(true);
+  });
+
   test("waits while a human tool has no result", () => {
     expect(
       shouldAutomaticallyContinueAssistant({

apps/nextjs/src/components/assistant/assistant-auto-submit.ts

@@ -1,6 +1,35 @@
-import { lastAssistantMessageIsCompleteWithApprovalResponses, lastAssistantMessageIsCompleteWithToolCalls } from "ai";
+import {
+  isToolUIPart,
+  lastAssistantMessageIsCompleteWithApprovalResponses,
+  lastAssistantMessageIsCompleteWithToolCalls,
+} from "ai";
+import type { UIMessage } from "ai";
 
 type AutoSubmitOptions = Parameters<typeof lastAssistantMessageIsCompleteWithToolCalls>[0];
 
+const lastAssistantMessageHasCompleteLocalApprovalResponses = (messages: UIMessage[]) => {
+  const message = messages.at(-1);
+  if (message?.role !== "assistant") return false;
+
+  const lastStepStartIndex = message.parts.reduce(
+    (lastIndex, part, index) => (part.type === "step-start" ? index : lastIndex),
+    -1,
+  );
+  const localToolParts = message.parts
+    .slice(lastStepStartIndex + 1)
+    .filter(isToolUIPart)
+    .filter((part) => !part.providerExecuted);
+
+  return (
+    localToolParts.some((part) => part.state === "approval-responded") &&
+    localToolParts.every(
+      (part) =>
+        part.state === "output-available" || part.state === "output-error" || part.state === "approval-responded",
+    )
+  );
+};
+
 export const shouldAutomaticallyContinueAssistant = (options: AutoSubmitOptions) =>
-  lastAssistantMessageIsCompleteWithApprovalResponses(options) || lastAssistantMessageIsCompleteWithToolCalls(options);
+  lastAssistantMessageIsCompleteWithApprovalResponses(options) ||
+  lastAssistantMessageHasCompleteLocalApprovalResponses(options.messages) ||
+  lastAssistantMessageIsCompleteWithToolCalls(options);

apps/nextjs/src/components/assistant/assistant-human-tools.tsx

@@ -83,7 +83,6 @@ export const AssistantAskUserTool = ({
     completed: result !== undefined,
     confirm: () => {
       if (!affirmativeResult) return;
-      setSubmitting(true);
       addResult(affirmativeResult);
     },
   });
@@ -299,7 +298,6 @@ export const AssistantConfigureAppTool = ({
     completed: result !== undefined,
     confirm: () => {
       if (!initialValues) return;
-      setSubmitting(true);
       addResult(initialValues);
     },
   });

apps/nextjs/src/components/assistant/assistant-panel.tsx

@@ -868,7 +868,6 @@ const ToolPart = ({
     completed: !awaitingApproval,
     confirm: () => {
       respondToApproval({ approved: true, reason: "Approved automatically by the user." });
-      setApprovalResponse("approve");
     },
   });