diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/ClassifyLabelController.java b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/ClassifyLabelController.java index 798de71df9..dc422433f0 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/ClassifyLabelController.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/ClassifyLabelController.java @@ -9,6 +9,7 @@ import java.util.Map; import java.util.Set; import org.apache.pdfbox.pdmodel.PDDocument; +import org.apache.pdfbox.pdmodel.PDDocumentInformation; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.core.io.Resource; import org.springframework.http.MediaType; @@ -20,12 +21,13 @@ import org.springframework.web.bind.annotation.RestController; import org.springframework.web.multipart.MultipartFile; import io.github.pixee.security.Filenames; -import io.swagger.v3.oas.annotations.Hidden; import io.swagger.v3.oas.annotations.Operation; import io.swagger.v3.oas.annotations.tags.Tag; import lombok.extern.slf4j.Slf4j; +import stirling.software.common.model.tool.ToolFormat; +import stirling.software.common.model.tool.ToolIO; import stirling.software.common.service.CustomPDFDocumentFactory; import stirling.software.common.service.PdfMetadataService; import stirling.software.common.service.UserServiceInterface; @@ -48,11 +50,13 @@ import tools.jackson.databind.node.ObjectNode; *

Runs as a Classification-policy pipeline step: it reads a bounded page window, asks the AI * engine to classify the document against the built-in label set, and stores the engine's JSON * answer — minus the transport-only {@code outcome} field — in the custom Info-dictionary key - * {@link PdfMetadataService#CLASSIFICATION_KEY}. Returns the labelled PDF. Not intended for direct - * client use. + * {@link PdfMetadataService#CLASSIFICATION_KEY}. Returns the labelled PDF. + * + *

Published in the API spec rather than hidden, so the tool-model generator emits it and a + * pipeline can name it as a step like any other tool. Classification is a thing a pipeline does, + * not a thing only the Classification policy may do. */ @Slf4j -@Hidden @RestController @RequestMapping("/api/v1/ai/tools") @Tag(name = "AI Tools", description = "Dispatchable AI-backed tools.") @@ -99,19 +103,31 @@ public class ClassifyLabelController { } @PostMapping(value = "/classify-and-label", consumes = MediaType.MULTIPART_FORM_DATA_VALUE) + // PDF in, the same PDF out with a verdict on it, so a chain can be checked across this step. + @ToolIO(accepts = ToolFormat.PDF, produces = ToolFormat.PDF) @Operation( summary = "Classify a PDF and label its metadata", description = "Reads the first two and last two pages, classifies the document via the AI" + " engine, and stores the result in the StirlingPDFClassification" - + " metadata field. Dispatched by the Classification policy; not" - + " intended for direct client use.") + + " metadata field. A document that already carries a verdict is" + + " passed through untouched unless reclassify=true.") public ResponseEntity classifyAndLabel( - @RequestParam("fileInput") MultipartFile fileInput) throws IOException { + @RequestParam("fileInput") MultipartFile fileInput, + @RequestParam(value = "reclassify", defaultValue = "false") boolean reclassify) + throws IOException { aiFeatureGate.requireClassify(); try (PDDocument document = pdfDocumentFactory.load(fileInput, true)) { String fileName = safeFileName(fileInput.getOriginalFilename()); + if (!reclassify && isClassified(document)) { + // Classifying twice costs a second engine call and charges for it, and a document + // that already carries a verdict has nothing new to learn. A pipeline can run this + // step over a mixed batch without paying for the ones already done. + log.debug("[classify-and-label] {} already classified; passing through", fileName); + return WebResponseUtils.pdfDocToWebResponse(document, fileName, tempFileManager); + } + List allowed = resolveAllowedLabels(); if (allowed.isEmpty()) { // No vocabulary to classify against: pass the file through unlabelled rather than @@ -135,6 +151,25 @@ public class ClassifyLabelController { } } + /** + * Whether a verdict is already on the document. + * + *

This only reads back what a previous run of this step wrote. It is not a statement that + * the verdict is trustworthy: the key is ordinary PDF metadata that whoever supplied the file + * can set. Skipping the engine on the strength of it is safe because the cost of being wrong is + * a missing re-classification, not a wrong decision. Anything that makes a SECURITY decision + * from this field - routing a document somewhere on the strength of its label, say - must + * classify with {@code reclassify=true} rather than trust what arrived. + */ + private static boolean isClassified(PDDocument document) { + PDDocumentInformation info = document.getDocumentInformation(); + if (info == null) { + return false; + } + String existing = info.getCustomMetadataValue(PdfMetadataService.CLASSIFICATION_KEY); + return existing != null && !existing.isBlank(); + } + private List extractWindow(PDDocument document) throws IOException { List pages = new ArrayList<>(); for (int pageNumber : windowPageNumbers(document.getNumberOfPages(), WINDOW_PAGES)) { diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/policy/model/Policy.java b/app/proprietary/src/main/java/stirling/software/proprietary/policy/model/Policy.java index a3d43b712d..14b1eb325c 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/policy/model/Policy.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/policy/model/Policy.java @@ -85,6 +85,11 @@ public record Policy( return new Policy(id, name, owner, enabled, inputs, steps, resolved, outputIds, teamId); } + /** A copy under a different owner (e.g. moving a seed off a placeholder name). */ + public Policy withOwner(String newOwner) { + return new Policy(id, name, newOwner, enabled, inputs, steps, output, outputIds, teamId); + } + /** A copy referencing the given saved output destinations. */ public Policy withOutputIds(List newOutputIds) { return new Policy(id, name, owner, enabled, inputs, steps, output, newOutputIds, teamId); diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/policy/seed/DefaultClassificationPolicySeeder.java b/app/proprietary/src/main/java/stirling/software/proprietary/policy/seed/DefaultClassificationPolicySeeder.java index 63ebae833a..9b347366bc 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/policy/seed/DefaultClassificationPolicySeeder.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/policy/seed/DefaultClassificationPolicySeeder.java @@ -22,8 +22,8 @@ import stirling.software.proprietary.security.repository.TeamRepository; import stirling.software.proprietary.security.service.TeamService; /** - * Seeds an enabled Classification policy per team so classification is on by default. Idempotent; - * skips the internal team. + * Seeds an enabled Classification policy per team; idempotent, skips the internal team. Left + * unowned: nobody created it, and an owner here would have to name a real user. */ @Slf4j @Component @@ -34,6 +34,11 @@ public class DefaultClassificationPolicySeeder { private static final String CLASSIFY_ENDPOINT = "/api/v1/ai/tools/classify-and-label"; private static final String POLICY_NAME = "Classification Policy"; + /** + * Pre-existing seeds used this placeholder, which was never a user; see {@link #repairOwner}. + */ + private static final String LEGACY_OWNER = "system"; + private final PolicyStore policyStore; private final TeamRepository teamRepository; @@ -46,9 +51,8 @@ public class DefaultClassificationPolicySeeder { .ifPresent(team -> seedIfMissing(team.getId(), team.getName())); } - // Any team created at runtime (admin-created, SaaS sign-ups). Seeds inside the team's own - // transaction: rollback still leaves no policy behind, and the store's pessimistic lock needs a - // live transaction, which AFTER_COMMIT cannot offer. + // Seeds inside the new team's own transaction: rollback leaves no policy behind, and the + // store's pessimistic lock needs a live transaction, which AFTER_COMMIT cannot offer. @TransactionalEventListener(phase = TransactionPhase.BEFORE_COMMIT) public void onTeamCreated(TeamCreatedEvent event) { seedIfMissing(event.teamId(), event.teamName()); @@ -58,16 +62,33 @@ public class DefaultClassificationPolicySeeder { if (teamId == null || TeamService.INTERNAL_TEAM_NAME.equals(teamName)) { return; } - boolean alreadySeeded = + Policy existing = policyStore.findByTeam(teamId).stream() - .anyMatch(DefaultClassificationPolicySeeder::isClassification); - if (alreadySeeded) { + .filter(DefaultClassificationPolicySeeder::isClassification) + .findFirst() + .orElse(null); + if (existing != null) { + repairOwner(existing); return; } policyStore.save(defaultPolicy(teamId)); log.info("Seeded default Classification policy for team {}", teamId); } + /** + * Clear an owner seeded as a placeholder name. An owner someone deliberately set is left alone. + */ + private void repairOwner(Policy policy) { + if (!LEGACY_OWNER.equals(policy.owner())) { + return; + } + policyStore.save(policy.withOwner(null)); + log.info( + "Cleared placeholder owner '{}' on Classification policy {}", + LEGACY_OWNER, + policy.id()); + } + private static boolean isClassification(Policy policy) { return policy.output() != null && CATEGORY.equals(policy.output().options().get("categoryId")); @@ -85,7 +106,9 @@ public class DefaultClassificationPolicySeeder { return new Policy( null, POLICY_NAME, - "system", + // Nobody created this - it is seeded. A name here would have to be a real user, and + // every consumer of owner already handles its absence. + null, true, List.of(), List.of(new PipelineStep(CLASSIFY_ENDPOINT, Map.of())), diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/controller/api/ClassifyLabelControllerTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/controller/api/ClassifyLabelControllerTest.java index 74749289a7..637f05b284 100644 --- a/app/proprietary/src/test/java/stirling/software/proprietary/controller/api/ClassifyLabelControllerTest.java +++ b/app/proprietary/src/test/java/stirling/software/proprietary/controller/api/ClassifyLabelControllerTest.java @@ -14,6 +14,7 @@ import static org.mockito.Mockito.when; import java.util.List; import org.apache.pdfbox.pdmodel.PDDocument; +import org.apache.pdfbox.pdmodel.PDDocumentInformation; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.ArgumentCaptor; @@ -76,7 +77,7 @@ class ClassifyLabelControllerTest { .thenReturn("{\"outcome\":\"classification\",\"labels\":[\"invoice\"]}"); try { - controller.classifyAndLabel(file); + controller.classifyAndLabel(file, false); } catch (Exception ignored) { // WebResponseUtils.pdfDocToWebResponse needs a real temp file; the engine call and // metadata write we assert on have already happened by the time it runs. @@ -89,6 +90,52 @@ class ClassifyLabelControllerTest { return objectMapper.readTree(body.getValue()); } + /** Stubs a document that already carries a verdict, as a second run over a batch would see. */ + private MultipartFile alreadyClassifiedDocument() throws Exception { + PDDocument document = mock(PDDocument.class); + PDDocumentInformation info = mock(PDDocumentInformation.class); + when(document.getDocumentInformation()).thenReturn(info); + when(info.getCustomMetadataValue(PdfMetadataService.CLASSIFICATION_KEY)) + .thenReturn("{\"labels\":[\"invoice\"]}"); + MultipartFile file = mock(MultipartFile.class); + when(file.getOriginalFilename()).thenReturn("invoice.pdf"); + when(pdfDocumentFactory.load(any(MultipartFile.class), eq(true))).thenReturn(document); + return file; + } + + @Test + void classifyAndLabel_skipsADocumentThatAlreadyCarriesAVerdict() throws Exception { + withLabels(List.of(new ClassificationLabel("invoice", "Invoice", null))); + MultipartFile file = alreadyClassifiedDocument(); + + try { + controller.classifyAndLabel(file, false); + } catch (Exception ignored) { + // The response needs a real temp file; the decision under test happens before it. + } + + // No second engine call, and no charge for one: re-classifying buys the same answer twice. + verify(aiEngineClient, never()).post(anyString(), anyString(), any()); + verify(pdfMetadataService, never()).setClassificationMetadata(any(), anyString()); + } + + @Test + void classifyAndLabel_reclassifiesWhenAskedTo() throws Exception { + withLabels(List.of(new ClassificationLabel("invoice", "Invoice", null))); + MultipartFile file = alreadyClassifiedDocument(); + when(pdfContentExtractor.extractPageTextRaw(any(), eq(1))).thenReturn("Invoice total"); + when(aiEngineClient.post(eq("/api/v1/documents/classify"), anyString(), isNull())) + .thenReturn("{\"outcome\":\"classification\",\"labels\":[\"receipt\"]}"); + + try { + controller.classifyAndLabel(file, true); + } catch (Exception ignored) { + // As above. + } + + verify(aiEngineClient).post(eq("/api/v1/documents/classify"), anyString(), isNull()); + } + @Test void classifyAndLabel_writesClassificationWithoutOutcome() throws Exception { withLabels(List.of(new ClassificationLabel("invoice", "Invoice", null))); diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/policy/seed/DefaultClassificationPolicySeederTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/policy/seed/DefaultClassificationPolicySeederTest.java index 49159e7d6e..f6e82bd011 100644 --- a/app/proprietary/src/test/java/stirling/software/proprietary/policy/seed/DefaultClassificationPolicySeederTest.java +++ b/app/proprietary/src/test/java/stirling/software/proprietary/policy/seed/DefaultClassificationPolicySeederTest.java @@ -36,10 +36,14 @@ class DefaultClassificationPolicySeederTest { } private static Policy classificationPolicy(Long teamId) { + return classificationPolicy(teamId, null); + } + + private static Policy classificationPolicy(Long teamId, String owner) { return new Policy( "p1", "Classification Policy", - "system", + owner, true, List.of(), List.of(), @@ -77,6 +81,29 @@ class DefaultClassificationPolicySeederTest { verify(policyStore, never()).save(any()); } + @Test + void clearsAPlaceholderOwnerSeededBeforeOwnersHadToBeReal() { + when(policyStore.findByTeam(7L)).thenReturn(List.of(classificationPolicy(7L, "system"))); + + seeder().onTeamCreated(new TeamCreatedEvent(7L, "Acme")); + + // "system" was never a user row, and a step dispatch authenticates as the owner. Absence + // is handled everywhere; a placeholder name is not. + ArgumentCaptor saved = ArgumentCaptor.forClass(Policy.class); + verify(policyStore).save(saved.capture()); + assertThat(saved.getValue().owner()).isNull(); + assertThat(saved.getValue().id()).isEqualTo("p1"); + } + + @Test + void leavesADeliberatelyChosenOwnerAlone() { + when(policyStore.findByTeam(7L)).thenReturn(List.of(classificationPolicy(7L, "alice"))); + + seeder().onTeamCreated(new TeamCreatedEvent(7L, "Acme")); + + verify(policyStore, never()).save(any()); + } + @Test void doesNotSeedForTheInternalTeam() { seeder().onTeamCreated(new TeamCreatedEvent(2L, "Internal")); diff --git a/frontend/editor/public/locales/en-US/translation.toml b/frontend/editor/public/locales/en-US/translation.toml index d2eda294b2..88814dcf08 100644 --- a/frontend/editor/public/locales/en-US/translation.toml +++ b/frontend/editor/public/locales/en-US/translation.toml @@ -4592,6 +4592,10 @@ desc = "Change document restrictions and permissions" tags = "permissions,restrictions,rights,access control,allow,deny,printing,copying,editing,modify permissions,security settings,user rights" title = "Change Permissions" +[home.classify] +desc = "Identify what kind of document this is and tag it." +title = "Classify" + [home.compare] desc = "Compares and shows the differences between 2 PDF Documents" tags = "difference,compare,diff,compare PDFs,compare documents,find differences,show differences,changes,what changed,track changes,revisions,version compare,side by side,contrast,delta" @@ -10841,6 +10845,7 @@ searchPlaceholder = "Search tools..." [toolPicker.subcategories] advancedFormatting = "Advanced Formatting" +ai = "AI" automation = "Automation" developerTools = "Developer Tools" documentReview = "Document Review" diff --git a/frontend/editor/public/og-metadata.json b/frontend/editor/public/og-metadata.json index f9733bd83f..b1eadb3bce 100644 --- a/frontend/editor/public/og-metadata.json +++ b/frontend/editor/public/og-metadata.json @@ -195,6 +195,11 @@ "title": "Compress - Stirling PDF", "description": "Compress PDFs to reduce their file size." }, + "classify": { + "image": "/og_images/home.png", + "title": "Classify - Stirling PDF", + "description": "Identify what kind of document this is and tag it." + }, "extractPages": { "image": "/og_images/extract-pages.png", "title": "Extract Pages - Stirling PDF", @@ -575,6 +580,7 @@ "/remove-cert-sign": "removeCertSign", "/unlock-p-d-f-forms": "unlockPDFForms", "/compress": "compress", + "/classify": "classify", "/extract-pages": "extractPages", "/reorganize-pages": "reorganizePages", "/extract-images": "extractImages", diff --git a/frontend/editor/public/og-metadata.saas.json b/frontend/editor/public/og-metadata.saas.json index be2e7a3e92..14e30f827c 100644 --- a/frontend/editor/public/og-metadata.saas.json +++ b/frontend/editor/public/og-metadata.saas.json @@ -196,6 +196,11 @@ "title": "Compress - Stirling PDF", "description": "Compress PDFs to reduce their file size." }, + "classify": { + "image": "/og_images/home.png", + "title": "Classify - Stirling PDF", + "description": "Identify what kind of document this is and tag it." + }, "extractPages": { "image": "/og_images/extract-pages.png", "title": "Extract Pages - Stirling PDF", @@ -588,6 +593,7 @@ "/remove-cert-sign": "removeCertSign", "/unlock-p-d-f-forms": "unlockPDFForms", "/compress": "compress", + "/classify": "classify", "/extract-pages": "extractPages", "/reorganize-pages": "reorganizePages", "/extract-images": "extractImages", diff --git a/frontend/editor/scripts/generate-tool-api-types.mts b/frontend/editor/scripts/generate-tool-api-types.mts index 23daa1dd76..ec38dc37d7 100644 --- a/frontend/editor/scripts/generate-tool-api-types.mts +++ b/frontend/editor/scripts/generate-tool-api-types.mts @@ -11,13 +11,8 @@ import { dirname, resolve } from "node:path"; import { parseArgs } from "node:util"; import { compile, type JSONSchema } from "json-schema-to-typescript"; -// The API namespaces whose endpoints a pipeline can reference. `/api/v1/ai/tools/` -// is absent from the spec, so it cannot appear here. Extend this list when other -// namespaces become tools. -// -// `/api/v1/filter/` and `/api/v1/integration/` are included even though neither is a -// user-facing tool: a stored pipeline can contain one, and ToolEndpoint keys the I/O -// table, so leaving them out would stop a chain being checked past such a step. +// Endpoints a pipeline can reference. filter/integration are not user-facing tools but a stored +// pipeline can contain one; the AI namespace is admitted one endpoint at a time, not wholesale. const ALLOWED_PATH_PREFIXES = [ "/api/v1/general/", "/api/v1/misc/", @@ -25,6 +20,7 @@ const ALLOWED_PATH_PREFIXES = [ "/api/v1/convert/", "/api/v1/filter/", "/api/v1/integration/", + "/api/v1/ai/tools/classify-and-label", ]; // File plumbing, not user parameters: `fileInput` and `file` are the uploaded primary document diff --git a/frontend/editor/src/core/data/classifyIsAPipelineTask.test.tsx b/frontend/editor/src/core/data/classifyIsAPipelineTask.test.tsx new file mode 100644 index 0000000000..f5281f80ae --- /dev/null +++ b/frontend/editor/src/core/data/classifyIsAPipelineTask.test.tsx @@ -0,0 +1,76 @@ +/** + * Classification is a thing a pipeline can do, not a thing only the Classification policy may do. + * + * The chain that used to stop it: `getExecutableTools` drops any tool whose endpoint is not a + * member of the generated `ToolEndpoint` union; that union comes from the OpenAPI spec, gated by + * the generator's namespace allowlist; and the classify controller was `@Hidden`, so it never + * reached the spec at all. These tests pin each link, because any one of them silently removes the + * step from the builder's picker rather than failing loudly. + */ +import { describe, expect, test, vi } from "vitest"; +import { renderHook } from "@testing-library/react"; +import { useTranslatedToolCatalog } from "@app/data/useTranslatedToolRegistry"; +import { getExecutableTools } from "@app/hooks/tools/shared/toolAutomation"; +import { isToolEndpoint } from "@app/hooks/tools/shared/toolApiMapping"; +import { TOOL_IO } from "@app/types/toolIO"; +import { filterToolRegistryByQuery } from "@app/utils/toolSearch"; + +vi.mock("react-i18next", () => ({ + useTranslation: () => ({ + t: (key: string, fallback?: string) => fallback ?? key, + i18n: { changeLanguage: vi.fn(), language: "en-US" }, + }), + Trans: ({ children }: { children?: unknown }) => children, +})); + +const CLASSIFY_ENDPOINT = "/api/v1/ai/tools/classify-and-label"; + +describe("classify as a pipeline task", () => { + test("the classify endpoint is a generated ToolEndpoint", () => { + // Fails if the controller goes back to @Hidden, or the generator's allowlist drops the + // /api/v1/ai/tools/ namespace, or nobody regenerated after either. + expect(isToolEndpoint(CLASSIFY_ENDPOINT)).toBe(true); + }); + + test("the builder offers it as a step", () => { + const { result } = renderHook(() => useTranslatedToolCatalog()); + + const executable = getExecutableTools(result.current.regularTools); + const classify = executable.find((tool) => tool.toolId === "classify"); + + expect(classify).toBeDefined(); + expect(classify?.endpoint).toBe(CLASSIFY_ENDPOINT); + }); + + test("it declares PDF in, PDF out, so a chain can be checked across it", () => { + // Without this the builder shows "Can't check what this step accepts" and validation stops + // dead at the step - the I/O table is keyed by endpoint and comes from @ToolIO in the spec. + expect(TOOL_IO[CLASSIFY_ENDPOINT]).toEqual({ + accepts: ["PDF"], + produces: "PDF", + arity: "SISO", + }); + }); + + test("it is offered to pipelines but kept out of the editor's tool list", () => { + const { result } = renderHook(() => useTranslatedToolCatalog()); + + // There is no interactive classify tool to open - it only means something inside a pipeline. + expect(result.current.regularTools.classify?.hiddenFromToolList).toBe(true); + expect( + filterToolRegistryByQuery(result.current.regularTools, "").some( + (ranked) => ranked.item[0] === "classify", + ), + ).toBe(false); + }); + + test("it does not re-classify by default", () => { + const { result } = renderHook(() => useTranslatedToolCatalog()); + + const config = result.current.regularTools.classify?.operationConfig; + + // The step is idempotent unless asked otherwise: a second run on a classified document + // would be a second engine call, and a second charge, for the same answer. + expect(config?.defaultParameters).toEqual({ reclassify: false }); + }); +}); diff --git a/frontend/editor/src/core/data/toolsTaxonomy.ts b/frontend/editor/src/core/data/toolsTaxonomy.ts index f9c2dde02f..740b8a13ec 100644 --- a/frontend/editor/src/core/data/toolsTaxonomy.ts +++ b/frontend/editor/src/core/data/toolsTaxonomy.ts @@ -29,6 +29,7 @@ import { ProprietaryToolId } from "@app/types/proprietaryToolId"; import { PrototypeToolId } from "@app/types/prototypeToolId"; export enum SubcategoryId { + AI = "ai", SIGNING = "signing", DOCUMENT_SECURITY = "documentSecurity", VERIFICATION = "verification", @@ -71,6 +72,8 @@ export type ToolRegistryEntry = { > | null; // Whether this tool supports automation (defaults to true) supportsAutomate?: boolean; + // Keep out of the editor's tool list: a step only a pipeline runs, with no UI to open. + hiddenFromToolList?: boolean; // Synonyms for search (optional) synonyms?: string[]; // Version status indicator (e.g., "alpha", "beta") @@ -90,6 +93,8 @@ export type ProprietaryToolRegistry = Record< export type PrototypeToolRegistry = Record; export const SUBCATEGORY_ORDER: SubcategoryId[] = [ + // First: AI steps are the ones a user is least likely to know exist. + SubcategoryId.AI, SubcategoryId.SIGNING, SubcategoryId.DOCUMENT_SECURITY, SubcategoryId.VERIFICATION, @@ -104,6 +109,7 @@ export const SUBCATEGORY_ORDER: SubcategoryId[] = [ ]; export const SUBCATEGORY_COLOR_MAP: Record = { + [SubcategoryId.AI]: "var(--category-color-automation)", // Pink [SubcategoryId.SIGNING]: "var(--category-color-signing)", // Green [SubcategoryId.DOCUMENT_SECURITY]: "var(--category-color-security)", // Orange [SubcategoryId.VERIFICATION]: "var(--category-color-verification)", // Orange diff --git a/frontend/editor/src/core/data/useTranslatedToolRegistry.tsx b/frontend/editor/src/core/data/useTranslatedToolRegistry.tsx index 5ae8075d0f..5e81427d55 100644 --- a/frontend/editor/src/core/data/useTranslatedToolRegistry.tsx +++ b/frontend/editor/src/core/data/useTranslatedToolRegistry.tsx @@ -20,6 +20,7 @@ import { import { adjustContrastOperationConfig } from "@app/hooks/tools/adjustContrast/useAdjustContrastOperation"; import { getSynonyms } from "@app/utils/toolSynonyms"; import { useProprietaryToolRegistry } from "@app/data/useProprietaryToolRegistry"; +import { classifyOperationConfig } from "@app/hooks/tools/classify/useClassifyOperation"; import { compressOperationConfig } from "@app/hooks/tools/compress/useCompressOperation"; import { splitOperationConfig } from "@app/hooks/tools/split/useSplitOperation"; import { addPasswordOperationConfig } from "@app/hooks/tools/addPassword/useAddPasswordOperation"; @@ -1338,6 +1339,30 @@ export function useTranslatedToolCatalog(): TranslatedToolCatalog { synonyms: getSynonyms(t, "compare"), supportsAutomate: false, }, + classify: { + icon: ( + + ), + name: t("home.classify.title", "Classify"), + // No interactive UI: this is a pipeline step, registered so a pipeline can name it. + component: null, + description: t( + "home.classify.desc", + "Identify what kind of document this is and tag it.", + ), + categoryId: ToolCategoryId.ADVANCED_TOOLS, + subcategoryId: SubcategoryId.AI, + maxFiles: -1, + endpoints: ["classify-and-label"], + operationConfig: asRegistryConfig(classifyOperationConfig), + automationSettings: null, + // Pipeline-only: there is no interactive classify tool to open in the editor. + hiddenFromToolList: true, + }, compress: { icon: ( The tool reads a window of the document, asks the AI engine what kind of document it is, and + * writes the verdict into the PDF's metadata. It has no interactive UI - there is nothing for a + * user to set beyond whether to redo work already done - so it exists in the registry purely so a + * pipeline can name it, the same way the Classification policy always could. + */ +const ENDPOINT = "/api/v1/ai/tools/classify-and-label" satisfies ToolEndpoint; +type ClassifyApiParams = ToolApiParams[typeof ENDPOINT]; + +export interface ClassifyParameters { + /** + * Classify again even when the document already carries a verdict. Off by default: re-running + * costs a second engine call, and a document that has been classified has nothing new to say. + */ + reclassify: boolean; +} + +export const defaultParameters: ClassifyParameters = { reclassify: false }; + +export const classifyToApiParams = ( + parameters: ClassifyParameters, +): ClassifyApiParams => ({ reclassify: parameters.reclassify }); + +export const classifyFromApiParams = ( + apiParams: ClassifyApiParams, +): Partial => ({ + reclassify: apiParams.reclassify ?? defaultParameters.reclassify, +}); + +export const buildClassifyFormData = ( + parameters: ClassifyParameters, + file: File, +): FormData => + objectToFormData(classifyToApiParams(parameters), { fileInput: file }); + +export const classifyOperationConfig = defineSingleFileTool({ + buildFormData: buildClassifyFormData, + toApiParams: classifyToApiParams, + fromApiParams: classifyFromApiParams, + operationType: "classify", + endpoint: ENDPOINT, + defaultParameters, +}); diff --git a/frontend/editor/src/core/services/fileStorage.ts b/frontend/editor/src/core/services/fileStorage.ts index 27b8746604..b0cede48cb 100644 --- a/frontend/editor/src/core/services/fileStorage.ts +++ b/frontend/editor/src/core/services/fileStorage.ts @@ -10,6 +10,7 @@ import { StirlingFile, StirlingFileStub, createStirlingFile, + type ClassificationConfidence, } from "@app/types/fileContext"; import { indexedDBManager, @@ -36,6 +37,8 @@ export interface StoredStirlingFileRecord extends BaseFileMetadata { // group by label without re-reading PDF bytes, and it survives versioning. // See StirlingFileStub.classificationLabels. classificationLabels?: string[]; + // See StirlingFileStub.classificationConfidence. + classificationConfidence?: ClassificationConfidence; } export interface StorageStats { @@ -695,6 +698,7 @@ class FileStorageService { folderId: record.folderId ?? null, createdAt: record.createdAt || Date.now(), classificationLabels: record.classificationLabels, + classificationConfidence: record.classificationConfidence, }; resolve(stub); @@ -762,6 +766,7 @@ class FileStorageService { folderId: record.folderId ?? null, createdAt: record.createdAt || Date.now(), classificationLabels: record.classificationLabels, + classificationConfidence: record.classificationConfidence, }); } cursor.continue(); @@ -861,6 +866,7 @@ class FileStorageService { folderId: record.folderId ?? null, createdAt: record.createdAt || Date.now(), classificationLabels: record.classificationLabels, + classificationConfidence: record.classificationConfidence, }); } cursor.continue(); diff --git a/frontend/editor/src/core/types/fileContext.ts b/frontend/editor/src/core/types/fileContext.ts index 3dc00c306d..c482e8e54a 100644 --- a/frontend/editor/src/core/types/fileContext.ts +++ b/frontend/editor/src/core/types/fileContext.ts @@ -9,6 +9,9 @@ import { generateId } from "@app/utils/generateId"; // Re-export FileId for convenience export type { FileId }; +/** How sure a classifier was about the labels it produced. */ +export type ClassificationConfidence = "none" | "low" | "medium" | "high"; + // Normalized state types export interface ProcessedFilePage { thumbnail?: string; @@ -61,6 +64,11 @@ export interface StirlingFileStub extends BaseFileMetadata { * unclassified files / non-SaaS builds. */ classificationLabels?: string[]; + /** + * How sure the local heuristic was about {@link classificationLabels}: a confident verdict + * stands, an unsure one escalates to the AI. Undefined when the labels came from the AI. + */ + classificationConfidence?: ClassificationConfidence; /** * This session proved the stored bytes unreadable (WebKit losing a blob's * backing store). The row renders as "data lost" instead of pretending the diff --git a/frontend/editor/src/core/types/toolApiTypes.ts b/frontend/editor/src/core/types/toolApiTypes.ts index cea05ca20a..55938d231e 100644 --- a/frontend/editor/src/core/types/toolApiTypes.ts +++ b/frontend/editor/src/core/types/toolApiTypes.ts @@ -207,6 +207,9 @@ export interface AddWatermarkRequest { */ widthSpacer?: number; } +export interface AiToolsClassifyAndLabelRequest { + reclassify?: boolean; +} export interface AutoRotatePdfRequest { /** * Minimum Tesseract OSD orientation confidence required before a correction is applied. Matches OCRmyPDF's --rotate-pages-threshold scale @@ -1485,6 +1488,7 @@ export interface UrlToPdfRequest { /** Endpoint path for a generated tool operation (the operation identity across languages). */ export type ToolEndpoint = + | "/api/v1/ai/tools/classify-and-label" | "/api/v1/convert/cbr/pdf" | "/api/v1/convert/cbz/pdf" | "/api/v1/convert/ebook/pdf" @@ -1587,6 +1591,7 @@ export type ToolEndpoint = /** Backend request-parameter model for each tool endpoint. */ export interface ToolApiParams { + "/api/v1/ai/tools/classify-and-label": AiToolsClassifyAndLabelRequest; "/api/v1/convert/cbr/pdf": ConvertCbrToPdfRequest; "/api/v1/convert/cbz/pdf": ConvertCbzToPdfRequest; "/api/v1/convert/ebook/pdf": ConvertEbookToPdfRequest; @@ -1690,6 +1695,7 @@ export interface ToolApiParams { /** Every generated tool endpoint, for iteration. */ export const TOOL_ENDPOINTS = [ + "/api/v1/ai/tools/classify-and-label", "/api/v1/convert/cbr/pdf", "/api/v1/convert/cbz/pdf", "/api/v1/convert/ebook/pdf", diff --git a/frontend/editor/src/core/types/toolIO.ts b/frontend/editor/src/core/types/toolIO.ts index 4a03a52102..9d672b2d6d 100644 --- a/frontend/editor/src/core/types/toolIO.ts +++ b/frontend/editor/src/core/types/toolIO.ts @@ -99,6 +99,11 @@ export interface ToolIOSpec { export type ToolIOTable = Partial>; export const TOOL_IO: ToolIOTable = { + "/api/v1/ai/tools/classify-and-label": { + accepts: ["PDF"], + produces: "PDF", + arity: "SISO", + }, "/api/v1/convert/cbr/pdf": { accepts: ["CBR"], produces: "PDF", diff --git a/frontend/editor/src/core/types/toolId.ts b/frontend/editor/src/core/types/toolId.ts index 13799ad869..0762c2e81d 100644 --- a/frontend/editor/src/core/types/toolId.ts +++ b/frontend/editor/src/core/types/toolId.ts @@ -50,6 +50,7 @@ export const CORE_REGULAR_TOOL_IDS = [ "removeCertSign", "unlockPDFForms", "compress", + "classify", "extractPages", "reorganizePages", "extractImages", diff --git a/frontend/editor/src/core/utils/toolSearch.ts b/frontend/editor/src/core/utils/toolSearch.ts index 41b40135a5..1a543974a3 100644 --- a/frontend/editor/src/core/utils/toolSearch.ts +++ b/frontend/editor/src/core/utils/toolSearch.ts @@ -15,7 +15,11 @@ export function filterToolRegistryByQuery( toolRegistry: Partial, query: string, ): RankedToolItem[] { - const entries = Object.entries(toolRegistry) as [ToolId, ToolRegistryEntry][]; + // The single funnel into the editor's tool list, so hiding here hides it everywhere the user + // browses - while getExecutableTools still offers it to a pipeline. + const entries = ( + Object.entries(toolRegistry) as [ToolId, ToolRegistryEntry][] + ).filter(([, tool]) => !tool?.hiddenFromToolList); if (!query.trim()) { return entries.map(([id, tool]) => ({ item: [id, tool] as [ToolId, ToolRegistryEntry], diff --git a/frontend/editor/src/proprietary/components/policies/useClientSideClassification.ts b/frontend/editor/src/proprietary/components/policies/useClientSideClassification.ts index aac4b7e3da..64fd849d87 100644 --- a/frontend/editor/src/proprietary/components/policies/useClientSideClassification.ts +++ b/frontend/editor/src/proprietary/components/policies/useClientSideClassification.ts @@ -1,5 +1,5 @@ -// With the AI engine off, the Classification policy runs here in the browser: -// each upload is labelled by the heuristic engine and metered for billing parity. +// The Classification policy's first pass: every upload is labelled locally before the AI is asked. +// The confidence reported here decides whether the AI is asked at all - see usePolicyAutoRun. import { useEffect, useRef, useState } from "react"; import { useAllFiles, useFileManagement } from "@app/contexts/FileContext"; @@ -7,7 +7,6 @@ import { useAppConfig } from "@app/contexts/AppConfigContext"; import { useIndexedDB } from "@app/contexts/IndexedDBContext"; import { fileStorage } from "@app/services/fileStorage"; import { useClassificationEnabled } from "@app/hooks/useClassificationEnabled"; -import { useAiEngineEnabled } from "@app/hooks/useAiEngineEnabled"; import { scheduleIdle } from "@app/utils/scheduleIdle"; import { usePolicies } from "@app/hooks/usePolicies"; import { classifyFileHeuristically } from "@app/services/heuristic/heuristicClassification"; @@ -15,12 +14,14 @@ import { meterClassificationRun } from "@app/services/classificationMeter"; import { isDispatched, markDispatched, + recordRunStart, + updateRun, } from "@app/components/policies/policyRunStore"; import type { FileId } from "@app/types/file"; import type { StirlingFile, StirlingFileStub } from "@app/types/fileContext"; +import type { HeuristicConfidence } from "@app/services/heuristic/types"; +import { CLASSIFICATION_CATEGORY_ID } from "@app/data/classificationPolicy"; -/** The category id of the Classification policy (see policyDefinitions). */ -const CLASSIFICATION_CATEGORY = "classification"; /** Files classified per idle pass, so a large library drains over several ticks. */ const CLASSIFY_BATCH = 3; /** How long to wait for an upload's bytes to land in IndexedDB (20 × 250ms ≈ 5s). @@ -47,9 +48,8 @@ export function useClientSideClassification(): void { const { bumpRevision } = useIndexedDB(); const { policies } = usePolicies(); const classificationEnabled = useClassificationEnabled(); - const aiEnabled = useAiEngineEnabled(); - // While app-config loads, aiEnabled reads false even on AI-on tenants; classifying - // in that window would double-run (and double-bill) files the server also labels. + // Still waited on: a verdict written before app-config lands would be acted on by the + // escalation decision before it knows whether the AI engine is even available. const { loading: configLoading } = useAppConfig(); // Files claimed this session, keyed id+lastModified so a new version is retried once. A claim is // taken synchronously right before classifying, so overlapping batches never double-classify. @@ -57,7 +57,9 @@ export function useClientSideClassification(): void { // Bumped after each batch to drain the next one. const [tick, setTick] = useState(0); - const policy = policies[CLASSIFICATION_CATEGORY]; + // TODO: keyed on the Classification CATEGORY, so a pipeline that merely contains a classify + // step gets no local pass - suppressing one step of a chain is not expressible today. + const policy = policies[CLASSIFICATION_CATEGORY_ID]; // Only when the admin has an active Classification policy - the same gate the AI path uses. const active = Boolean( policy?.configured && @@ -69,7 +71,8 @@ export function useClientSideClassification(): void { ); useEffect(() => { - if (configLoading || !classificationEnabled || aiEnabled || !active) { + // Runs whether or not the AI engine is on: it is the first pass either way, not a fallback. + if (configLoading || !classificationEnabled || !active) { return; } const claimKey = (s: StirlingFileStub) => @@ -95,17 +98,23 @@ export function useClientSideClassification(): void { // Re-validate at execution time - another batch may have claimed it since. if (claimed.current.has(key)) continue; claimed.current.add(key); - const labels = await classifyStub(stub.id as FileId, stub.name); + const verdict = await classifyStub( + stub.id as FileId, + stub.name, + stub.size ?? 0, + ); // Bytes never landed (file removed mid-wait): leave undelivered so a // reload (or new version) retries; the claim stops churn this session. - if (labels == null) continue; + if (verdict == null) continue; // Deliver unconditionally - a re-render must never discard a computed // (and already metered) result. Writes are idempotent. updateStirlingFileStub(stub.id as FileId, { - classificationLabels: labels, + classificationLabels: verdict.labels, + classificationConfidence: verdict.confidence, }); const ok = await fileStorage.updateFileMetadata(stub.id as FileId, { - classificationLabels: labels, + classificationLabels: verdict.labels, + classificationConfidence: verdict.confidence, }); if (ok) wrote = true; } @@ -122,7 +131,6 @@ export function useClientSideClassification(): void { fileStubs, active, classificationEnabled, - aiEnabled, configLoading, updateStirlingFileStub, bumpRevision, @@ -134,7 +142,8 @@ export function useClientSideClassification(): void { async function classifyStub( fileId: FileId, fileName: string, -): Promise { + fileSize: number, +): Promise<{ labels: string[]; confidence: HeuristicConfidence } | null> { let file: StirlingFile | null = null; for (let i = 0; i < FILE_WAIT_TRIES; i++) { file = await fileStorage.getStirlingFile(fileId).catch(() => null); @@ -149,10 +158,28 @@ async function classifyStub( } const debug = isClassificationDebug(); const startedAt = performance.now(); + // A local run is still a billable policy run, so it belongs in the activity feed; recorded only + // once the bytes are in hand, so a file whose bytes never land leaves no phantom row. + + // Read before recordRunStart, which takes the dispatch key itself and would otherwise always + // answer "already dispatched", silently stopping metering. + const alreadyMetered = isDispatched(CLASSIFICATION_CATEGORY_ID, fileId); + const runId = `local-${CLASSIFICATION_CATEGORY_ID}-${fileId}-${Date.now()}`; + recordRunStart({ + runId, + categoryId: CLASSIFICATION_CATEGORY_ID, + fileId: fileId as string, + fileName, + fileSize, + target: "local", + status: "RUNNING", + outputs: [], + error: null, + startedAt: Date.now(), + }); try { const result = await classifyFileHeuristically(file, { explain: debug }); const { labels } = result; - const alreadyMetered = isDispatched(CLASSIFICATION_CATEGORY, fileId); const ms = Math.round(performance.now() - startedAt); const verdict = labels.length > 0 @@ -174,12 +201,22 @@ async function classifyStub( labels, }); } - markDispatched(CLASSIFICATION_CATEGORY, fileId); - return labels; + markDispatched(CLASSIFICATION_CATEGORY_ID, fileId); + // Labels, no output file - the same settle shape the server-run classification uses. + updateRun(runId, { + status: "COMPLETED", + imported: true, + outputFileIds: [fileId as string], + }); + return { labels, confidence: result.confidence }; } catch (err) { // Never persist a verdict for an unreadable file - the failure may be // environmental, so it must stay eligible to retry (and meter) later. console.warn(`[Classify] ${fileName}: could not be read, will retry`, err); + updateRun(runId, { + status: "FAILED", + error: err instanceof Error ? err.message : String(err), + }); return null; } } diff --git a/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.batch.test.tsx b/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.batch.test.tsx index bc4fa54dff..08768f1aa7 100644 --- a/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.batch.test.tsx +++ b/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.batch.test.tsx @@ -1,5 +1,6 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import { renderHook, act } from "@testing-library/react"; +import type { ClassificationConfidence } from "@app/types/fileContext"; /** * Batch integration test (61 files, two chained upload policies) driving the real @@ -12,7 +13,11 @@ const FILE_COUNT = 61; // the workbench, mirrored into useAllFiles. consumeFiles mutates it in place // (input id → output id) exactly as the real silent reducer would. const mocks = vi.hoisted(() => ({ - workspace: [] as Array<{ id: string; classificationLabels?: string[] }>, + workspace: [] as Array<{ + id: string; + classificationLabels?: string[]; + classificationConfidence?: ClassificationConfidence; + }>, consumeSilentCalls: 0, consumeNonSilentCalls: 0, persistCalls: 0, @@ -117,10 +122,20 @@ function Harness() { return null; } +/** The heuristic verdict that escalates to the AI classifier; only "high" stands alone. */ +const LOW = "low" as const; + function replaceInWorkspace(inputIds: string[], outputIds: string[]) { + // A versioned output carries its input's heuristic verdict; the escalation decision is about the + // document, not about which step produced the current bytes. + const inherited = + mocks.workspace.find((s) => inputIds.includes(s.id)) + ?.classificationConfidence ?? LOW; mocks.workspace = mocks.workspace .filter((s) => !inputIds.includes(s.id)) - .concat(outputIds.map((id) => ({ id }))); + .concat( + outputIds.map((id) => ({ id, classificationConfidence: inherited })), + ); } beforeEach(() => { @@ -138,6 +153,7 @@ beforeEach(() => { mocks.workspace = Array.from({ length: FILE_COUNT }, (_, i) => ({ id: `file-${i}`, + classificationConfidence: LOW, })); mocks.listPolicyRuns.mockResolvedValue([]); diff --git a/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.race.test.tsx b/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.race.test.tsx index b67648e7fc..84b6f58bf7 100644 --- a/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.race.test.tsx +++ b/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.race.test.tsx @@ -1,5 +1,6 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; import { renderHook, act } from "@testing-library/react"; +import type { ClassificationConfidence } from "@app/types/fileContext"; /** * Mid-run race: classification is in flight (its labelled output is still @@ -18,6 +19,7 @@ const mocks = vi.hoisted(() => ({ sourceFileIds?: string[]; derivedFromTool?: boolean; classificationLabels?: string[]; + classificationConfidence?: ClassificationConfidence; }>, runStoredPolicy: vi.fn(), getPolicyRun: vi.fn(), @@ -114,7 +116,7 @@ beforeEach(() => { resetPolicyRuns(); vi.clearAllMocks(); - mocks.workspace = [{ id: "file-0" }]; + mocks.workspace = [{ id: "file-0", classificationConfidence: "low" }]; mocks.listPolicyRuns.mockResolvedValue([]); mocks.getStirlingFile.mockResolvedValue( diff --git a/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.ts b/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.ts index 7a2c4bda2e..8140691b6d 100644 --- a/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.ts +++ b/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.ts @@ -1,18 +1,6 @@ /** - * Auto-run controller: every enabled policy enforces on every uploaded file. - * Watches the session's files and fires a real backend run - * (`POST /api/v1/policies/{id}/run`) per file, polling it to completion and - * recording progress in {@link policyRunStore} for the activity feed. - * - * When several policies enforce on the same trigger they run as an ordered chain: - * the first fires on the upload, and each subsequent policy fires on the previous - * one's output once it lands — so their effects accumulate in the admin-defined - * order rather than racing to fork the same version. - * - * Headless — call it from {@link PolicyAutoRunController}, which is mounted once - * wherever the editor is open so enforcement happens regardless of whether the - * policy panel is on screen. Each (policy, file) pair runs exactly once (tracked - * in the run store), so re-renders and remounts don't re-fire. + * Headless auto-run controller: one backend run per (policy, file), fired exactly once and polled. + * Policies sharing a trigger run as an ordered chain so their effects accumulate. */ import { useCallback, useEffect, useMemo, useRef } from "react"; @@ -39,7 +27,12 @@ import { dispatchPaygLimitReached } from "@app/services/usageLimitBridge"; import type { FileId } from "@app/types/file"; import { createStirlingFilesAndStubs } from "@app/services/fileStubHelpers"; import { readClassificationLabelsFromFile } from "@app/services/fileClassification"; -import { isClassificationCategory } from "@app/data/policyCategories"; +import { + policyDeliversOutputFiles, + policyRequiresAiEngine, + policyRewritesDocument, + shouldDispatchToAi, +} from "@app/data/classificationPolicy"; import { acquireDispatchSlot, releaseDispatchSlot, @@ -68,10 +61,8 @@ const POLL_MS = 2000; * sitting on an indeterminate spinner for a full poll interval. */ const FIRST_POLL_MS = 500; -/** The server aborts any single tool step that runs longer than its internal-API - * read timeout, then fails the run — so a run can legitimately stay in flight - * for up to this long per step. The client must keep polling at least that long, - * or it abandons a run the server is still working on (which reads as a hang). */ +/** Server's per-step abort budget - poll at least this long per step, or we abandon + * a run the server is still working on. */ const STEP_TIMEOUT_MS = 300_000; /** Slack on top of the per-step budget: queueing before the first step starts and @@ -91,11 +82,8 @@ const POLICY_QUEUE_FULL = "POLICY_QUEUE_FULL"; const MAX_QUEUE_RETRIES = 5; const QUEUE_RETRY_BASE_MS = 4000; -/** Consecutive "run not found" responses before giving up. The run state lives - * in memory on the server, so a restart or a second instance behind the load - * balancer makes a live run's status return 404 — and it won't come back. We - * tolerate a brief blip (e.g. a poll racing a just-dispatched run, or one hop - * to an instance that hasn't seen it) then fail, rather than polling forever. */ +/** Consecutive 404s before failing. Run state is in-memory server-side, so a restart + * or a hop to another instance loses it permanently - tolerate a blip, not forever. */ const MAX_NOT_FOUND = 3; /** A 404 (run status gone, or output file gone), across the web (axios) and @@ -118,24 +106,20 @@ function failRun(runId: string, message: string): void { updateRun(runId, { status: "FAILED", error: message, errorCode: null }); } -/** How long to wait for an upload's bytes to land in IndexedDB before giving up - * (20 × 250ms ≈ 5s). The stub can surface in the file list a beat before its - * bytes are committed, so a too-eager fetch would otherwise miss the file. */ +/** Wait for an upload's bytes to land in IndexedDB (~5s): the stub surfaces in the + * file list before its bytes are committed, so an eager fetch would miss the file. */ const FILE_WAIT_TRIES = 20; const FILE_WAIT_MS = 250; -/** - * A policy that changed nothing (redaction matched no text, say) completes with no - * output: nothing to deliver, but finished. Left unimported, the file's badge and - * its blocking overlay spin forever. - */ +/** A policy that changed nothing completes with no output; left unimported its badge + * and blocking overlay spin forever. */ export function finishedWithNothingToDeliver(run: PolicyRunRecord): boolean { return ( run.status === "COMPLETED" && !run.imported && (run.outputs?.length ?? 0) === 0 && - // Classification has its own settle path: labels, no output file. - !isClassificationCategory(run.categoryId) + // An annotating policy settles on labels, not an output file. + policyDeliversOutputFiles(run.categoryId) ); } @@ -155,34 +139,23 @@ export function usePolicyAutoRun(): void { const { policies } = usePolicies(); const aiEnabled = useAiEngineEnabled(); const runs = usePolicyRuns(); - // Live view of the workspace files, read inside the import effect WITHOUT making - // it a dependency. The silent consume that delivers an output mutates fileStubs, - // so if the import effect depended on fileStubs it would re-fire on its own - // delivery — an infinite import cascade (and a bumpRevision storm that trips - // React's max-update-depth). The effect only needs to fire when `runs` changes. + // Read in the import effect via ref, not as a dependency: delivery mutates fileStubs, + // so depending on them would re-fire the effect on its own delivery (infinite cascade). const fileStubsRef = useRef(fileStubs); fileStubsRef.current = fileStubs; - // Keys (run ids / dispatch keys) currently in flight, so the effects never - // double-fire across re-renders while their first async step is pending. + // Keys in flight, so effects never double-fire across re-renders while async work pends. const polling = useRef>(new Set()); const importing = useRef>(new Set()); const dispatching = useRef>(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 - // limit sentinel on the run's errorCode; when a run we polled finishes blocked, broadcast a - // window event. A saas-layer listener (which can read the wallet + open the modal — this - // proprietary hook can't import the saas modal API) decides free-limit vs spend-cap. Dedupe per - // run so a folder-watch burst opens the modal once, not once per file. + // Server-side runs never hit the apiClient 402 interceptor, so we broadcast the limit + // sentinel for a saas listener to open the modal. Deduped per run. const firedLimitModal = useRef>(new Set()); - // Active upload policies in execution order. When several enforce on upload they - // run as a chain — the first fires on the upload, each subsequent one on the - // previous policy's output — so their effects accumulate in a defined order - // instead of racing to fork the same version. Mirrors the dispatch filter - // (incl. the editor-source gate) so the chain honours the same eligibility. + // Active upload policies in chain order, so effects accumulate instead of racing to fork + // the same version. Mirrors the dispatch filter so the chain honours the same eligibility. const orderedUploadCategories = useMemo( () => Object.entries(policies) @@ -195,41 +168,40 @@ export function usePolicyAutoRun(): void { s.sources.length === 0 || s.sources.includes("editor")) && (s.runOn ?? "upload") === "upload" && - // Non-AI systems classify in the browser (useClientSideClassification), so keep the - // Classification policy out of the server chain when the AI engine is off. - !(id === "classification" && !aiEnabled), + // An escalation-only policy has nothing to do with no engine to escalate to. + !(policyRequiresAiEngine(id) && !aiEnabled), ) - // Classification runs last: it's non-blocking, so an enforcement policy - // running after it would fork a new version and drop the user's edits. + // Annotating policies run last: a rewriting one after them would fork a new + // version from the pre-annotation state and drop their labels. .sort(([idA, a], [idB, b]) => { - const ca = isClassificationCategory(idA) ? 1 : 0; - const cb = isClassificationCategory(idB) ? 1 : 0; - if (ca !== cb) return ca - cb; + const ra = policyRewritesDocument(idA) ? 0 : 1; + const rb = policyRewritesDocument(idB) ? 0 : 1; + if (ra !== rb) return ra - rb; return (a.order ?? 0) - (b.order ?? 0); }) .map(([id]) => id), [policies, aiEnabled], ); - // Runs whose chain-continuation we've already handled this session, so the next - // policy is dispatched exactly once per completed run. + // Chain-continuations handled this session, so the next policy fires once per run. const chained = useRef>(new Set()); // Latest policies, read from inside the stable retry callback (which has no deps). const policiesRef = useRef(policies); policiesRef.current = policies; + // Latest stubs for the chaining effect, which keys off runs and must not depend on stubs. + const stubsRef = useRef(fileStubs); + stubsRef.current = fileStubs; // Per-file (dispatchKey) count of consecutive queue-rejection retries, so backoff escalates and // eventually gives up. Survives the run-id changing on each retry; reset on any real outcome. const queueRetries = useRef>(new Map()); - // A queue-rejected run is just backpressure — drop the rejected record and fire a fresh run in - // its place after a growing backoff (one feed row, not a new one per attempt). Once the budget is - // spent, leave the last failure standing so the activity feed offers a manual Retry. + // Queue rejection is backpressure: replace the record with a fresh run after a backoff, so the + // feed keeps one row. Budget spent, leave the failure standing for a manual Retry. 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. + // A reconciled run has no local fileId to re-dispatch; leave it failed. if (!rec.fileId) return; const key = dispatchKey(rec.categoryId, rec.fileId); const attempts = queueRetries.current.get(key) ?? 0; @@ -278,42 +250,38 @@ export function usePolicyAutoRun(): void { [scheduleQueueRetry], ); - // Dispatch: fire only the FIRST upload policy on each not-yet-run file. The rest - // of the chain is dispatched by the chaining effect below, each on the previous - // policy's output, so the policies apply cumulatively in order. + // Fire only the FIRST upload policy per file; the chaining effect below runs the rest + // on each previous output, so policies apply cumulatively in order. useEffect(() => { const firstCategory = orderedUploadCategories[0]; if (!firstCategory) return; const backendId = policies[firstCategory]?.backendId; if (!backendId) return; 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. + // Input-mode policies cover uploads only; tool-produced files are left to + // export-mode policies at export time. if (stub.derivedFromTool) continue; const key = dispatchKey(firstCategory, stub.id); - // Skip if already run (persisted) or a dispatch is in flight — the - // in-memory guard prevents double-firing during the async wait. + // Skip if already run (persisted) or in flight - the in-memory guard covers the async wait. if ( isDispatched(firstCategory, stub.id) || dispatching.current.has(key) ) { continue; } + // A confident local verdict stands; only an unsure one is escalated to the engine. + if (!shouldDispatchToAi(firstCategory, stub)) continue; dispatching.current.add(key); void runPolicyOnFile(firstCategory, backendId, stub.id, stub.name) .catch(() => { - // runPolicyOnFile handles its own failures; this is just a backstop - // so an unexpected rejection never becomes an unhandled rejection. + // Backstop: runPolicyOnFile handles its own failures. }) .finally(() => dispatching.current.delete(key)); } }, [fileStubs, policies, orderedUploadCategories]); - // Chain: once a run has completed AND its output landed in the workspace, fire the - // next upload policy on that output. Only chains on success (a failed run has no - // output), and only once per run. isDispatched guards re-dispatch across reloads. + // Once a run's output lands, fire the next upload policy on it - success only, once per + // run. isDispatched guards re-dispatch across reloads. useEffect(() => { for (const run of runs) { if (run.status !== "COMPLETED" || !run.imported) continue; @@ -332,10 +300,15 @@ export function usePolicyAutoRun(): void { // Next policy not ready yet (still reconciling) — retry when policies change. if (!backendId) continue; chained.current.add(run.runId); - // Chain onto EVERY output, not just the first — a run that produced multiple files (split, - // ZIP-unpacked) must apply the next policy to all of them, or outputs 2..N silently skip it. + // Chain onto EVERY output: a run that produced several files (split, ZIP-unpacked) + // would otherwise silently skip the next policy on outputs 2..N. for (const outputId of outputIds) { if (isDispatched(nextCategory, outputId as FileId)) continue; + const outputStub = stubsRef.current.find((s) => s.id === outputId); + // Nothing to escalate: either the heuristic already answered confidently, or it has + // not reported yet and this effect re-runs when the verdict lands. + if (outputStub && !shouldDispatchToAi(nextCategory, outputStub)) + continue; void runPolicyOnFile( nextCategory, backendId, @@ -345,7 +318,7 @@ export function usePolicyAutoRun(): void { ).catch(() => {}); } } - }, [runs, policies, orderedUploadCategories]); + }, [runs, policies, orderedUploadCategories, fileStubs]); // Poll each in-flight run to a terminal state. useEffect(() => { @@ -358,11 +331,10 @@ export function usePolicyAutoRun(): void { } }, [runs, onRunFinished]); - // Import each completed run's outputs into the workspace (each output once), - // so the enforced file appears in the app rather than only on the backend. + // Import each completed run's outputs once, so the enforced file appears in the app. useEffect(() => { for (const run of runs) { - const classification = isClassificationCategory(run.categoryId); + const deliversFiles = policyDeliversOutputFiles(run.categoryId); if ( run.status !== "COMPLETED" || run.imported || @@ -375,13 +347,10 @@ export function usePolicyAutoRun(): void { continue; } importing.current.add(run.runId); - // Classification is metadata-only: stamp labels onto the current leaf of - // the file it ran on (no version fork). See importClassificationLabels. - if (classification) { - // Targets are resolved by importClassificationLabels AT WRITE TIME (not - // snapshotted here): its download/parse is an async window during which - // a manual tool run can consume the input and fork a new leaf, and a - // stale snapshot would no-op on the dead id and lose the labels. + // An annotating policy writes labels onto the current leaf; no version fork. + if (!deliversFiles) { + // Resolved at write time, not snapshotted: a tool run during the async parse can fork + // a new leaf, and a stale id would no-op and lose the labels. void importClassificationLabels( run, () => @@ -390,8 +359,7 @@ export function usePolicyAutoRun(): void { ).finally(() => importing.current.delete(run.runId)); continue; } - // Honour the policy's output mode: a new file, or a new version of the - // input file it ran on (needs that input's stub, still in the workspace). + // Output mode: a new file, or a new version of the input (needs its stub in the workspace). const outputMode = policies[run.categoryId]?.outputMode ?? "new_version"; const outputName = policies[run.categoryId]?.outputName ?? ""; const outputNamePosition = policies[run.categoryId]?.outputNamePosition; @@ -410,9 +378,8 @@ export function usePolicyAutoRun(): void { firstUploadCategory: orderedUploadCategories[0], }).finally(() => importing.current.delete(run.runId)); } - // NB: fileStubs is intentionally NOT a dependency — it's read via a ref so a - // delivery's own workspace mutation can't re-trigger this effect (see the ref - // declaration above). The effect fires on run completions, which is all it needs. + // NB: fileStubs is read via a ref, not a dependency, so a delivery's own workspace + // mutation can't re-trigger this effect. }, [ runs, addFiles, @@ -422,11 +389,8 @@ export function usePolicyAutoRun(): void { orderedUploadCategories, ]); - // 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. + // The server owns runs, so rediscover any this client never recorded and let the effects + // above collect their outputs. Waits for policies so runs can be attributed to a category. useEffect(() => { if (reconciled.current) return; if (Object.keys(policies).length === 0) return; @@ -457,23 +421,18 @@ interface ImportContext { outputMode: "new_file" | "new_version"; /** Rename rule. Empty → keep the input's filename. */ outputName: string; - /** Where the rename is applied: before ("prefix") or after ("suffix") the - * base filename. Defaults to "suffix" when absent. */ + /** Rename position around the base filename; defaults to "suffix" when absent. */ outputNamePosition?: "prefix" | "suffix" | "auto-number"; /** The input file's stub — required to version it; absent if it's been removed. */ parentStub: StirlingFileStub | undefined; - /** The first upload policy in the chain — the only one the dispatch effect ever - * fires. Every policy output is marked dispatched for it so a downstream policy's - * output is never mistaken for a fresh upload and re-enforced (an endless loop). */ + /** The only policy the dispatch effect fires; every output is marked dispatched for it + * so a downstream output is never mistaken for a fresh upload and re-enforced. */ firstUploadCategory: string | 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 for feed visibility (polled if still live, but never auto-imported — see the - * `imported` note below). Server-excluded ad-hoc runs and runs we can't map to a configured - * category are skipped. + * Fold server-side runs into the local store: patch tracked ones, adopt untracked ones for feed + * visibility only. Unmappable and ad-hoc runs are skipped. */ function applyOutputName( inputFileName: string, @@ -489,8 +448,7 @@ function applyOutputName( : `${base}_${outputName}${ext}`; } -/** The next upload policy after {@code categoryId} in the chain, or undefined if - * it's last or no longer in the ordered set (e.g. paused since it ran). */ +/** Next upload policy in the chain, or undefined if last or no longer eligible. */ function nextUploadCategory( orderedUploadCategories: string[], categoryId: string, @@ -522,13 +480,11 @@ async function reconcileServerRuns( addReconciledRun({ runId: view.runId, categoryId, - // No local input link: a run rediscovered purely from the server was never recorded by - // this client, so it can't be tied back to a workspace/storage file (and isn't retried). + // Server-only run: never recorded here, so it can't be tied to a file (and isn't retried). fileId: "", fileName: view.outputs[0]?.fileName ?? "", fileSize: 0, - // Rediscovered from the SaaS run registry (listPolicyRuns), so its outputs - // live on the cloud backend. + // From the SaaS run registry, so its outputs live on the cloud backend. target: "saas", status: view.status, outputs: view.outputs, @@ -560,10 +516,8 @@ interface ClassificationImportContext { bumpRevision: () => void; } -/** Workspace stubs to tag with a classification run's labels: the file it ran - * on plus any live descendants, so an edit made during the async run (which - * forks a new leaf) still shows the tags. Empty once the document has left the - * workspace (closed, or a reconciled run with no local input link). */ +/** The run's file plus live descendants, so an edit during the run (which forks a new leaf) + * still shows the tags. Empty once the document has left the workspace. */ export function classificationLabelTargetStubs( runFileId: string, stubs: ReadonlyArray, @@ -576,19 +530,14 @@ export function classificationLabelTargetStubs( ); } -/** Attempts to read a completed run's labels before giving up, and the backoff - * between them (delay × attempt). The import effect only re-runs when the run - * store changes, so a transient read failure has to be retried HERE: bailing - * out would leave the run unsettled and the file's "running" pill spinning - * until unrelated policy activity happened to nudge the effect. */ +/** Label-read attempts and backoff. Retried HERE because the import effect only re-runs on + * run-store changes, so bailing out would leave the file's "running" pill spinning. */ const LABEL_READ_ATTEMPTS = 3; const LABEL_READ_RETRY_MS = 2000; /** - * Read classification labels out of a completed run's output PDF. A 404 means - * that output aged out, so it's skipped; any other failure is transient and - * retried with backoff. Returns null when there are genuinely no labels to - * apply (including a run with no outputs), so the caller can settle the run. + * Read labels from a completed run's output PDF: a 404 means it aged out and is skipped, other + * failures retry with backoff. Null means no labels to apply, so the caller settles the run. */ async function readRunLabels(run: PolicyRunRecord): Promise { for (let attempt = 0; attempt < LABEL_READ_ATTEMPTS; attempt++) { @@ -609,21 +558,13 @@ async function readRunLabels(run: PolicyRunRecord): Promise { // Every output was read (or had aged out): there are no labels to apply. if (!transientFailure) return null; } - // Out of attempts. Settle the run unlabelled rather than spin forever; the - // file keeps its classification badge, just without tags. + // Out of attempts: settle unlabelled rather than spin forever - the badge stays, tags don't. return null; } /** - * Stamp `labels` onto the run's live descendants in place (workspace + storage) - * — no versioned child, no history entry, only tags. Returns the tagged ids. - * - * Runs twice, because `resolveTargets` reads a rendered snapshot of the - * workspace: a CONSUME_FILES that was dispatched but not yet rendered when the - * first pass ran leaves its target already gone by the time UPDATE_FILE_RECORD - * is processed, so that stamp no-ops and the labels would be silently lost. The - * second pass sees the forked leaf and tags it. Each id is stamped at most once - * across both passes, so the pass costs nothing when no consume raced. + * Stamp `labels` in place (workspace + storage) - tags only, no versioned child. Two passes + * because a consume racing the first would strand its target and lose the labels. */ async function stampClassificationLabels( labels: string[], @@ -634,10 +575,8 @@ async function stampClassificationLabels( const tagged = new Set(); for (let pass = 0; pass < 2; pass++) { - // Resolve and stamp the store in one synchronous block — no await between - // them, so a target can't be consumed in between. A consume AFTER the stamp - // is safe too: the CONSUME_FILES reducer carries classificationLabels onto - // the new leaf. + // Resolve and stamp synchronously so no consume lands in between; a consume after the + // stamp is safe, as the reducer carries the labels onto the new leaf. const fresh = resolveTargets().filter((s) => !tagged.has(s.id)); for (const stub of fresh) { tagged.add(stub.id); @@ -651,25 +590,20 @@ async function stampClassificationLabels( } if (mutated) ctx.bumpRevision(); - // Yield a macrotask so React processes this pass's stamps (and any consume - // that raced them) before the next pass re-resolves. + // Yield a macrotask so React processes this pass's stamps before the next re-resolves. if (pass === 0) await new Promise((resolve) => setTimeout(resolve)); } return Array.from(tagged); } -/** - * Deliver a classification run: read its labels and tag the live document with - * them. Metadata-only — nothing is versioned. - */ +/** Deliver a classification run: read its labels and tag the live document. Nothing is versioned. */ async function importClassificationLabels( run: PolicyRunRecord, resolveTargets: () => StirlingFileStub[], ctx: ClassificationImportContext, ): Promise { if (resolveTargets().length === 0) { - // The document left the workspace (closed, or a server-reconciled run with - // no local input link) — nothing to tag. + // The document left the workspace - nothing to tag. updateRun(run.runId, { imported: true }); return; } @@ -678,9 +612,8 @@ async function importClassificationLabels( labels && labels.length > 0 ? await stampClassificationLabels(labels, resolveTargets, ctx) : []; - // Settle either way so it stops re-importing. outputFileIds are the TAGGED - // workspace files (no forked version), so their policy badge persists. Safe - // to chain-key on: classification is always last, so nothing chains off it. + // Settle either way so it stops re-importing. outputFileIds are the TAGGED files, so their + // badge persists; safe to chain-key on, as classification is always last. updateRun(run.runId, { imported: true, importedFileIds: run.outputs.map((o) => o.fileId), @@ -689,16 +622,8 @@ async function importClassificationLabels( } /** - * 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, - * so a partial failure retries only the missing files on a later tick and the - * ones that succeeded are never added twice. `imported` flips true only once - * every output has landed. - * - * Delivery honours the policy's output mode: "new_version" replaces the input - * file with a versioned child (its history chain), "new_file" adds the output - * as a standalone file. Versioning falls back to a new file if the input is - * gone (no parent stub). + * Deliver a run's outputs per-output, so a partial failure retries only the missing files and + * successes are never added twice. Honours the output mode; versioning needs the parent stub. */ async function importOutputs( run: PolicyRunRecord, @@ -711,9 +636,8 @@ async function importOutputs( return; } - // Keep the input's original filename unless a rename rule is set — without a - // rule the backend's auto-suffixed name (e.g. "_watermarked_sanitized") would - // otherwise rename every output. + // Keep the input's filename unless a rename rule is set, else the backend's auto-suffixed + // name renames every output. const targetName = ctx.outputName ? applyOutputName( run.fileName, diff --git a/frontend/editor/src/proprietary/components/viewer/Viewer.tsx b/frontend/editor/src/proprietary/components/viewer/Viewer.tsx index 1f04ea22ca..d5a2e7537d 100644 --- a/frontend/editor/src/proprietary/components/viewer/Viewer.tsx +++ b/frontend/editor/src/proprietary/components/viewer/Viewer.tsx @@ -8,7 +8,7 @@ import { usePolicyRuns, type PolicyRunRecord, } from "@app/components/policies/policyRunStore"; -import { isClassificationCategory } from "@app/data/policyCategories"; +import { isClassificationCategory } from "@app/data/classificationPolicy"; import { PolicyEnforcementOverlay } from "@app/components/viewer/PolicyEnforcementOverlay"; type SignatureOverlayPassThrough = Pick< diff --git a/frontend/editor/src/proprietary/data/classificationPolicy.test.ts b/frontend/editor/src/proprietary/data/classificationPolicy.test.ts new file mode 100644 index 0000000000..d3cb55c043 --- /dev/null +++ b/frontend/editor/src/proprietary/data/classificationPolicy.test.ts @@ -0,0 +1,91 @@ +import { describe, it, expect } from "vitest"; +import { + isClassificationCategory, + orderRewritesFirst, + policyDeliversOutputFiles, + policyRequiresAiEngine, + policyRewritesDocument, + shouldDispatchToAi, +} from "@app/data/classificationPolicy"; +import type { StirlingFileStub } from "@app/types/fileContext"; + +const stub = ( + confidence?: StirlingFileStub["classificationConfidence"], +): StirlingFileStub => + ({ classificationConfidence: confidence }) as StirlingFileStub; + +describe("isClassificationCategory", () => { + it("recognises the classification category and nothing else", () => { + expect(isClassificationCategory("classification")).toBe(true); + expect(isClassificationCategory("security")).toBe(false); + expect(isClassificationCategory("")).toBe(false); + }); +}); + +describe("policy capabilities", () => { + it("treats classification as annotating, everything else as rewriting", () => { + expect(policyRewritesDocument("security")).toBe(true); + expect(policyRewritesDocument("classification")).toBe(false); + // A builder pipeline (no catalogue category) runs tools, so it rewrites. + expect(policyRewritesDocument("pipeline-abc123")).toBe(true); + }); + + it("expects output files from rewriting policies only", () => { + expect(policyDeliversOutputFiles("security")).toBe(true); + expect(policyDeliversOutputFiles("classification")).toBe(false); + }); + + it("marks classification as the AI-escalation policy", () => { + expect(policyRequiresAiEngine("classification")).toBe(true); + expect(policyRequiresAiEngine("security")).toBe(false); + }); +}); + +describe("orderRewritesFirst", () => { + it("moves annotating policies to the end, preserving other order", () => { + expect( + orderRewritesFirst(["classification", "security", "compliance"]), + ).toEqual(["security", "compliance", "classification"]); + }); + + it("leaves an order without an annotating policy untouched", () => { + expect(orderRewritesFirst(["security", "compliance"])).toEqual([ + "security", + "compliance", + ]); + }); + + it("is a no-op when the annotating policy is already last", () => { + expect(orderRewritesFirst(["security", "classification"])).toEqual([ + "security", + "classification", + ]); + }); + + it("handles the annotating policy as the only one", () => { + expect(orderRewritesFirst(["classification"])).toEqual(["classification"]); + }); +}); + +describe("shouldDispatchToAi", () => { + it("always dispatches a policy that is not classification", () => { + expect(shouldDispatchToAi("security", stub())).toBe(true); + expect(shouldDispatchToAi("security", stub("high"))).toBe(true); + }); + + it("holds back until the local heuristic has reported", () => { + // Not a skip: dispatching now races the local pass and pays for a free answer; + // the caller re-evaluates once the verdict lands. + expect(shouldDispatchToAi("classification", stub())).toBe(false); + }); + + it("lets a confident local verdict stand", () => { + expect(shouldDispatchToAi("classification", stub("high"))).toBe(false); + }); + + it("escalates anything less than confident", () => { + expect(shouldDispatchToAi("classification", stub("medium"))).toBe(true); + expect(shouldDispatchToAi("classification", stub("low"))).toBe(true); + expect(shouldDispatchToAi("classification", stub("none"))).toBe(true); + }); +}); diff --git a/frontend/editor/src/proprietary/data/classificationPolicy.ts b/frontend/editor/src/proprietary/data/classificationPolicy.ts new file mode 100644 index 0000000000..4327b40d39 --- /dev/null +++ b/frontend/editor/src/proprietary/data/classificationPolicy.ts @@ -0,0 +1,69 @@ +/** + * Everything specific to the built-in Classification policy, in one module. The generic policy + * runner asks the capability questions below instead of naming classification itself, so a second + * annotating policy needs a change here rather than in the runner. + * + * These are still keyed on the category id rather than a property each policy declares. That is + * deliberate for now: policies are becoming pipelines with labels behind a separate enforcement + * layer, which removes the category concept these would be declared against. Classification also + * stays genuinely privileged - it is the only policy with a browser-side implementation, so it can + * answer without the server. Ordering and output shape belong in that rework (an in-place output + * mode, and a run result that can carry findings as well as files), not in a flag added here first. + */ + +import type { + ClassificationConfidence, + StirlingFileStub, +} from "@app/types/fileContext"; + +/** Catalogue category id of the built-in Classification policy. */ +export const CLASSIFICATION_CATEGORY_ID = "classification"; + +export function isClassificationCategory(categoryId: string): boolean { + return categoryId === CLASSIFICATION_CATEGORY_ID; +} + +/** + * Whether the policy rewrites the document rather than only annotating it. Annotating policies are + * ordered last: a rewriting one after them would fork from the pre-annotation version. + */ +export function policyRewritesDocument(categoryId: string): boolean { + return !isClassificationCategory(categoryId); +} + +/** Whether a completed run is expected to deliver output files (annotators deliver labels). */ +export function policyDeliversOutputFiles(categoryId: string): boolean { + return policyRewritesDocument(categoryId); +} + +/** Whether the policy's server-side run exists only to escalate to the AI engine. */ +export function policyRequiresAiEngine(categoryId: string): boolean { + return isClassificationCategory(categoryId); +} + +/** Order annotating policies last; everything else keeps the order it was given. */ +export function orderRewritesFirst(categoryIds: string[]): string[] { + return [ + ...categoryIds.filter(policyRewritesDocument), + ...categoryIds.filter((id) => !policyRewritesDocument(id)), + ]; +} + +/** + * The one heuristic verdict trusted to stand on its own; anything less escalates to the AI, which + * overwrites it. Deliberately strict - a wrong label costs more than an engine call. + */ +const TRUSTED_CONFIDENCE: ClassificationConfidence = "high"; + +/** + * Whether the AI classifier should be asked about this file. Only once the heuristic has reported: + * dispatching before then races the first pass and bills for an answer it was about to produce. + */ +export function shouldDispatchToAi( + categoryId: string, + stub: StirlingFileStub, +): boolean { + if (!isClassificationCategory(categoryId)) return true; + const confidence = stub.classificationConfidence; + return confidence != null && confidence !== TRUSTED_CONFIDENCE; +} diff --git a/frontend/editor/src/proprietary/data/policyCategories.test.ts b/frontend/editor/src/proprietary/data/policyCategories.test.ts deleted file mode 100644 index 4b4fca2b7e..0000000000 --- a/frontend/editor/src/proprietary/data/policyCategories.test.ts +++ /dev/null @@ -1,41 +0,0 @@ -import { describe, it, expect } from "vitest"; -import { - isClassificationCategory, - pinClassificationLast, -} from "@app/data/policyCategories"; - -describe("isClassificationCategory", () => { - it("recognises the classification category and nothing else", () => { - expect(isClassificationCategory("classification")).toBe(true); - expect(isClassificationCategory("security")).toBe(false); - expect(isClassificationCategory("")).toBe(false); - }); -}); - -describe("pinClassificationLast", () => { - it("moves classification to the end, preserving other order", () => { - expect( - pinClassificationLast(["classification", "security", "compliance"]), - ).toEqual(["security", "compliance", "classification"]); - }); - - it("leaves an order without classification untouched", () => { - expect(pinClassificationLast(["security", "compliance"])).toEqual([ - "security", - "compliance", - ]); - }); - - it("is a no-op when classification is already last", () => { - expect(pinClassificationLast(["security", "classification"])).toEqual([ - "security", - "classification", - ]); - }); - - it("handles classification as the only policy", () => { - expect(pinClassificationLast(["classification"])).toEqual([ - "classification", - ]); - }); -}); diff --git a/frontend/editor/src/proprietary/data/policyCategories.ts b/frontend/editor/src/proprietary/data/policyCategories.ts deleted file mode 100644 index 820a93366b..0000000000 --- a/frontend/editor/src/proprietary/data/policyCategories.ts +++ /dev/null @@ -1,21 +0,0 @@ -/** The classification policy's catalog category id. */ -export const CLASSIFICATION_CATEGORY_ID = "classification"; - -/** - * Classification is metadata-only: it runs async (never blocks), never forks a - * version, and always runs last. This predicate gates that special handling. - */ -export function isClassificationCategory(categoryId: string): boolean { - return categoryId === CLASSIFICATION_CATEGORY_ID; -} - -/** - * Move classification to the end of an execution order (others keep their order), - * so a persisted/displayed order can't place it anywhere but last. - */ -export function pinClassificationLast(orderedCategoryIds: string[]): string[] { - return [ - ...orderedCategoryIds.filter((id) => !isClassificationCategory(id)), - ...orderedCategoryIds.filter((id) => isClassificationCategory(id)), - ]; -} diff --git a/frontend/editor/src/proprietary/hooks/usePolicies.ts b/frontend/editor/src/proprietary/hooks/usePolicies.ts index 389f98cd65..382a7f9222 100644 --- a/frontend/editor/src/proprietary/hooks/usePolicies.ts +++ b/frontend/editor/src/proprietary/hooks/usePolicies.ts @@ -35,7 +35,7 @@ import { removePolicy, } from "@app/services/policyBackend"; import { reorderPolicies as reorderBackendPolicies } from "@app/services/policyApi"; -import { pinClassificationLast } from "@app/data/policyCategories"; +import { orderRewritesFirst } from "@app/data/classificationPolicy"; import type { PolicyToStore } from "@app/services/policyPipeline"; import type { PoliciesByCategory, @@ -327,9 +327,9 @@ export function usePolicies() { * first for an instant re-render; the next reconcile re-reads the server order. */ const reorderPolicies = useCallback((orderedCategoryIds: string[]) => { - // Pin classification last so the persisted/server order matches execution - // (it always runs last — see usePolicyAutoRun). - const ordered = pinClassificationLast(orderedCategoryIds); + // Annotating policies last, so the persisted order matches execution order + // (see usePolicyAutoRun). + const ordered = orderRewritesFirst(orderedCategoryIds); persistPolicyOrder(ordered); const current = loadPolicies(); const backendIds = ordered diff --git a/frontend/editor/src/proprietary/hooks/usePolicyFileBadges.ts b/frontend/editor/src/proprietary/hooks/usePolicyFileBadges.ts index 52779b9cf8..85b2249cbb 100644 --- a/frontend/editor/src/proprietary/hooks/usePolicyFileBadges.ts +++ b/frontend/editor/src/proprietary/hooks/usePolicyFileBadges.ts @@ -4,7 +4,7 @@ import type { PolicyRunRecord } from "@app/components/policies/policyRunStore"; import { useAllFiles } from "@app/contexts/FileContext"; import { loadPolicyCatalog } from "@app/services/policyCatalog"; import { policyAccentVar } from "@app/components/policies/policyStatus"; -import { isClassificationCategory } from "@app/data/policyCategories"; +import { isClassificationCategory } from "@app/data/classificationPolicy"; import type { FileItemPolicyRef } from "@app/components/shared/PolicyBadges"; /** Minimal provenance shape needed to resolve a file's inherited badges. */ diff --git a/frontend/editor/src/proprietary/services/heuristic/types.ts b/frontend/editor/src/proprietary/services/heuristic/types.ts index fbd8128ba4..5170feccd5 100644 --- a/frontend/editor/src/proprietary/services/heuristic/types.ts +++ b/frontend/editor/src/proprietary/services/heuristic/types.ts @@ -1,5 +1,7 @@ // Shared types for the client-side heuristic (non-AI) document classifier. +import type { ClassificationConfidence } from "@app/types/fileContext"; + /** Input document for the heuristic engine. */ export interface HeuristicDoc { fileName: string; @@ -11,7 +13,7 @@ export interface HeuristicDoc { } // "none" = no match or non-English; a real runtime value, not just a type state. -export type HeuristicConfidence = "none" | "low" | "medium" | "high"; +export type HeuristicConfidence = ClassificationConfidence; /** One scored candidate label with the rule hits that produced its score (debug only). */ export interface LabelScoreExplanation {