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.
This commit is contained in:
James Brunton
2026-06-17 16:18:50 +00:00
committed by GitHub
parent 3750111ffc
commit 13af10a6d1
18 changed files with 469 additions and 72 deletions
@@ -587,60 +587,13 @@ export async function consumeFiles(
);
}
// Mark input files as processed in storage (no longer leaf nodes)
if (
!outputStirlingFileStubs.reduce(
(areAllV1, stub) => areAllV1 && stub.versionNumber == 1,
true,
)
) {
await Promise.all(
inputFileIds.map(async (fileId) => {
try {
await fileStorage.markFileAsProcessed(fileId);
if (DEBUG)
console.log(
`📄 Marked file ${fileId} as processed (no longer leaf)`,
);
} catch (error) {
if (DEBUG)
console.warn(
`📄 Failed to mark file ${fileId} as processed:`,
error,
);
}
}),
);
}
// Save output files directly to fileStorage with complete metadata
for (let i = 0; i < outputStirlingFiles.length; i++) {
const stirlingFile = outputStirlingFiles[i];
const stub = outputStirlingFileStubs[i];
try {
// Use fileStorage directly with complete metadata from stub
await fileStorage.storeStirlingFile(stirlingFile, stub);
if (DEBUG)
console.log(
`📄 Saved StirlingFile ${stirlingFile.name} directly to storage with complete metadata:`,
{
fileId: stirlingFile.fileId,
versionNumber: stub.versionNumber,
originalFileId: stub.originalFileId,
parentFileId: stub.parentFileId,
toolChainLength: stub.toolHistory?.length || 0,
},
);
} catch (error) {
console.error(
"Failed to persist output file to fileStorage:",
stirlingFile.name,
error,
);
}
}
// Persist the durable half (mark inputs non-leaf + store output versions) via the shared
// storage helper, so a policy run recovered after a reload versions the file identically.
await fileStorage.persistVersionedOutputs(
inputFileIds,
outputStirlingFiles,
outputStirlingFileStubs,
);
// Dispatch the consume action with pre-created stubs (no processing needed)
dispatch({
@@ -730,6 +730,53 @@ class FileStorageService {
}
}
/**
* Persist output files as versions of their inputs: mark each input non-leaf (unless the
* outputs are v1 originals, i.e. nothing was versioned) and store each output with its stub.
* This is the durable half of {@link consumeFiles}, shared so a versioned result can be written
* even when the input isn't in the active workspace (e.g. a policy run recovered after a reload).
* Storage-only callers must bump the IndexedDB revision afterwards so the file views re-read;
* {@link consumeFiles} instead updates workspace state via its dispatch.
*/
async persistVersionedOutputs(
inputFileIds: FileId[],
outputStirlingFiles: StirlingFile[],
outputStirlingFileStubs: StirlingFileStub[],
): Promise<void> {
if (outputStirlingFiles.length !== outputStirlingFileStubs.length) {
throw new Error(
`Mismatch between output files (${outputStirlingFiles.length}) and stubs (${outputStirlingFileStubs.length})`,
);
}
const allV1 = outputStirlingFileStubs.every(
(stub) => stub.versionNumber === 1,
);
if (!allV1) {
await Promise.all(
inputFileIds.map((fileId) =>
this.markFileAsProcessed(fileId).catch((error) => {
// Best-effort: a missing/locked input shouldn't block storing the outputs.
console.warn(`Failed to mark file ${fileId} as processed:`, error);
}),
),
);
}
await Promise.all(
outputStirlingFiles.map((file, i) =>
this.storeStirlingFile(file, outputStirlingFileStubs[i]).catch(
(error) =>
console.error(
"Failed to persist output file to storage:",
file.name,
error,
),
),
),
);
}
/**
* Mark a file as leaf (opposite of markFileAsProcessed)
* Used when promoting a file back to "recent" status
@@ -81,6 +81,7 @@ vi.mock("../../services/fileStorage", () => ({
});
}),
storeStirlingFile: vi.fn().mockResolvedValue(undefined),
persistVersionedOutputs: vi.fn().mockResolvedValue(undefined),
getAllFileMetadata: vi.fn().mockResolvedValue([]),
cleanup: vi.fn().mockResolvedValue(undefined),
},
@@ -213,9 +214,9 @@ describe("Convert Tool Integration Tests", () => {
expect(result.current.isLoading).toBe(false);
expect(result.current.errorMessage).toBe(null);
// The output file must be persisted via fileStorage.storeStirlingFile
// The output file must be persisted via fileStorage.persistVersionedOutputs
// so downstream tools see it in the registry.
expect(fileStorage.storeStirlingFile).toHaveBeenCalled();
expect(fileStorage.persistVersionedOutputs).toHaveBeenCalled();
});
test("should handle API error responses correctly", async () => {
@@ -79,6 +79,7 @@ vi.mock("../../services/fileStorage", () => ({
});
}),
storeStirlingFile: vi.fn().mockResolvedValue(undefined),
persistVersionedOutputs: vi.fn().mockResolvedValue(undefined),
getAllFileMetadata: vi.fn().mockResolvedValue([]),
cleanup: vi.fn().mockResolvedValue(undefined),
},
@@ -173,9 +174,9 @@ describe("Convert Tool - Smart Detection Integration Tests", () => {
},
);
// The output file must be persisted via fileStorage.storeStirlingFile
// The output file must be persisted via fileStorage.persistVersionedOutputs
// so downstream tools see it in the registry.
expect(fileStorage.storeStirlingFile).toHaveBeenCalled();
expect(fileStorage.persistVersionedOutputs).toHaveBeenCalled();
});
test("should handle unknown file type with file-to-pdf fallback", async () => {
@@ -256,6 +256,13 @@ export async function mockAppApis(
await page.route("**/api/v1/policies", (route: Route) =>
route.fulfill({ json: [] }),
);
// The auto-run controller also reconciles server-side runs on load via
// GET /api/v1/policies/runs. The glob above doesn't cover this sub-path, so
// stub it empty too; otherwise the request hits the absent backend and the
// console error fails the page's no-unexpected-output guard.
await page.route("**/api/v1/policies/runs", (route: Route) =>
route.fulfill({ json: [] }),
);
}
/**
@@ -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);
@@ -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();
});
});
@@ -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";
@@ -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);
}
}
@@ -104,3 +104,14 @@ export async function getPolicyRun(runId: string): Promise<PolicyRunView> {
);
return res.data;
}
/**
* The caller's in-flight and recently-finished stored-policy runs (server-owned,
* within the run-retention window). Used to reconcile on load: a run started
* before a refresh/crash is rediscovered here and its outputs collected, so a
* finished run is never orphaned on the backend.
*/
export async function listPolicyRuns(): Promise<PolicyRunView[]> {
const res = await apiClient.get<PolicyRunView[]>("/api/v1/policies/runs");
return res.data;
}
@@ -73,6 +73,8 @@ export interface BackendResultFile {
/** Read-only view returned by the run status endpoint (mirrors PolicyRunView). */
export interface PolicyRunView {
runId: string;
/** ID of the stored policy that produced the run; null for ad-hoc pipelines. */
policyId: string | null;
status: PolicyRunStatus;
currentStep: number;
stepCount: number;
@@ -89,6 +91,8 @@ export interface PolicyRunView {
*/
errorSubscribed?: boolean | null;
outputs: BackendResultFile[];
/** When the run was created (epoch millis); lets a rediscovered run show its real age. */
createdAt: number;
}
/** Resolve a frontend operation id to its backend tool endpoint path. */