diff --git a/apps/web/src/features/i18n/locales/en-US.json b/apps/web/src/features/i18n/locales/en-US.json index b4cc878dee..3996eaf2d2 100644 --- a/apps/web/src/features/i18n/locales/en-US.json +++ b/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", diff --git a/apps/web/src/features/shared/ai-image-generator/ai-image-generator.tsx b/apps/web/src/features/shared/ai-image-generator/ai-image-generator.tsx index 3725b8c266..e32865e414 100644 --- a/apps/web/src/features/shared/ai-image-generator/ai-image-generator.tsx +++ b/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(null); const [selectedPower, setSelectedPower] = useState(null); const [generatedUrl, setGeneratedUrl] = useState(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(null); + // Bounded auto-poll while the attempt is pending server-side. + const autoFetchTimerRef = useRef | 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 })} )} @@ -352,7 +441,11 @@ export function AiImageGenerator({ onInsert, showInsertAction = true, suggestedP {deliveryPending && !isGenerating && (
- {i18next.t("ai-image-generator.finishing")} + {i18next.t( + pendingPhase === "generating" + ? "ai-image-generator.still-generating" + : "ai-image-generator.finishing" + )}
)} diff --git a/apps/web/src/specs/features/ai-image-generator/pending-autopoll.spec.tsx b/apps/web/src/specs/features/ai-image-generator/pending-autopoll.spec.tsx new file mode 100644 index 0000000000..fd7764747e --- /dev/null +++ b/apps/web/src/specs/features/ai-image-generator/pending-autopoll.spec.tsx @@ -0,0 +1,196 @@ +import { createTestQueryClient, renderWithQueryClient } from "@/specs/test-utils"; +import { fireEvent, screen, waitFor } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi, type Mock } from "vitest"; + +vi.mock("@/utils", async () => ({ + ...(await vi.importActual("@/utils")), + random: vi.fn(), + getAccessToken: vi.fn(() => "mock-token"), + ensureValidToken: vi.fn(async () => "mock-token") +})); + +import { useActiveAccount } from "@/core/hooks/use-active-account"; +import { useGenerateImage } from "@ecency/sdk"; +import { AiImageGenerator } from "@/features/shared/ai-image-generator/ai-image-generator"; + +type HttpishError = Error & { status: number; data: Record }; + +// err shaped like the SDK mutation throws it: HTTP status + parsed body. +function httpError(status: number, data: Record): HttpishError { + return Object.assign(new Error(`failed with status ${status}`), { status, data }); +} + +function seededClient() { + const queryClient = createTestQueryClient(); + // Same key shape as the SDK's QueryKeys.ai.prices(). + queryClient.setQueryData(["ai", "prices"], { + prices: [{ aspect_ratio: "1:1", cost: 150 }], + power: [{ power: 1, multiplier: 1 }] + }); + return queryClient; +} + +const stillInProgress = (retryAfter: number) => + httpError(409, { error: "in_progress", retry_after: retryAfter }); + +describe("AiImageGenerator pending auto-poll", () => { + const generateMock = vi.fn(); + + beforeEach(() => { + generateMock.mockReset(); + (useGenerateImage as Mock).mockReturnValue({ + mutateAsync: generateMock, + isPending: false + }); + (useActiveAccount as Mock).mockReturnValue({ + activeUser: { username: "alice" }, + username: "alice" + }); + }); + + async function renderAndGenerate() { + const view = renderWithQueryClient(, { + queryClient: seededClient() + }); + fireEvent.change(screen.getByPlaceholderText("ai-image-generator.prompt-placeholder"), { + target: { value: "a fox" } + }); + fireEvent.click(screen.getByText("ai-image-generator.generate-button")); + await waitFor(() => expect(generateMock).toHaveBeenCalledTimes(1)); + return view; + } + + it("keeps the idempotency key on 409 in_progress and auto-polls until the image arrives", async () => { + generateMock + .mockRejectedValueOnce(stillInProgress(0.01)) + .mockResolvedValueOnce({ url: "https://images.test/done.png" }); + + await renderAndGenerate(); + + // The in-flight state is shown and the retry happens automatically with the SAME key. + await waitFor(() => expect(generateMock).toHaveBeenCalledTimes(2)); + expect(generateMock.mock.calls[1][0].idempotency_key).toBe( + generateMock.mock.calls[0][0].idempotency_key + ); + await waitFor(() => + expect(screen.getByText("ai-image-generator.result-title")).toBeTruthy() + ); + }); + + it("shows the still-generating notice while a 409 attempt is pending", async () => { + // Bottomless in_progress: every poll answers 409 again. + generateMock.mockRejectedValue(stillInProgress(60)); + + await renderAndGenerate(); + + await waitFor(() => + expect(screen.getByText("ai-image-generator.still-generating")).toBeTruthy() + ); + }); + + it("keeps the finishing notice for a 202 delivery_pending answer", async () => { + generateMock.mockRejectedValue( + httpError(202, { error: "delivery_pending", retry_after: 60 }) + ); + + await renderAndGenerate(); + + await waitFor(() => + expect(screen.getByText("ai-image-generator.finishing")).toBeTruthy() + ); + }); + + it("drops the key on a hard failure so the next attempt is a fresh request", async () => { + generateMock + .mockRejectedValueOnce(httpError(500, { error: "generation_failed" })) + .mockRejectedValueOnce(httpError(500, { error: "generation_failed" })); + + await renderAndGenerate(); + + fireEvent.click(screen.getByText("ai-image-generator.generate-button")); + await waitFor(() => expect(generateMock).toHaveBeenCalledTimes(2)); + expect(generateMock.mock.calls[1][0].idempotency_key).not.toBe( + generateMock.mock.calls[0][0].idempotency_key + ); + }); + + it("drops a pending response whose inputs changed mid-flight instead of polling the new ones", async () => { + let rejectFirst!: (e: unknown) => void; + generateMock.mockImplementationOnce( + () => new Promise((_resolve, reject) => (rejectFirst = reject)) + ); + + await renderAndGenerate(); + + // The user edits the prompt while the request is unsettled, abandoning the attempt. + fireEvent.change(screen.getByPlaceholderText("ai-image-generator.prompt-placeholder"), { + target: { value: "a different fox" } + }); + rejectFirst(stillInProgress(0.01)); + + // Without the stale-attempt gate the 409 would arm a poll that auto-submits the NEW + // prompt under a fresh key: an unrequested, billed generation. + await new Promise((resolve) => setTimeout(resolve, 150)); + expect(generateMock).toHaveBeenCalledTimes(1); + expect(screen.queryByText("ai-image-generator.still-generating")).toBeNull(); + }); + + it("stops polling when the dialog unmounts mid-request", async () => { + let rejectFirst!: (e: unknown) => void; + generateMock.mockImplementationOnce( + () => new Promise((_resolve, reject) => (rejectFirst = reject)) + ); + + const view = await renderAndGenerate(); + view.unmount(); + rejectFirst(stillInProgress(0.01)); + + // The resolution of an unmounted dialog must not schedule further polls. + await new Promise((resolve) => setTimeout(resolve, 150)); + expect(generateMock).toHaveBeenCalledTimes(1); + }); + + it("replaces an armed poll timer on manual fetch and never overlaps requests", async () => { + let resolveSecond!: (v: unknown) => void; + generateMock + .mockRejectedValueOnce(stillInProgress(0.2)) + .mockImplementationOnce(() => new Promise((resolve) => (resolveSecond = resolve))); + + await renderAndGenerate(); + + // Manual fetch while the poll timer is armed. Whichever of the two fires first, the + // other must be swallowed: exactly one request runs at a time. + await waitFor(() => + expect(screen.getByText("ai-image-generator.finishing-retry")).toBeTruthy() + ); + fireEvent.click(screen.getByText("ai-image-generator.finishing-retry")); + await waitFor(() => expect(generateMock).toHaveBeenCalledTimes(2)); + + // The timer window passes while the second request hangs: no concurrent third call. + await new Promise((resolve) => setTimeout(resolve, 350)); + expect(generateMock).toHaveBeenCalledTimes(2); + + resolveSecond({ url: "https://images.test/done.png" }); + await waitFor(() => + expect(screen.getByText("ai-image-generator.result-title")).toBeTruthy() + ); + }); + + it("stops automatic polling at the budget and keeps the key for the manual fetch", async () => { + generateMock.mockRejectedValue(stillInProgress(0.001)); + + await renderAndGenerate(); + + // 1 initial request + 24 automatic polls, then the budget is spent. + await waitFor(() => expect(generateMock).toHaveBeenCalledTimes(25), { timeout: 15000 }); + await new Promise((resolve) => setTimeout(resolve, 150)); + expect(generateMock).toHaveBeenCalledTimes(25); + + // The manual button remains and still replays the SAME paid attempt. + fireEvent.click(screen.getByText("ai-image-generator.finishing-retry")); + await waitFor(() => expect(generateMock).toHaveBeenCalledTimes(26)); + expect(generateMock.mock.calls[25][0].idempotency_key).toBe( + generateMock.mock.calls[0][0].idempotency_key + ); + }, 20000); +}); diff --git a/apps/web/src/specs/setup-any-spec.ts b/apps/web/src/specs/setup-any-spec.ts index 340558807e..320a3f6409 100644 --- a/apps/web/src/specs/setup-any-spec.ts +++ b/apps/web/src/specs/setup-any-spec.ts @@ -88,6 +88,10 @@ vi.mock("@ecency/sdk", async () => ({ })), getBoostPlusPricesQueryOptions: vi.fn(() => ({ queryKey: ["boost-prices"], queryFn: vi.fn() })), getPointsQueryOptions: vi.fn(() => ({ queryKey: ["points"], queryFn: vi.fn() })), + // Key shape matches the SDK's QueryKeys.ai.prices(). + getAiGeneratePriceQueryOptions: vi.fn(() => ({ queryKey: ["ai", "prices"], queryFn: vi.fn() })), + useGenerateImage: vi.fn(() => ({ mutateAsync: vi.fn(), isPending: false })), + useAddImage: vi.fn(() => ({ mutateAsync: vi.fn(async () => ({})) })), getProMembersQueryOptions: vi.fn(() => ({ queryKey: ["accounts", "pro-members"], queryFn: vi.fn() })), getPostTipsQueryOptions: vi.fn((author: string, permlink: string) => ({ queryKey: ["posts", "tips", author, permlink], queryFn: vi.fn() })), getTrendingTagsQueryOptions: vi.fn((limit?: number) => ({