diff --git a/frontend/editor/src/core/contexts/FileContext.tsx b/frontend/editor/src/core/contexts/FileContext.tsx index 45847cbe98..08f867f7b6 100644 --- a/frontend/editor/src/core/contexts/FileContext.tsx +++ b/frontend/editor/src/core/contexts/FileContext.tsx @@ -69,6 +69,7 @@ import type { RemovePasswordParameters } from "@app/hooks/tools/removePassword/u import { useResolutionContinuation } from "@app/hooks/tools/shared/useResolutionContinuation"; import apiClient from "@app/services/apiClient"; import { reportFilesRemoved } from "@app/services/failureReporting"; +import { setPendingUnlocks } from "@app/services/pendingUnlocks"; import { processResponse } from "@app/utils/toolResponseProcessor"; import { ToolOperation } from "@app/types/file"; import { handlePasswordError } from "@app/utils/toolErrorHandler"; @@ -185,6 +186,17 @@ function FileContextInner({ } }, [activeEncryptedFileId, state.files.ids]); + // Published so an upload policy holds off until the user has answered the prompt: running now + // would fail on a document they are about to decrypt, and leave a row about a version that no + // longer exists once they have. + useEffect(() => { + setPendingUnlocks( + activeEncryptedFileId + ? [activeEncryptedFileId, ...encryptedQueue] + : encryptedQueue, + ); + }, [activeEncryptedFileId, encryptedQueue]); + useEffect(() => { setUnlockPassword(""); setUnlockError(null); diff --git a/frontend/editor/src/core/services/pendingUnlocks.ts b/frontend/editor/src/core/services/pendingUnlocks.ts new file mode 100644 index 0000000000..0a6bc61a4d --- /dev/null +++ b/frontend/editor/src/core/services/pendingUnlocks.ts @@ -0,0 +1,44 @@ +/** + * The uploads still waiting on their unlock prompt. Published by the workbench, read by anything + * that would otherwise act on a document the user is in the middle of decrypting. + * + * A module store rather than context: the reader is a policy hook in another layer, and it needs + * the answer during an effect rather than as a render input. + */ + +const pending = new Set(); +const listeners = new Set<() => void>(); + +/** Replaces the set wholesale, since the prompt queue is authoritative about who is waiting. */ +export function setPendingUnlocks(fileIds: readonly string[]): void { + const next = new Set(fileIds); + if (next.size === pending.size && [...next].every((id) => pending.has(id))) { + return; + } + pending.clear(); + for (const id of next) pending.add(id); + version += 1; + for (const listener of listeners) listener(); +} + +/** + * Whether this document is still awaiting an unlock decision. False once the user has unlocked it + * (the document is replaced by a decrypted version) or skipped it (they have chosen to go on). + */ +export function isAwaitingUnlock(fileId: string): boolean { + return pending.has(fileId); +} + +let version = 0; + +/** Changes whenever the set does, so a subscriber can re-run work it skipped. */ +export function pendingUnlocksVersion(): number { + return version; +} + +export function subscribeToPendingUnlocks(listener: () => void): () => void { + listeners.add(listener); + return () => { + listeners.delete(listener); + }; +} diff --git a/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.pendingUnlock.test.tsx b/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.pendingUnlock.test.tsx new file mode 100644 index 0000000000..a50bc5ed5e --- /dev/null +++ b/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.pendingUnlock.test.tsx @@ -0,0 +1,115 @@ +/** + * An encrypted upload opens the unlock prompt and dispatches its policy in the same tick, so the + * two race. If the user wins, the run fails on a document they have already replaced, billing for + * it and leaving a row about a version that no longer exists. The dispatch waits for the answer. + */ +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { renderHook, act } from "@testing-library/react"; + +vi.mock("@app/hooks/useAiEngineEnabled", () => ({ + useAiEngineEnabled: () => false, +})); +const fileStubs: { id: string; name: string; derivedFromTool?: boolean }[] = []; +vi.mock("@app/contexts/FileContext", () => ({ + useAllFiles: () => ({ fileStubs }), + useFileManagement: () => ({ addFiles: vi.fn() }), + useFileContext: () => ({ consumeFiles: vi.fn() }), +})); +vi.mock("@app/hooks/usePolicies", () => ({ + usePolicies: () => ({ + policies: { + security: { + configured: true, + status: "active", + backendId: "backend-sec", + runOn: "upload", + order: 0, + }, + }, + }), +})); +vi.mock("@app/services/policyApi", () => ({ + runStoredPolicy: vi.fn(), + getPolicyRun: vi.fn(), + downloadPolicyOutput: vi.fn(), + resolvePolicyRunTarget: () => "saas", +})); +vi.mock("@app/services/fileStorage", () => ({ + fileStorage: { getStirlingFile: vi.fn(), getStirlingFileStub: vi.fn() }, +})); +vi.mock("@app/contexts/IndexedDBContext", () => ({ + useIndexedDB: () => ({ bumpRevision: vi.fn() }), +})); + +import { usePolicyAutoRun } from "@app/components/policies/usePolicyAutoRun"; +import { resetPolicyRuns } from "@app/components/policies/policyRunStore"; +import { setPendingUnlocks } from "@app/services/pendingUnlocks"; +import { runStoredPolicy } from "@app/services/policyApi"; +import { fileStorage } from "@app/services/fileStorage"; + +const runStored = vi.mocked(runStoredPolicy); +const getFile = vi.mocked(fileStorage.getStirlingFile); + +function setFileStubs(next: typeof fileStubs) { + fileStubs.length = 0; + fileStubs.push(...next); +} + +beforeEach(() => { + vi.useFakeTimers(); + resetPolicyRuns(); + setPendingUnlocks([]); + runStored.mockReset().mockResolvedValue("run-sec"); + getFile.mockReset().mockResolvedValue({ size: 1460 } as never); + setFileStubs([]); +}); +// Cleared in beforeEach, not here: notifying a component RTL has not unmounted yet +// would be a state update outside act. +afterEach(() => vi.useRealTimers()); + +async function settle() { + await act(async () => { + await vi.advanceTimersByTimeAsync(1); + }); +} + +describe("an upload waiting on its unlock prompt", () => { + it("holds the run back while the prompt is open", async () => { + setPendingUnlocks(["file-locked"]); + setFileStubs([{ id: "file-locked", name: "locked.pdf" }]); + + renderHook(() => usePolicyAutoRun()); + await settle(); + + expect(runStored).not.toHaveBeenCalled(); + }); + + it("runs once the prompt is answered, so skipping still records the failure", async () => { + // Skipping releases the file encrypted: the run fails, and that failure is the row the + // bell offers Decrypt and retry on. Holding it back forever would lose that entirely. + setPendingUnlocks(["file-locked"]); + setFileStubs([{ id: "file-locked", name: "locked.pdf" }]); + + renderHook(() => usePolicyAutoRun()); + await settle(); + expect(runStored).not.toHaveBeenCalled(); + + act(() => setPendingUnlocks([])); + await settle(); + + expect(runStored).toHaveBeenCalledWith( + "backend-sec", + expect.anything(), + "file-locked", + ); + }); + + it("leaves an upload nobody is prompting about alone", async () => { + setFileStubs([{ id: "file-plain", name: "plain.pdf" }]); + + renderHook(() => usePolicyAutoRun()); + await settle(); + + expect(runStored).toHaveBeenCalledTimes(1); + }); +}); diff --git a/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.ts b/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.ts index 68a70cfd8c..1e9f10b524 100644 --- a/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.ts +++ b/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.ts @@ -3,7 +3,13 @@ * Policies sharing a trigger run as an ordered chain so their effects accumulate. */ -import { useCallback, useEffect, useMemo, useRef } from "react"; +import { + useCallback, + useEffect, + useMemo, + useRef, + useSyncExternalStore, +} from "react"; import { useAllFiles, useFileManagement, @@ -11,6 +17,11 @@ import { } from "@app/contexts/FileContext"; import { fileStorage } from "@app/services/fileStorage"; import { refreshNotificationsNow } from "@app/hooks/useNotifications"; +import { + isAwaitingUnlock, + pendingUnlocksVersion, + subscribeToPendingUnlocks, +} from "@app/services/pendingUnlocks"; import { useIndexedDB } from "@app/contexts/IndexedDBContext"; import i18n from "@app/i18n"; import { @@ -168,6 +179,14 @@ export function usePolicyAutoRun(): void { // Chain-continuations handled this session, so the next policy fires once per run. const chained = useRef>(new Set()); + // Answering a prompt has to re-run the dispatch effect, or a released file waits for the + // next unrelated render to be picked up. + const unlocksVersion = useSyncExternalStore( + subscribeToPendingUnlocks, + pendingUnlocksVersion, + pendingUnlocksVersion, + ); + // Latest policies, read from inside the stable retry callback (which has no deps). const policiesRef = useRef(policies); policiesRef.current = policies; @@ -245,6 +264,10 @@ export function usePolicyAutoRun(): void { // Input-mode policies cover uploads only; tool-produced files are left to // export-mode policies at export time. if (stub.derivedFromTool) continue; + // Held while the unlock prompt is open: the run would fail on a document the user is + // about to decrypt, bill for it, and leave a row about a version soon replaced. Skipping + // the prompt releases it, so a document nobody unlocks still records its failure. + if (isAwaitingUnlock(stub.id)) continue; const key = dispatchKey(firstCategory, stub.id); // Skip if already run (persisted) or in flight - the in-memory guard covers the async wait. if ( @@ -271,7 +294,7 @@ export function usePolicyAutoRun(): void { }) .finally(() => dispatching.current.delete(key)); } - }, [fileStubs, policies, orderedUploadCategories]); + }, [fileStubs, policies, orderedUploadCategories, unlocksVersion]); // Once a run's output lands, fire the next upload policy on it - success only, once per // run. isDispatched guards re-dispatch across reloads.