fix(ai-image): keep the idempotency key on 409 in_progress and auto-poll for the finished image

apps/web/src/features/i18n/locales/en-US.json

@@ -4058,7 +4058,8 @@
     "use-suggestion": "Use: {{suggestion}}",
     "select-power": "Quality boost",
     "finishing": "Your image is generated and finishing upload. Tap below to fetch it — you won't be charged again.",
-    "finishing-retry": "Fetch image"
+    "finishing-retry": "Fetch image",
+    "still-generating": "Your image is still generating. It will be fetched automatically with no extra charge."
   },
   "ai-usage": {
     "menu": "AI usage",

apps/web/src/features/shared/ai-image-generator/ai-image-generator.tsx

@@ -51,6 +51,12 @@ function makeIdempotencyKey(): string {
   return `k${Date.now().toString(36)}${Math.random().toString(36).slice(2, 12)}`;
 }
 
+// Bounded automatic polling for a generation the backend answered 202/409 for: it is paid
+// (or still running) server-side and a retry with the SAME key only fetches it. Past the
+// budget the manual button remains, still carrying the key.
+const AUTO_FETCH_MAX_ATTEMPTS = 24;
+const AUTO_FETCH_DEFAULT_DELAY_S = 5;
+
 export function AiImageGenerator({ onInsert, showInsertAction = true, suggestedPrompt }: Props) {
   const { activeUser } = useActiveAccount();
   const username = activeUser?.username;
@@ -82,15 +88,48 @@ export function AiImageGenerator({ onInsert, showInsertAction = true, suggestedP
   const [selectedRatio, setSelectedRatio] = useState<string | null>(null);
   const [selectedPower, setSelectedPower] = useState<AiImagePowerTier | null>(null);
   const [generatedUrl, setGeneratedUrl] = useState<string | null>(null);
-  // Set when the backend returns 202: the image is paid for and finishing upload; the next
-  // Generate click retries with the same key to fetch it (no new charge).
-  const [deliveryPending, setDeliveryPending] = useState(false);
+  // Set while the backend reports the attempt as still in flight: "generating" for a 409
+  // in_progress (prediction still running), "finishing" for a 202 delivery_pending (paid,
+  // upload pending). Either way the SAME key fetches it with no new charge.
+  const [pendingPhase, setPendingPhase] = useState<"generating" | "finishing" | null>(null);
+  const deliveryPending = pendingPhase !== null;
 
   // Idempotency key for the current attempt. Kept stable across a delivery-pending retry so
   // the backend recovers the paid generation; reset whenever the inputs change (below) so a
   // genuinely new request gets a fresh key.
   const idempotencyKeyRef = useRef<string | null>(null);
 
+  // Bounded auto-poll while the attempt is pending server-side.
+  const autoFetchTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+  const autoFetchAttemptsRef = useRef(0);
+
+  // Identifies the attempt (user + inputs + key) that the outstanding request and any
+  // armed timer belong to. Bumped whenever those change, and on unmount, so a resolution
+  // or timer from an earlier attempt is dropped instead of acting on state it no longer
+  // owns: an old 202/409 must never schedule a poll that would submit the NEW inputs
+  // under a fresh key (an unrequested, billed generation) or bill another account.
+  const attemptRef = useRef(0);
+  // Synchronous re-entrancy guard: a timer firing around a manual fetch must not start a
+  // second concurrent request racing the same key's pending/success state.
+  const inFlightRef = useRef(false);
+
+  const clearAutoFetch = useCallback(() => {
+    if (autoFetchTimerRef.current) {
+      clearTimeout(autoFetchTimerRef.current);
+      autoFetchTimerRef.current = null;
+    }
+  }, []);
+
+  useEffect(
+    () => () => {
+      // Unmount invalidates the in-flight resolution too, not just the armed timer --
+      // otherwise its 202/409 handler re-arms and keeps polling a dead dialog.
+      attemptRef.current += 1;
+      clearAutoFetch();
+    },
+    [clearAutoFetch]
+  );
+
   useEffect(() => {
     if (prices && prices.length > 0 && !selectedRatio) {
       const defaultRatio = prices.find((p) => p.aspect_ratio === "16:9");
@@ -104,11 +143,15 @@ export function AiImageGenerator({ onInsert, showInsertAction = true, suggestedP
     }
   }, [powerTiers, selectedPower]);
 
-  // Changing any request input starts a fresh attempt: drop the reused key + pending state.
+  // Changing any request input (or the account) starts a fresh attempt: invalidate the
+  // outstanding request/timer and drop the reused key + pending state.
   useEffect(() => {
+    attemptRef.current += 1;
     idempotencyKeyRef.current = null;
-    setDeliveryPending(false);
-  }, [prompt, selectedRatio, selectedPower]);
+    autoFetchAttemptsRef.current = 0;
+    clearAutoFetch();
+    setPendingPhase(null);
+  }, [prompt, selectedRatio, selectedPower, username, clearAutoFetch]);
 
   const charsRemaining = 1000 - prompt.length;
 
@@ -136,9 +179,21 @@ export function AiImageGenerator({ onInsert, showInsertAction = true, suggestedP
     (deliveryPending || !isInsufficientBalance) &&
     !isGenerating;
 
+  // Always points at the latest handleGenerate so an auto-poll timer never fires a stale
+  // closure (e.g. one holding an outdated points balance).
+  const handleGenerateRef = useRef<() => void>(() => {});
+
   const handleGenerate = useCallback(async () => {
     if (!selectedRatio || !prompt.trim()) return;
 
+    // A manual click replaces any armed poll timer, and nothing ever runs two requests
+    // for the same attempt concurrently -- overlapping completions would race the
+    // pending/success state and duplicate the gallery insert.
+    clearAutoFetch();
+    if (inFlightRef.current) return;
+    inFlightRef.current = true;
+    const attempt = attemptRef.current;
+
     try {
       const token = username ? await ensureValidToken(username) : undefined;
       if (!token) {
@@ -159,26 +214,55 @@ export function AiImageGenerator({ onInsert, showInsertAction = true, suggestedP
         idempotency_key: idempotencyKeyRef.current,
       });
 
-      setDeliveryPending(false);
+      // Inputs, account, or mount state changed while this request was in flight: the
+      // response belongs to an abandoned attempt, so drop it without touching state.
+      if (attempt !== attemptRef.current) return;
+
+      setPendingPhase(null);
       idempotencyKeyRef.current = null;
+      autoFetchAttemptsRef.current = 0;
+      clearAutoFetch();
       setGeneratedUrl(result.url);
       success(i18next.t("ai-image-generator.success"));
 
       // Auto-add to user's gallery (non-blocking)
       addToGallery({ url: result.url, code: token }).catch(() => {});
     } catch (err: any) {
+      // Same stale-attempt gate as the success path: this failure (or pending answer)
+      // belongs to an abandoned attempt and must not touch state or schedule a poll.
+      if (attempt !== attemptRef.current) return;
+
       const status = err?.status;
       const data = err?.data;
 
-      if (status === 202) {
-        // Image is paid for and finishing upload — keep the key so the next click fetches it.
-        setDeliveryPending(true);
+      // 202 = paid, upload finishing. 409 in_progress = the prediction is still running
+      // server-side. Both mean: keep the key (a retry only fetches, never re-bills) and
+      // poll automatically at the backend's suggested cadence, up to a bounded budget.
+      const stillInFlight =
+        status === 202 || (status === 409 && data?.error === "in_progress");
+      if (stillInFlight) {
+        setPendingPhase(status === 202 ? "finishing" : "generating");
+        if (autoFetchAttemptsRef.current < AUTO_FETCH_MAX_ATTEMPTS) {
+          autoFetchAttemptsRef.current += 1;
+          const delayS =
+            typeof data?.retry_after === "number" && data.retry_after > 0
+              ? data.retry_after
+              : AUTO_FETCH_DEFAULT_DELAY_S;
+          clearAutoFetch();
+          autoFetchTimerRef.current = setTimeout(() => {
+            autoFetchTimerRef.current = null;
+            if (attempt !== attemptRef.current) return;
+            handleGenerateRef.current();
+          }, delayS * 1000);
+        }
         return;
       }
 
       // Any hard failure ends this attempt: drop the key so a retry is a fresh request.
       idempotencyKeyRef.current = null;
-      setDeliveryPending(false);
+      autoFetchAttemptsRef.current = 0;
+      clearAutoFetch();
+      setPendingPhase(null);
 
       if (status === 402) {
         error(
@@ -194,8 +278,13 @@ export function AiImageGenerator({ onInsert, showInsertAction = true, suggestedP
       } else {
         error(i18next.t("ai-image-generator.error-generic"));
       }
+    } finally {
+      inFlightRef.current = false;
     }
-  }, [selectedRatio, selectedPower, prompt, username, generateImage, addToGallery, cost]);
+  }, [selectedRatio, selectedPower, prompt, username, generateImage, addToGallery, cost,
+      clearAutoFetch]);
+
+  handleGenerateRef.current = handleGenerate;
 
   const handleGenerateAgain = useCallback(() => {
     setGeneratedUrl(null);
@@ -280,7 +369,7 @@ export function AiImageGenerator({ onInsert, showInsertAction = true, suggestedP
             onClick={() => setPrompt(suggestedPrompt.slice(0, 1000))}
           >
             {i18next.t("ai-image-generator.use-suggestion", {
-              suggestion: suggestedPrompt.length > 80 ? suggestedPrompt.slice(0, 80) + "\u2026" : suggestedPrompt
+              suggestion: suggestedPrompt.length > 80 ? suggestedPrompt.slice(0, 80) + "…" : suggestedPrompt
             })}
           </button>
         )}
@@ -352,7 +441,11 @@ export function AiImageGenerator({ onInsert, showInsertAction = true, suggestedP
 
       {deliveryPending && !isGenerating && (
         <div className="text-sm rounded-lg border border-blue-dark-sky/30 bg-blue-dark-sky/5 text-blue-dark-sky px-3 py-2">
-          {i18next.t("ai-image-generator.finishing")}
+          {i18next.t(
+            pendingPhase === "generating"
+              ? "ai-image-generator.still-generating"
+              : "ai-image-generator.finishing"
+          )}
         </div>
       )}