Redesign policy running (#6609)

# Description of Changes
Redesign policy running so the server is in charge of policy IDs and
running, to make it impossible to have the frontend miss the results.
This solves a minor bug that we currently have in policies, where if you
load a file and then refresh while the policy is running, you'll never
receive the outputted file.

frontend/editor/src/proprietary/components/policies/policyRunStore.ts

@@ -133,6 +133,18 @@ export function recordRunStart(record: PolicyRunRecord) {
   emit();
 }
 
+/**
+ * Add a run discovered on the backend that this client has no local record of (e.g. it was
+ * started before a refresh recorded it). Unlike {@link recordRunStart} this adds no dispatch key:
+ * the run already exists server-side, so re-dispatch isn't the concern; we only want it polled and
+ * its outputs imported. No-op if the run is already known (reconcile patches those via updateRun).
+ */
+export function addReconciledRun(record: PolicyRunRecord) {
+  if (state.runs.some((r) => r.runId === record.runId)) return;
+  state = { ...state, runs: [record, ...state.runs].slice(0, MAX_RUNS) };
+  emit();
+}
+
 /** Mark a (policy, file) pair dispatched without a run (e.g. dispatch failed). */
 export function markDispatched(categoryId: string, fileId: string) {
   const key = dispatchKey(categoryId, fileId);

frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.import.test.tsx

@@ -0,0 +1,186 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { renderHook, act } from "@testing-library/react";
+
+// Drive a COMPLETED run's import against the REAL run store, stubbing the contexts/network the
+// hook reaches into. Spies are hoisted so the vi.mock factories share stable references we can
+// assert on. The focus is which delivery path importOutputs takes for "new_version" output:
+//  - input in the active workspace  → consumeFiles (versions in place)
+//  - input only in storage (reload) → fileStorage.persistVersionedOutputs (+ bumpRevision)
+//  - input gone entirely            → addFiles (new file)
+const mocks = vi.hoisted(() => ({
+  fileStubs: [] as Array<{ id: string }>,
+  addFiles: vi.fn(),
+  consumeFiles: vi.fn(),
+  bumpRevision: vi.fn(),
+  persistVersionedOutputs: vi.fn(),
+  getStirlingFile: vi.fn(),
+  getStirlingFileStub: vi.fn(),
+  downloadPolicyOutput: vi.fn(),
+  listPolicyRuns: vi.fn(),
+  createStirlingFilesAndStubs: vi.fn(),
+}));
+
+vi.mock("@app/constants/featureFlags", () => ({ POLICIES_ENABLED: true }));
+vi.mock("@app/contexts/FileContext", () => ({
+  useAllFiles: () => ({ fileStubs: mocks.fileStubs }),
+  useFileManagement: () => ({ addFiles: mocks.addFiles }),
+  useFileContext: () => ({ consumeFiles: mocks.consumeFiles }),
+}));
+vi.mock("@app/contexts/IndexedDBContext", () => ({
+  useIndexedDB: () => ({ bumpRevision: mocks.bumpRevision }),
+}));
+vi.mock("@app/hooks/usePolicies", () => ({
+  usePolicies: () => ({
+    policies: {
+      security: {
+        configured: true,
+        status: "active",
+        backendId: "backend-1",
+        runOn: "upload",
+        outputMode: "new_version",
+        outputName: "",
+      },
+    },
+  }),
+}));
+vi.mock("@app/services/policyApi", () => ({
+  runStoredPolicy: vi.fn(),
+  getPolicyRun: vi.fn(),
+  listPolicyRuns: mocks.listPolicyRuns,
+  downloadPolicyOutput: mocks.downloadPolicyOutput,
+}));
+vi.mock("@app/services/fileStorage", () => ({
+  fileStorage: {
+    getStirlingFile: mocks.getStirlingFile,
+    getStirlingFileStub: mocks.getStirlingFileStub,
+    persistVersionedOutputs: mocks.persistVersionedOutputs,
+  },
+}));
+vi.mock("@app/services/fileStubHelpers", () => ({
+  createStirlingFilesAndStubs: mocks.createStirlingFilesAndStubs,
+}));
+
+import { usePolicyAutoRun } from "@app/components/policies/usePolicyAutoRun";
+import {
+  recordRunStart,
+  getRun,
+  resetPolicyRuns,
+} from "@app/components/policies/policyRunStore";
+
+/** Seed a finished run (one output) so the import effect fires on mount. */
+function recordCompletedRun() {
+  recordRunStart({
+    runId: "run-1",
+    categoryId: "security",
+    fileId: "file-1",
+    fileName: "doc.pdf",
+    fileSize: 1234,
+    status: "COMPLETED",
+    outputs: [{ fileId: "out-file-1", fileName: "doc.pdf" }],
+    error: null,
+    startedAt: 0,
+  });
+}
+
+/** Render the hook and wait for the run's outputs to finish importing. */
+async function runImport() {
+  renderHook(() => usePolicyAutoRun());
+  await act(async () => {
+    await vi.waitFor(() => expect(getRun("run-1")?.imported).toBe(true));
+  });
+}
+
+beforeEach(() => {
+  localStorage.clear();
+  resetPolicyRuns();
+  vi.clearAllMocks();
+  mocks.fileStubs = [];
+  mocks.listPolicyRuns.mockResolvedValue([]);
+  mocks.getStirlingFileStub.mockResolvedValue(null);
+  mocks.persistVersionedOutputs.mockResolvedValue(undefined);
+  mocks.consumeFiles.mockResolvedValue(undefined);
+  mocks.addFiles.mockResolvedValue([{ fileId: "out-1" }]);
+  mocks.downloadPolicyOutput.mockResolvedValue(
+    new Blob(["x"], { type: "application/pdf" }),
+  );
+  mocks.createStirlingFilesAndStubs.mockResolvedValue({
+    stirlingFiles: [{ name: "doc.pdf" }],
+    stubs: [{ id: "out-1" }],
+  });
+});
+
+describe("auto-run import: new-version output delivery", () => {
+  it("versions the input in storage when it's recovered after a reload (no second file)", async () => {
+    // Reload case: the workspace is empty, but the input still persists in IndexedDB.
+    mocks.fileStubs = [];
+    mocks.getStirlingFileStub.mockResolvedValue({
+      id: "file-1",
+      versionNumber: 1,
+    });
+
+    recordCompletedRun();
+    await runImport();
+
+    expect(mocks.persistVersionedOutputs).toHaveBeenCalledWith(
+      ["file-1"],
+      expect.any(Array),
+      expect.any(Array),
+    );
+    expect(mocks.bumpRevision).toHaveBeenCalled();
+    expect(mocks.consumeFiles).not.toHaveBeenCalled();
+    expect(mocks.addFiles).not.toHaveBeenCalled();
+  });
+
+  it("versions the input in the workspace when it's open (consumeFiles, not a storage write)", async () => {
+    mocks.fileStubs = [{ id: "file-1" }];
+
+    recordCompletedRun();
+    await runImport();
+
+    expect(mocks.consumeFiles).toHaveBeenCalledWith(
+      ["file-1"],
+      expect.any(Array),
+      expect.any(Array),
+    );
+    expect(mocks.persistVersionedOutputs).not.toHaveBeenCalled();
+    expect(mocks.addFiles).not.toHaveBeenCalled();
+  });
+
+  it("falls back to adding a new file when the input is gone from storage too", async () => {
+    mocks.fileStubs = [];
+    mocks.getStirlingFileStub.mockResolvedValue(null);
+
+    recordCompletedRun();
+    await runImport();
+
+    expect(mocks.addFiles).toHaveBeenCalled();
+    expect(mocks.persistVersionedOutputs).not.toHaveBeenCalled();
+    expect(mocks.consumeFiles).not.toHaveBeenCalled();
+  });
+
+  it("adopts a server-only run, dating it from the server's createdAt (not now)", async () => {
+    // A run the client never recorded (true orphan): reconcile adopts it from the server. With no
+    // local input link it delivers as a new file, and its age comes from the server, not Date.now().
+    mocks.listPolicyRuns.mockResolvedValue([
+      {
+        runId: "srv-1",
+        policyId: "backend-1",
+        status: "COMPLETED",
+        currentStep: 1,
+        stepCount: 1,
+        error: null,
+        outputs: [{ fileId: "out-file-1", fileName: "doc.pdf" }],
+        createdAt: 1000,
+      },
+    ]);
+
+    renderHook(() => usePolicyAutoRun());
+    await act(async () => {
+      await vi.waitFor(() => expect(getRun("srv-1")?.imported).toBe(true));
+    });
+
+    expect(getRun("srv-1")?.startedAt).toBe(1000);
+    expect(mocks.addFiles).toHaveBeenCalled();
+    expect(mocks.persistVersionedOutputs).not.toHaveBeenCalled();
+  });
+});

frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.retry.test.tsx

@@ -27,7 +27,10 @@ vi.mock("@app/services/policyApi", () => ({
   downloadPolicyOutput: vi.fn(),
 }));
 vi.mock("@app/services/fileStorage", () => ({
-  fileStorage: { getStirlingFile: vi.fn() },
+  fileStorage: { getStirlingFile: vi.fn(), getStirlingFileStub: vi.fn() },
+}));
+vi.mock("@app/contexts/IndexedDBContext", () => ({
+  useIndexedDB: () => ({ bumpRevision: vi.fn() }),
 }));
 
 import { usePolicyAutoRun } from "@app/components/policies/usePolicyAutoRun";

frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.ts

@@ -18,10 +18,12 @@ import {
   useFileContext,
 } from "@app/contexts/FileContext";
 import { fileStorage } from "@app/services/fileStorage";
+import { useIndexedDB } from "@app/contexts/IndexedDBContext";
 import { POLICIES_ENABLED } from "@app/constants/featureFlags";
 import {
   runStoredPolicy,
   getPolicyRun,
+  listPolicyRuns,
   downloadPolicyOutput,
 } from "@app/services/policyApi";
 import type {
@@ -32,8 +34,10 @@ import { dispatchPaygLimitReached } from "@app/services/usageLimitBridge";
 import type { FileId } from "@app/types/file";
 import { createStirlingFilesAndStubs } from "@app/services/fileStubHelpers";
 import type { StirlingFile, StirlingFileStub } from "@app/types/fileContext";
+import type { PoliciesByCategory } from "@app/types/policies";
 import { usePolicies } from "@app/hooks/usePolicies";
 import {
+  addReconciledRun,
   dispatchKey,
   getRun,
   isDispatched,
@@ -116,6 +120,7 @@ export function usePolicyAutoRun(): void {
   const { fileStubs } = useAllFiles();
   const { addFiles } = useFileManagement();
   const { consumeFiles } = useFileContext();
+  const { bumpRevision } = useIndexedDB();
   const { policies } = usePolicies();
   const runs = usePolicyRuns();
   // Keys (run ids / dispatch keys) currently in flight, so the effects never
@@ -123,6 +128,8 @@ export function usePolicyAutoRun(): void {
   const polling = useRef<Set<string>>(new Set());
   const importing = useRef<Set<string>>(new Set());
   const dispatching = useRef<Set<string>>(new Set());
+  // Reconcile against the backend exactly once per mount.
+  const reconciled = useRef(false);
 
   // A policy's tool calls run server-side, so a usage-limit 402 never reaches the apiClient
   // interceptor (and thus never pops the modal that direct calls get). The backend surfaces the
@@ -145,6 +152,9 @@ export function usePolicyAutoRun(): void {
   const scheduleQueueRetry = useCallback((runId: string) => {
     const rec = getRun(runId);
     if (!rec) return;
+    // A run rediscovered from the server (reconciled) has no local input fileId, so it can't be
+    // re-dispatched; leave it failed rather than spinning on a file we can't resolve.
+    if (!rec.fileId) return;
     const key = dispatchKey(rec.categoryId, rec.fileId);
     const attempts = queueRetries.current.get(key) ?? 0;
     const backendId = policiesRef.current[rec.categoryId]?.backendId;
@@ -267,12 +277,25 @@ export function usePolicyAutoRun(): void {
       void importOutputs(run, {
         addFiles,
         consumeFiles,
+        bumpRevision,
         outputMode,
         outputName,
         parentStub,
       }).finally(() => importing.current.delete(run.runId));
     }
   }, [runs, addFiles, consumeFiles, policies, fileStubs]);
+
+  // Reconcile against the backend on load. The server owns runs (durable, user-scoped),
+  // so a run started before this client recorded it, or before a refresh/crash, is
+  // rediscovered here; the poll + import effects above then collect its outputs rather
+  // than leaving them orphaned. Waits until policies are known so server runs can be
+  // attributed to their category.
+  useEffect(() => {
+    if (!POLICIES_ENABLED || reconciled.current) return;
+    if (Object.keys(policies).length === 0) return;
+    reconciled.current = true;
+    void reconcileServerRuns(policies);
+  }, [policies]);
 }
 
 interface ImportContext {
@@ -282,6 +305,8 @@ interface ImportContext {
     outputs: StirlingFile[],
     stubs: StirlingFileStub[],
   ) => Promise<unknown>;
+  /** Bump the IndexedDB revision so the file views re-read after a storage-only version write. */
+  bumpRevision: () => void;
   /** "new_file" adds the output as a separate file; "new_version" versions the input. */
   outputMode: "new_file" | "new_version";
   /** Rename rule. Empty → keep the input's filename; set → use the policy's
@@ -291,6 +316,61 @@ interface ImportContext {
   parentStub: StirlingFileStub | undefined;
 }
 
+/**
+ * Pull the caller's server-side runs and fold them into the local store. For a run we already
+ * track, patch its status/outputs (preserving local import progress + attribution); for one we
+ * don't, adopt it so the poll/import effects pick it up. Server-excluded ad-hoc runs and runs we
+ * can't map to a configured category are skipped.
+ */
+async function reconcileServerRuns(
+  policies: PoliciesByCategory,
+): Promise<void> {
+  let serverRuns;
+  try {
+    serverRuns = await listPolicyRuns();
+  } catch {
+    return; // offline / backend down; local cache stands.
+  }
+  for (const view of serverRuns) {
+    // No-ops unless the run is already tracked, so this only patches known runs.
+    updateRun(view.runId, {
+      status: view.status,
+      outputs: view.outputs,
+      error: view.error,
+    });
+    // No-ops if already tracked, so this only adopts runs we'd otherwise have lost.
+    const categoryId = categoryForPolicy(view.policyId, policies);
+    if (!categoryId) continue;
+    addReconciledRun({
+      runId: view.runId,
+      categoryId,
+      // No local input link: a run rediscovered purely from the server (never recorded by this
+      // client) can't be tied back to a workspace/storage file, so its output is delivered as a
+      // new file rather than a version, and it isn't retried. The recorded-run path (real fileId)
+      // covers the common refresh case; this only bites true orphans (storage wipe / other device).
+      fileId: "",
+      fileName: view.outputs[0]?.fileName ?? "",
+      fileSize: 0,
+      status: view.status,
+      outputs: view.outputs,
+      error: view.error,
+      // Use the server's creation time, not now, so a rediscovered run shows its real age.
+      startedAt: view.createdAt,
+    });
+  }
+}
+
+/** The category whose configured policy produced this run, if any. */
+function categoryForPolicy(
+  policyId: string | null,
+  policies: PoliciesByCategory,
+): string | undefined {
+  if (!policyId) return undefined;
+  return Object.entries(policies).find(
+    ([, s]) => s.backendId === policyId,
+  )?.[0];
+}
+
 /**
  * Fetch a completed run's not-yet-imported output files and deliver them to the
  * workspace. Per-output, via allSettled: each output is tracked once delivered,
@@ -343,22 +423,44 @@ async function importOutputs(
   // them, so they retry (without having been added).
   const files = fetched.map((f) => f.file);
   // Workspace fileIds of the delivered outputs — the policy badge marks these
-  // (the policy's output), not the input it ran on. Set in both branches below.
+  // (the policy's output), not the input it ran on. Set in every branch below.
   let deliveredIds: string[];
-  if (ctx.outputMode === "new_version" && ctx.parentStub) {
+  // For new-version output, resolve the input's stub from the active workspace, or from storage
+  // when the workspace is empty (e.g. after a reload, where the run is recovered but the input
+  // still persists in IndexedDB). Versioning it there keeps the result identical to the no-reload
+  // case (one leaf) instead of adding the output as a second file.
+  const parentStub =
+    ctx.outputMode === "new_version"
+      ? (ctx.parentStub ??
+        (await fileStorage.getStirlingFileStub(run.fileId as FileId)) ??
+        undefined)
+      : undefined;
+  if (parentStub) {
     // Replace the input file with a versioned child (preserves its history).
     // The version records "automate" as its origin tool — a policy is a
     // multi-tool automation, not any single tool (redact/watermark/sanitize/…).
     const { stirlingFiles, stubs } = await createStirlingFilesAndStubs(
       files,
-      ctx.parentStub,
+      parentStub,
       "automate",
     );
     // Mark the outputs handled BEFORE adding them, so the auto-run never enforces
     // the policy on its own output — that would version endlessly in a loop.
     for (const s of stubs) markDispatched(run.categoryId, s.id);
     deliveredIds = stubs.map((s) => s.id as string);
-    await ctx.consumeFiles([run.fileId as FileId], stirlingFiles, stubs);
+    if (ctx.parentStub) {
+      // Input is in the active workspace: version it there (workspace + storage).
+      await ctx.consumeFiles([run.fileId as FileId], stirlingFiles, stubs);
+    } else {
+      // Input is only in storage (run recovered after a reload): version it at the
+      // storage layer, then refresh the file views.
+      await fileStorage.persistVersionedOutputs(
+        [run.fileId as FileId],
+        stirlingFiles,
+        stubs,
+      );
+      ctx.bumpRevision();
+    }
   } else {
     const added = await ctx.addFiles(files);
     // Same loop-guard for new-file output: the produced file is a new workspace
@@ -424,8 +526,9 @@ export async function runPolicyOnFile(
       startedAt: Date.now(),
     });
   } catch {
-    // Dispatch failed (offline / backend error). Mark dispatched so we don't
-    // hammer; the absent run simply won't appear in the activity feed.
+    // Dispatch failed (offline / backend error). Mark dispatched so we don't hammer;
+    // the absent run simply won't appear in the activity feed. If the backend did
+    // start a run we never recorded, reconcileServerRuns rediscovers it.
     markDispatched(categoryId, fileId);
   }
 }