Policies: enforce input on uploads only; badge follows edited files (#6627)

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

@@ -83,6 +83,11 @@ export function usePolicyAutoRun(): void {
     );
     for (const [categoryId, s] of active) {
       for (const stub of fileStubs) {
+        // Input-mode policies enforce only on files that actually entered the
+        // system as an upload — not on files a tool/automation produced in-app
+        // (versioned edits or independent artifacts like convert/split/merge).
+        // Those are enforced only by export-mode policies, at export time.
+        if (stub.derivedFromTool) continue;
         const key = dispatchKey(categoryId, stub.id);
         // Skip if already run (persisted) or a dispatch is in flight — the
         // in-memory guard prevents double-firing during the async wait.

frontend/editor/src/proprietary/hooks/usePolicyFileBadges.test.ts

@@ -0,0 +1,113 @@
+import { describe, it, expect } from "vitest";
+import { buildPolicyBadgeMap } from "@app/hooks/usePolicyFileBadges";
+import type { PolicyRunRecord } from "@app/components/policies/policyRunStore";
+
+const NOW = 1_000_000;
+const labels = new Map([
+  ["security", "Security"],
+  ["watermark", "Watermark"],
+]);
+
+function run(overrides: Partial<PolicyRunRecord>): PolicyRunRecord {
+  return {
+    runId: "r",
+    categoryId: "security",
+    fileId: "in",
+    fileName: "in.pdf",
+    fileSize: 1,
+    status: "COMPLETED",
+    outputs: [],
+    outputFileIds: ["out"],
+    error: null,
+    startedAt: NOW - 1_000, // recent by default
+    ...overrides,
+  };
+}
+
+describe("buildPolicyBadgeMap — badge follows the document onto derived files", () => {
+  it("badges a policy's direct output, and marks it recent within the window", () => {
+    const map = buildPolicyBadgeMap([run({})], [{ id: "out" }], labels, NOW);
+    const badges = map.get("out") ?? [];
+    expect(badges.map((b) => b.id)).toEqual(["security"]);
+    expect(badges[0].recent).toBe(true);
+  });
+
+  it("a versioned edit inherits the badge via parentFileId (never glows)", () => {
+    const map = buildPolicyBadgeMap(
+      [run({})],
+      [{ id: "out" }, { id: "edit", parentFileId: "out" }],
+      labels,
+      NOW,
+    );
+    const edit = map.get("edit") ?? [];
+    expect(edit.map((b) => b.id)).toEqual(["security"]);
+    expect(edit[0].recent).toBe(false);
+  });
+
+  it("SPLIT parts inherit the badge via sourceFileIds, though they have no parent", () => {
+    // Split consumes the policy output "out" → two fresh roots, no parentFileId,
+    // each recording sourceFileIds=["out"]. "out" itself is gone from the
+    // workbench (consumed) but still lives in the run store.
+    const map = buildPolicyBadgeMap(
+      [run({})],
+      [
+        { id: "part1", sourceFileIds: ["out"] },
+        { id: "part2", sourceFileIds: ["out"] },
+      ],
+      labels,
+      NOW,
+    );
+    expect((map.get("part1") ?? []).map((b) => b.id)).toEqual(["security"]);
+    expect((map.get("part2") ?? []).map((b) => b.id)).toEqual(["security"]);
+    expect((map.get("part1") ?? [])[0].recent).toBe(false);
+  });
+
+  it("resolves transitively when an intermediate edit was consumed/removed", () => {
+    // redact → edit (consumed) → split. The split part's sourceFileIds carries
+    // the original output id directly, so the badge still resolves.
+    const map = buildPolicyBadgeMap(
+      [run({})],
+      [{ id: "part", sourceFileIds: ["editGone", "out"] }],
+      labels,
+      NOW,
+    );
+    expect((map.get("part") ?? []).map((b) => b.id)).toEqual(["security"]);
+  });
+
+  it("MERGE output inherits every input's badge", () => {
+    const map = buildPolicyBadgeMap(
+      [
+        run({ runId: "r1", categoryId: "security", outputFileIds: ["a"] }),
+        run({ runId: "r2", categoryId: "watermark", outputFileIds: ["b"] }),
+      ],
+      [{ id: "merged", sourceFileIds: ["a", "b"] }],
+      labels,
+      NOW,
+    );
+    expect((map.get("merged") ?? []).map((b) => b.id).sort()).toEqual([
+      "security",
+      "watermark",
+    ]);
+  });
+
+  it("a file with no policy provenance gets no badge", () => {
+    const map = buildPolicyBadgeMap(
+      [run({})],
+      [{ id: "out" }, { id: "unrelated", sourceFileIds: ["someUpload"] }],
+      labels,
+      NOW,
+    );
+    expect(map.has("unrelated")).toBe(false);
+  });
+
+  it("inherited badges never glow even when the source run is recent", () => {
+    const map = buildPolicyBadgeMap(
+      [run({ startedAt: NOW })], // maximally recent
+      [{ id: "out" }, { id: "part", sourceFileIds: ["out"] }],
+      labels,
+      NOW,
+    );
+    expect((map.get("out") ?? [])[0].recent).toBe(true);
+    expect((map.get("part") ?? [])[0].recent).toBe(false);
+  });
+});

frontend/editor/src/proprietary/hooks/usePolicyFileBadges.ts

@@ -1,5 +1,7 @@
 import { useMemo } from "react";
 import { usePolicyRuns } from "@app/components/policies/policyRunStore";
+import type { PolicyRunRecord } from "@app/components/policies/policyRunStore";
+import { useAllFiles } from "@app/contexts/FileContext";
 import { loadPolicyCatalog } from "@app/services/policyCatalog";
 import { ROW_ACCENT } from "@app/components/policies/policyStatus";
 import type { FileItemPolicyRef } from "@app/components/shared/FileSidebarFileItem";
@@ -18,38 +20,110 @@ const ACCENT_VAR: Record<string, string> = {
   red: "var(--color-red)",
 };
 
+/** Minimal provenance shape needed to resolve a file's inherited badges. */
+type LineageStub = {
+  id: string;
+  parentFileId?: string;
+  sourceFileIds?: string[];
+};
+
+/** Merge a ref into a list, deduping by policy id. A direct (recent) hit wins
+ *  the glow over an inherited one for the same policy. */
+function mergeRef(list: FileItemPolicyRef[], ref: FileItemPolicyRef): void {
+  const existing = list.find((p) => p.id === ref.id);
+  if (!existing) {
+    list.push(ref);
+  } else if (ref.recent) {
+    existing.recent = true;
+  }
+}
+
+/**
+ * Pure core of {@link usePolicyFileBadges} (no React/storage deps, so it's
+ * unit-testable). Returns the policies that have run on each file, keyed by
+ * fileId.
+ *
+ * A policy run is pinned to a specific output fileId, so a later tool edit
+ * produces a NEW file that wouldn't carry the badge — "been through a policy"
+ * would vanish the moment you edit the file. To keep the badge with the
+ * document, every file also INHERITS the badges of the files it was derived
+ * from: its transitive `sourceFileIds` (recorded at the consume boundary, so it
+ * covers split/merge/convert too) plus, defensively, its `parentFileId`.
+ * Because `sourceFileIds` is transitive, a flat lookup suffices — no chain walk,
+ * and it survives a consumed intermediate. Inherited badges never glow
+ * (recent=false): only the original application does.
+ */
+export function buildPolicyBadgeMap(
+  runs: ReadonlyArray<PolicyRunRecord>,
+  stubs: ReadonlyArray<LineageStub>,
+  labelById: ReadonlyMap<string, string>,
+  now: number,
+): Map<string, FileItemPolicyRef[]> {
+  // Direct badges: a file that IS a policy run's output.
+  const directByFile = new Map<string, FileItemPolicyRef[]>();
+  for (const run of runs) {
+    const name = labelById.get(run.categoryId);
+    if (!name) continue;
+    const recent = now - run.startedAt < RECENT_MS;
+    for (const fileId of run.outputFileIds ?? []) {
+      const list = directByFile.get(fileId) ?? [];
+      if (!list.some((p) => p.id === run.categoryId)) {
+        list.push({
+          id: run.categoryId,
+          name,
+          accentColor: ACCENT_VAR[ROW_ACCENT[run.categoryId] ?? "blue"],
+          recent,
+        });
+        directByFile.set(fileId, list);
+      }
+    }
+  }
+
+  // Seed the result with the direct badges (cloning the arrays so the lineage
+  // pass can mutate freely without touching the direct map).
+  const result = new Map<string, FileItemPolicyRef[]>();
+  for (const [id, refs] of directByFile) {
+    if (refs.length)
+      result.set(
+        id,
+        refs.map((r) => ({ ...r })),
+      );
+  }
+
+  // Inheritance pass: a derived file carries the badges of every file it came
+  // from. `sourceFileIds` is the transitive provenance set (so a flat lookup
+  // catches even ancestors whose intermediate edits were consumed), and
+  // `parentFileId` is included defensively for any child not created via a
+  // consume. Inherited badges are marked recent=false (carried, not applied).
+  for (const stub of stubs) {
+    const sources = new Set<string>(stub.sourceFileIds ?? []);
+    if (stub.parentFileId) sources.add(stub.parentFileId);
+    for (const src of sources) {
+      const srcBadges = directByFile.get(src);
+      if (!srcBadges?.length) continue;
+      const list = result.get(stub.id) ?? [];
+      for (const ref of srcBadges) mergeRef(list, { ...ref, recent: false });
+      result.set(stub.id, list);
+    }
+  }
+
+  return result;
+}
+
 /**
  * Distinct policies that have produced each file, keyed by fileId, derived from
  * the reactive policy run store. Drives the file sidebar's shield badges. The
- * badge marks a policy's OUTPUT (the versioned/added result), not the input it
- * ran on — so it keys off each run's imported output fileIds. Shadows the core
- * stub via the {@code @app/*} alias cascade.
+ * badge follows a document down its tool-edit chain — see
+ * {@link buildPolicyBadgeMap}. Shadows the core stub via the {@code @app/*}
+ * alias cascade.
  */
 export function usePolicyFileBadges(): Map<string, FileItemPolicyRef[]> {
   const runs = usePolicyRuns();
+  const { fileStubs } = useAllFiles();
   return useMemo(() => {
     const labelById = new Map(
       loadPolicyCatalog().categories.map((c) => [c.id, c.label]),
     );
-    const now = Date.now();
-    const byFile = new Map<string, FileItemPolicyRef[]>();
-    for (const run of runs) {
-      const name = labelById.get(run.categoryId);
-      if (!name) continue;
-      const recent = now - run.startedAt < RECENT_MS;
-      for (const fileId of run.outputFileIds ?? []) {
-        const list = byFile.get(fileId) ?? [];
-        if (!list.some((p) => p.id === run.categoryId)) {
-          list.push({
-            id: run.categoryId,
-            name,
-            accentColor: ACCENT_VAR[ROW_ACCENT[run.categoryId] ?? "blue"],
-            recent,
-          });
-          byFile.set(fileId, list);
-        }
-      }
-    }
-    return byFile;
-  }, [runs]);
+    return buildPolicyBadgeMap(runs, fileStubs, labelById, Date.now());
+  }, [runs, fileStubs]);
 }