Harden AI engine PostHog telemetry against silent, permanent halts

The AI engine translates pydantic-ai OpenTelemetry spans into PostHog
`$ai_generation`/`$ai_trace` events via a custom SpanProcessor. Two gaps
in that path let LLM telemetry stop silently and stay at zero:

- `PostHogSpanProcessor.on_end` was unguarded. OpenTelemetry calls it
  synchronously from `Span.end()` with no exception handling of its own,
  and pydantic-ai ends the model span inside a `with` block, so any error
  while translating a span propagates into the model request itself. A
  telemetry defect could therefore break the agent run and/or halt
  emission for every subsequent span. It now translates and captures
  defensively and logs-and-drops on failure.

- The PostHog client was created without an `on_error` callback. Its
  `capture` is decorated `@no_throw` and delivery happens on a background
  consumer thread, so a failing or stalled delivery pipeline was
  completely invisible. We now log delivery failures and log the
  enabled/disabled state at startup, so a halt is noticeable in minutes
  instead of going unnoticed.

Adds unit tests covering the happy path, per-trace dedup, non-chat spans,
and that `on_end` never raises on delivery failure or a malformed span.

Generated-By: PostHog Code
Task-Id: 505d85dd-9cd2-4258-b76a-e027f4240277

app/allowed-licenses.json

@@ -208,6 +208,18 @@
       "moduleName": ".*",
       "moduleLicense": "The W3C License"
     },
+    {
+      "moduleName": "com.google.re2j:re2j",
+      "moduleLicense": "Go License"
+    },
+    {
+      "moduleName": "com.hubspot:algebra",
+      "moduleLicense": null
+    },
+    {
+      "moduleName": "com.hubspot.immutables:immutables-exceptions",
+      "moduleLicense": null
+    },
     {
       "moduleName": ".*",
       "moduleLicense": "UnRar License"

app/common/src/main/java/stirling/software/common/model/job/JobResult.java

@@ -4,8 +4,6 @@ import java.time.LocalDateTime;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CopyOnWriteArrayList;
 
 import com.fasterxml.jackson.annotation.JsonIgnore;
@@ -49,9 +47,6 @@ public class JobResult {
      */
     private final List<String> notes = new CopyOnWriteArrayList<>();
 
-    /** Key/value metadata that survives the write-through into the shared job store. */
-    private final Map<String, String> metadata = new ConcurrentHashMap<>();
-
     /**
      * Create a new JobResult with the given job ID
      *
@@ -166,16 +161,4 @@ public class JobResult {
     public List<String> getNotes() {
         return Collections.unmodifiableList(notes);
     }
-
-    /** Attach a metadata value, e.g. a policy id so cluster peers can identify a policy run. */
-    public void putMetadata(String key, String value) {
-        if (key != null && value != null) {
-            this.metadata.put(key, value);
-        }
-    }
-
-    /** An unmodifiable view of this job's metadata. */
-    public Map<String, String> getMetadata() {
-        return Collections.unmodifiableMap(metadata);
-    }
 }

app/common/src/main/java/stirling/software/common/service/TaskManager.java

@@ -230,18 +230,6 @@ public class TaskManager {
         return false;
     }
 
-    /** Attach metadata to a job and write it through to the shared store for cluster peers. */
-    public boolean putMetadata(String jobId, String key, String value) {
-        JobResult jobResult = jobResults.get(jobId);
-        if (jobResult != null) {
-            jobResult.putMetadata(key, value);
-            writeThrough(jobId, jobResult);
-            return true;
-        }
-        log.warn("Attempted to set metadata on non-existent job ID: {}", jobId);
-        return false;
-    }
-
     /**
      * Get statistics about all jobs in the system
      *
@@ -390,7 +378,7 @@ public class TaskManager {
                 fileIds.add(rf.getFileId());
             }
         }
-        Map<String, String> meta = new HashMap<>(result.getMetadata());
+        Map<String, String> meta = new HashMap<>();
         if (result.getNotes() != null && !result.getNotes().isEmpty()) {
             meta.put("notesCount", Integer.toString(result.getNotes().size()));
         }

app/core/src/main/resources/application.properties

@@ -96,8 +96,7 @@ spring.main.allow-bean-definition-overriding=true
 # spring-data-redis is on the classpath only for the optional Valkey backplane (which wires its own
 # factory); exclude Spring Boot's stock Redis auto-config so a default install doesn't create a dead
 # localhost:6379 factory that flips /actuator/health to DOWN.
-# Also exclude the repositories auto-config: in cluster mode it needs a redisTemplate bean we don't define.
-spring.autoconfigure.exclude=org.springframework.boot.data.redis.autoconfigure.DataRedisAutoConfiguration,org.springframework.boot.data.redis.autoconfigure.DataRedisReactiveAutoConfiguration,org.springframework.boot.data.redis.autoconfigure.DataRedisRepositoriesAutoConfiguration
+spring.autoconfigure.exclude=org.springframework.boot.data.redis.autoconfigure.DataRedisAutoConfiguration,org.springframework.boot.data.redis.autoconfigure.DataRedisReactiveAutoConfiguration
 
 # Set up a consistent temporary directory location
 java.io.tmpdir=${stirling.tempfiles.directory:${java.io.tmpdir}/stirling-pdf}

app/proprietary/build.gradle

@@ -66,6 +66,20 @@ dependencies {
 
     implementation "com.google.code.gson:gson:${gsonVersion}"
 
+    // jinjava/jjwt transitively request older Jackson 2 versions; declare the current
+    // version directly so it is selected consistently (root build.gradle pins are the fallback).
+    runtimeOnly "com.fasterxml.jackson.core:jackson-core:${jackson2Version}"
+    runtimeOnly "com.fasterxml.jackson.core:jackson-databind:${jackson2Version}"
+
+    implementation("com.hubspot.jinjava:jinjava:${jinjavaVersion}") {
+        // Compile-time-only annotation artifacts (class-retention annotations, not needed at
+        // runtime) whose declared licences (LGPL / none) fail the licence compatibility check.
+        exclude group: 'com.google.code.findbugs', module: 'annotations'
+        exclude group: 'org.derive4j', module: 'derive4j-annotation'
+        exclude group: 'com.hubspot.immutables', module: 'hubspot-style'
+        exclude group: 'com.hubspot.immutables', module: 'immutable-collection-encodings'
+    }
+
     api 'io.micrometer:micrometer-registry-prometheus'
 
     api "io.jsonwebtoken:jjwt-api:${jwtVersion}"

app/proprietary/src/main/java/stirling/software/proprietary/controller/api/CreatePdfAgentController.java

@@ -8,12 +8,14 @@ import java.util.List;
 
 import org.apache.pdfbox.pdmodel.PDDocument;
 import org.springframework.core.io.Resource;
+import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.PostMapping;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.RestController;
+import org.springframework.web.server.ResponseStatusException;
 
 import io.github.pixee.security.Filenames;
 import io.swagger.v3.oas.annotations.Hidden;
@@ -24,18 +26,24 @@ import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 
 import stirling.software.common.configuration.RuntimePathConfig;
+import stirling.software.common.model.ApplicationProperties;
 import stirling.software.common.service.CustomPDFDocumentFactory;
 import stirling.software.common.util.ProcessExecutor;
 import stirling.software.common.util.TempFile;
 import stirling.software.common.util.TempFileManager;
 import stirling.software.common.util.WebResponseUtils;
+import stirling.software.proprietary.model.api.ai.create.AiDocument;
+import stirling.software.proprietary.service.AiDocumentHtmlRenderer;
+
+import tools.jackson.core.JacksonException;
+import tools.jackson.databind.ObjectMapper;
 
 /**
- * Dispatchable tool that converts an AI-generated HTML string to a PDF via WeasyPrint.
+ * Dispatchable tool that converts an AI-generated document model to a PDF via WeasyPrint.
  *
  * <p>Called by {@link stirling.software.proprietary.service.AiWorkflowService} when the engine
- * emits a {@code CREATE_PDF_FROM_HTML_AGENT} plan step. The HTML comes from a trusted Jinja
- * template so sanitization is intentionally skipped.
+ * emits a {@code CREATE_PDF_FROM_HTML_AGENT} plan step. The engine supplies the document as
+ * structured fields; the HTML is built here from a fixed template.
  */
 @Slf4j
 @Hidden
@@ -48,6 +56,9 @@ public class CreatePdfAgentController {
     private final TempFileManager tempFileManager;
     private final CustomPDFDocumentFactory pdfDocumentFactory;
     private final RuntimePathConfig runtimePathConfig;
+    private final ApplicationProperties applicationProperties;
+    private final ObjectMapper objectMapper;
+    private final AiDocumentHtmlRenderer htmlRenderer;
 
     /**
      * Returns true only when WeasyPrint is definitively unavailable — either the binary could not
@@ -74,32 +85,42 @@ public class CreatePdfAgentController {
             value = "/create-pdf-from-html-agent",
             consumes = MediaType.MULTIPART_FORM_DATA_VALUE)
     @Operation(
-            summary = "Convert AI-generated HTML to a PDF",
+            summary = "Convert an AI-generated document to a PDF",
             description =
-                    "Accepts an HTML document as a plain-text parameter and returns a PDF."
-                            + " This endpoint is dispatched by the AI workflow orchestrator as a"
-                            + " plan step; it is not intended for direct client use.")
-    public ResponseEntity<Resource> createPdfFromHtml(
-            @RequestParam("htmlContent") String htmlContent,
-            @RequestParam("filename") String filename)
+                    "Accepts a structured document as a JSON parameter and returns a PDF. This"
+                            + " endpoint is dispatched by the AI workflow orchestrator as a plan"
+                            + " step; it is not intended for direct client use.")
+    public ResponseEntity<Resource> createPdf(
+            @RequestParam("document") String document, @RequestParam("filename") String filename)
             throws Exception {
 
+        if (!applicationProperties.getAiEngine().isEnabled()) {
+            throw new ResponseStatusException(HttpStatus.NOT_FOUND);
+        }
+
+        AiDocument model;
+        try {
+            model = objectMapper.readValue(document, AiDocument.class);
+        } catch (JacksonException e) {
+            throw new ResponseStatusException(HttpStatus.BAD_REQUEST);
+        }
+
+        String html = htmlRenderer.render(model);
+
         log.info(
-                "[create-pdf-agent] converting HTML to PDF via WeasyPrint — html_bytes={}",
-                htmlContent.length());
+                "[create-pdf-agent] converting document to PDF via WeasyPrint — html_bytes={}",
+                html.length());
 
         try (TempFile htmlFile = tempFileManager.createManagedTempFile(".html");
                 TempFile pdfFile = tempFileManager.createManagedTempFile(".pdf")) {
 
-            Files.writeString(htmlFile.getPath(), htmlContent, StandardCharsets.UTF_8);
+            Files.writeString(htmlFile.getPath(), html, StandardCharsets.UTF_8);
 
             List<String> command = new ArrayList<>();
             command.add(runtimePathConfig.getWeasyPrintPath());
             command.add("-e");
             command.add("utf-8");
             command.add("-v");
-            // SSRF: the HTML is self-contained and the engine validates style colours, so no
-            // external url() reaches WeasyPrint. For full isolation, run it network-isolated.
             command.add(htmlFile.getAbsolutePath());
             command.add(pdfFile.getAbsolutePath());
 
@@ -126,8 +147,8 @@ public class CreatePdfAgentController {
             // avoids materialising the whole document as a byte[] twice (read-all + re-serialise),
             // which matters for large generated documents.
             TempFile tempOut = tempFileManager.createManagedTempFile(".pdf");
-            try (PDDocument document = pdfDocumentFactory.load(pdfFile.getPath())) {
-                document.save(tempOut.getPath().toFile());
+            try (PDDocument pdDocument = pdfDocumentFactory.load(pdfFile.getPath())) {
+                pdDocument.save(tempOut.getPath().toFile());
             } catch (Exception e) {
                 tempOut.close();
                 throw e;

app/proprietary/src/main/java/stirling/software/proprietary/integration/crypto/CredentialEncryption.java

@@ -44,13 +44,10 @@ public class CredentialEncryption {
     private static volatile SecretKey key;
 
     private final String configuredKey;
-    private final boolean clusterEnabled;
 
     public CredentialEncryption(
-            @Value("${stirling.security.credentialEncryptionKey:}") String configuredKey,
-            @Value("${cluster.enabled:false}") boolean clusterEnabled) {
+            @Value("${stirling.security.credentialEncryptionKey:}") String configuredKey) {
         this.configuredKey = configuredKey;
-        this.clusterEnabled = clusterEnabled;
     }
 
     @PostConstruct
@@ -67,14 +64,6 @@ public class CredentialEncryption {
         if (configured != null && !configured.isBlank()) {
             return new SecretKeySpec(Base64.getDecoder().decode(configured.trim()), ALGORITHM);
         }
-        // Cluster nodes must share this key, so fail fast rather than generate a node-local one.
-        if (clusterEnabled) {
-            throw new IllegalStateException(
-                    "cluster.enabled=true requires a shared credential encryption key. Set"
-                            + " STIRLING_CREDENTIAL_ENCRYPTION_KEY (or"
-                            + " stirling.security.credentialEncryptionKey) to the same value on every"
-                            + " node.");
-        }
         return loadOrCreateKeyFile();
     }
 

app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/create/AiDocument.java

@@ -0,0 +1,35 @@
+package stirling.software.proprietary.model.api.ai.create;
+
+import java.util.List;
+
+import lombok.Data;
+
+@Data
+public class AiDocument {
+
+    private String title;
+    private String subtitle;
+    private String referenceNumber;
+    private Style style;
+    private List<Section> sections;
+
+    @Data
+    public static class Style {
+        private String primaryColor;
+        private String backgroundColor;
+        private String bodyTextColor;
+    }
+
+    @Data
+    public static class Section {
+        private String type;
+        private String heading;
+        private String body;
+        private List<List<String>> pairs;
+        private List<String> columns;
+        private List<List<String>> rows;
+        private List<String> totalRow;
+        private List<String> items;
+        private List<String> signatories;
+    }
+}

app/proprietary/src/main/java/stirling/software/proprietary/policy/controller/PolicyController.java

@@ -6,7 +6,6 @@ import java.util.Comparator;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBooleanProperty;
 import org.springframework.core.io.FileSystemResource;
@@ -41,8 +40,6 @@ import jakarta.validation.Valid;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 
-import stirling.software.common.cluster.JobStore;
-import stirling.software.common.cluster.JobStoreEntry;
 import stirling.software.common.model.ApplicationProperties;
 import stirling.software.common.model.job.JobResponse;
 import stirling.software.common.service.JobOwnershipService;
@@ -107,8 +104,6 @@ public class PolicyController {
     private final ApplicationProperties applicationProperties;
     private final TempFileManager tempFileManager;
     private final JobOwnershipService jobOwnershipService;
-    // Shared job store: lets the run endpoints see runs that executed on other nodes.
-    private final JobStore jobStore;
 
     @PostMapping(value = "/run", consumes = MediaType.MULTIPART_FORM_DATA_VALUE)
     @Operation(
@@ -179,19 +174,10 @@ public class PolicyController {
             description = "Returns the current status, step cursor, and output files of a run.")
     public ResponseEntity<PolicyRunView> status(@PathVariable String runId) {
         PolicyRun run = runRegistry.get(runId);
-        if (run != null) {
-            return ResponseEntity.ok(PolicyRunView.of(run));
+        if (run == null) {
+            return ResponseEntity.notFound().build();
         }
-        // Not local: read the run's shared projection so any node can serve its status.
-        if (ownedByCurrentUser(runId)) {
-            Optional<JobStoreEntry> entry = jobStore.get(runId);
-            if (entry.isPresent()
-                    && entry.get().resultMeta() != null
-                    && entry.get().resultMeta().containsKey("policyId")) {
-                return ResponseEntity.ok(PolicyRunView.ofEntry(entry.get()));
-            }
-        }
-        return ResponseEntity.notFound().build();
+        return ResponseEntity.ok(PolicyRunView.of(run));
     }
 
     @GetMapping("/runs")
@@ -204,26 +190,11 @@ public class PolicyController {
                             + " collected, rather than orphaned on the backend. Ad-hoc runs (no"
                             + " policy id) are excluded.")
     public List<PolicyRunView> listRuns() {
-        // Local runs first (they carry live step state); keyed by runId to dedupe shared entries.
-        Map<String, PolicyRunView> byRunId = new LinkedHashMap<>();
-        runRegistry.all().stream()
+        return runRegistry.all().stream()
                 .filter(run -> run.getPolicyId() != null)
                 .filter(run -> ownedByCurrentUser(run.getRunId()))
-                .forEach(run -> byRunId.put(run.getRunId(), PolicyRunView.of(run)));
-        // Then runs from other nodes, read from the shared job store.
-        for (JobStoreEntry entry : jobStore.all()) {
-            if (byRunId.containsKey(entry.jobId())) {
-                continue;
-            }
-            Map<String, String> meta = entry.resultMeta();
-            if (meta == null || !meta.containsKey("policyId")) {
-                continue; // ad-hoc job, not a stored-policy run
-            }
-            if (ownedByCurrentUser(entry.jobId())) {
-                byRunId.put(entry.jobId(), PolicyRunView.ofEntry(entry));
-            }
-        }
-        return List.copyOf(byRunId.values());
+                .map(PolicyRunView::of)
+                .toList();
     }
 
     /**

app/proprietary/src/main/java/stirling/software/proprietary/policy/engine/PolicyEngine.java

@@ -136,10 +136,6 @@ public class PolicyEngine {
         // ownership check passes. No-op when security is off.
         String runId = jobOwnershipService.createScopedJobKey(UUID.randomUUID().toString());
         taskManager.createTask(runId);
-        // Tag the shared job entry with the policy id so peers can list it as a policy run.
-        if (policyId != null) {
-            taskManager.putMetadata(runId, "policyId", policyId);
-        }
         PolicyRun run = new PolicyRun(runId, policyId, definition);
         registry.register(run);
         CompletableFuture<PolicyRun> completion = new CompletableFuture<>();

app/proprietary/src/main/java/stirling/software/proprietary/policy/model/PolicyRunView.java

@@ -1,9 +1,7 @@
 package stirling.software.proprietary.policy.model;
 
 import java.util.List;
-import java.util.Map;
 
-import stirling.software.common.cluster.JobStoreEntry;
 import stirling.software.common.model.job.ResultFile;
 
 /**
@@ -36,33 +34,4 @@ public record PolicyRunView(
                 run.getOutputs(),
                 run.getCreatedAt().toEpochMilli());
     }
-
-    /** Cross-node view from a shared job-store entry; step cursor is node-local so it reads 0. */
-    public static PolicyRunView ofEntry(JobStoreEntry entry) {
-        Map<String, String> meta = entry.resultMeta() == null ? Map.of() : entry.resultMeta();
-        PolicyRunStatus status =
-                switch (entry.state()) {
-                    case COMPLETE -> PolicyRunStatus.COMPLETED;
-                    case FAILED -> PolicyRunStatus.FAILED;
-                    case RUNNING, PENDING -> PolicyRunStatus.RUNNING;
-                };
-        List<ResultFile> outputs =
-                entry.fileIds() == null
-                        ? List.of()
-                        : entry.fileIds().stream()
-                                .map(id -> ResultFile.builder().fileId(id).build())
-                                .toList();
-        long createdAt = entry.createdAt() == null ? 0L : entry.createdAt().toEpochMilli();
-        return new PolicyRunView(
-                entry.jobId(),
-                meta.get("policyId"),
-                status,
-                0,
-                0,
-                entry.error(),
-                null,
-                null,
-                outputs,
-                createdAt);
-    }
 }

app/proprietary/src/main/java/stirling/software/proprietary/security/model/JwtSigningKeyEntity.java

@@ -1,52 +0,0 @@
-package stirling.software.proprietary.security.model;
-
-import java.io.Serializable;
-import java.time.LocalDateTime;
-
-import org.hibernate.annotations.CreationTimestamp;
-
-import jakarta.persistence.Column;
-import jakarta.persistence.Convert;
-import jakarta.persistence.Entity;
-import jakarta.persistence.Id;
-import jakarta.persistence.Table;
-
-import lombok.Getter;
-import lombok.NoArgsConstructor;
-import lombok.Setter;
-
-import stirling.software.proprietary.integration.crypto.EncryptedStringConverter;
-
-/** A JWT signing keypair in the shared DB; the private key is encrypted at rest. */
-@Entity
-@Table(name = "jwt_signing_keys")
-@NoArgsConstructor
-@Getter
-@Setter
-public class JwtSigningKeyEntity implements Serializable {
-
-    private static final long serialVersionUID = 1L;
-
-    @Id
-    @Column(name = "key_id", length = 128)
-    private String keyId;
-
-    // Base64 X.509 public key (non-secret).
-    @Column(name = "verifying_key", columnDefinition = "text", nullable = false)
-    private String verifyingKey;
-
-    // Base64 PKCS#8 private key, encrypted at rest.
-    @Convert(converter = EncryptedStringConverter.class)
-    @Column(name = "signing_key", columnDefinition = "text", nullable = false)
-    private String signingKey;
-
-    @CreationTimestamp
-    @Column(name = "created_at", updatable = false)
-    private LocalDateTime createdAt;
-
-    public JwtSigningKeyEntity(String keyId, String verifyingKey, String signingKey) {
-        this.keyId = keyId;
-        this.verifyingKey = verifyingKey;
-        this.signingKey = signingKey;
-    }
-}

app/proprietary/src/main/java/stirling/software/proprietary/security/repository/JwtSigningKeyRepository.java

@@ -1,26 +0,0 @@
-package stirling.software.proprietary.security.repository;
-
-import java.time.LocalDateTime;
-import java.util.List;
-import java.util.Optional;
-
-import org.springframework.data.jpa.repository.JpaRepository;
-import org.springframework.stereotype.Repository;
-
-import stirling.software.proprietary.security.model.JwtSigningKeyEntity;
-
-/** Shared-DB store of JWT signing keys - the source of truth every cluster node reads from. */
-@Repository
-public interface JwtSigningKeyRepository extends JpaRepository<JwtSigningKeyEntity, String> {
-
-    /** Newest first, so the most recently created key is the active signing key. */
-    List<JwtSigningKeyEntity> findAllByOrderByCreatedAtDesc();
-
-    /**
-     * The current active signing key: the single newest row. Used for cheap cluster convergence.
-     */
-    Optional<JwtSigningKeyEntity> findFirstByOrderByCreatedAtDesc();
-
-    /** Keys created before the cutoff, eligible for rotation cleanup. */
-    List<JwtSigningKeyEntity> findByCreatedAtBefore(LocalDateTime cutoff);
-}

app/proprietary/src/main/java/stirling/software/proprietary/security/service/KeyPairCleanupService.java

@@ -3,10 +3,8 @@ package stirling.software.proprietary.security.service;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.time.Duration;
 import java.time.LocalDateTime;
 import java.util.List;
-import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 
 import org.springframework.beans.factory.annotation.Autowired;
@@ -19,8 +17,6 @@ import jakarta.annotation.PostConstruct;
 
 import lombok.extern.slf4j.Slf4j;
 
-import stirling.software.common.cluster.DistributedLock;
-import stirling.software.common.cluster.DistributedLock.LockHandle;
 import stirling.software.common.configuration.InstallationPathConfig;
 import stirling.software.common.model.ApplicationProperties;
 import stirling.software.proprietary.security.model.JwtVerificationKey;
@@ -30,23 +26,15 @@ import stirling.software.proprietary.security.model.JwtVerificationKey;
 @ConditionalOnBooleanProperty("v2")
 public class KeyPairCleanupService {
 
-    // Cluster-wide single-writer: keys live in the shared DB, so only one node may prune + rotate
-    // per cycle. Otherwise every node runs this and they race to delete each other's keys.
-    private static final String CLEANUP_LOCK = "jwt-key-cleanup";
-    private static final Duration LOCK_LEASE = Duration.ofMinutes(5);
-
     private final KeyPersistenceService keyPersistenceService;
     private final ApplicationProperties.Security.Jwt jwtProperties;
-    private final DistributedLock distributedLock;
 
     @Autowired
     public KeyPairCleanupService(
             KeyPersistenceService keyPersistenceService,
-            ApplicationProperties applicationProperties,
-            DistributedLock distributedLock) {
+            ApplicationProperties applicationProperties) {
         this.keyPersistenceService = keyPersistenceService;
         this.jwtProperties = applicationProperties.getSecurity().getJwt();
-        this.distributedLock = distributedLock;
     }
 
     @Transactional
@@ -56,28 +44,7 @@ public class KeyPairCleanupService {
         if (!jwtProperties.isEnableKeyCleanup() || !keyPersistenceService.isKeystoreEnabled()) {
             return;
         }
-        // A lock-backend error must never fail this @PostConstruct/scheduled run: degrade to
-        // "skip this cycle" so a transient Valkey blip can't stop a node from booting.
-        Optional<LockHandle> lock;
-        try {
-            lock = distributedLock.tryAcquire(CLEANUP_LOCK, LOCK_LEASE);
-        } catch (RuntimeException e) {
-            log.warn(
-                    "Could not acquire the JWT key-cleanup lock ({}); skipping this cycle",
-                    e.getMessage());
-            return;
-        }
-        // No lock means another node is already pruning; skip until the next tick.
-        if (lock.isEmpty()) {
-            log.debug("Another node holds the JWT key-cleanup lock; skipping this cycle");
-            return;
-        }
-        try (LockHandle held = lock.get()) {
-            runCleanup();
-        }
-    }
 
-    private void runCleanup() {
         LocalDateTime cutoffDate =
                 LocalDateTime.now().minusDays(jwtProperties.getKeyRetentionDays());
 

app/proprietary/src/main/java/stirling/software/proprietary/security/service/KeyPersistenceService.java

@@ -18,17 +18,15 @@ import java.time.LocalDateTime;
 import java.time.format.DateTimeFormatter;
 import java.util.Base64;
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
-import java.util.UUID;
-import java.util.concurrent.ConcurrentHashMap;
 
-import org.springframework.beans.factory.annotation.Value;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.cache.Cache;
 import org.springframework.cache.CacheManager;
-import org.springframework.context.annotation.DependsOn;
-import org.springframework.scheduling.annotation.Scheduled;
+import org.springframework.cache.annotation.CacheEvict;
+import org.springframework.cache.caffeine.CaffeineCache;
 import org.springframework.stereotype.Service;
+import org.springframework.transaction.annotation.Transactional;
 
 import jakarta.annotation.PostConstruct;
 
@@ -36,39 +34,27 @@ import lombok.extern.slf4j.Slf4j;
 
 import stirling.software.common.configuration.InstallationPathConfig;
 import stirling.software.common.model.ApplicationProperties;
-import stirling.software.proprietary.security.model.JwtSigningKeyEntity;
 import stirling.software.proprietary.security.model.JwtVerificationKey;
-import stirling.software.proprietary.security.repository.JwtSigningKeyRepository;
 
-/** Persists JWT signing keys in the shared DB so all nodes sign/verify with the same key. */
 @Slf4j
 @Service
-// CredentialEncryption must init first: startup persists an encrypted private key.
-@DependsOn("credentialEncryption")
 public class KeyPersistenceService implements KeyPersistenceServiceInterface {
 
     public static final String KEY_SUFFIX = ".key";
     public static final String PUB_KEY_SUFFIX = ".pub";
 
     private final ApplicationProperties.Security.Jwt jwtProperties;
+    private final CacheManager cacheManager;
     private final Cache verifyingKeyCache;
-    private final JwtSigningKeyRepository keyRepository;
-    private final boolean clusterEnabled;
-
-    // kid -> KeyPair; safe to cache since key material is immutable.
-    private final Map<String, KeyPair> keyPairCache = new ConcurrentHashMap<>();
 
     private volatile JwtVerificationKey activeKey;
 
+    @Autowired
     public KeyPersistenceService(
-            ApplicationProperties applicationProperties,
-            CacheManager cacheManager,
-            JwtSigningKeyRepository keyRepository,
-            @Value("${cluster.enabled:false}") boolean clusterEnabled) {
+            ApplicationProperties applicationProperties, CacheManager cacheManager) {
         this.jwtProperties = applicationProperties.getSecurity().getJwt();
+        this.cacheManager = cacheManager;
         this.verifyingKeyCache = cacheManager.getCache("verifyingKeys");
-        this.keyRepository = keyRepository;
-        this.clusterEnabled = clusterEnabled;
     }
 
     @PostConstruct
@@ -77,86 +63,138 @@ public class KeyPersistenceService implements KeyPersistenceServiceInterface {
             log.info("JWT keystore is disabled - keys will be generated in memory");
             return;
         }
+
         try {
-            importLegacyDiskKeysIfPresent();
-            loadKeysFromDb();
+            ensurePrivateKeyDirectoryExists();
+            loadExistingKeysFromDisk();
         } catch (Exception e) {
-            log.error("Failed to initialize keystore, generating a fresh keypair", e);
-            generateAndStoreKeypair();
+            log.error("Failed to initialize keystore, using in-memory generation", e);
         }
     }
 
     /**
-     * Cluster convergence: adopt the newest signing key in the shared DB as this node's active key.
-     * Runs on every node so a key a peer just minted becomes the shared active signer within one
-     * interval, keeping cluster rotation equivalent to single-node. Cluster-only: a single node
-     * always holds its own newest key, so this is skipped entirely off-cluster.
+     * Load all existing JWT keys from disk into memory on startup.
+     *
+     * <p>This ensures tokens signed with previous keys remain valid after server restart. If no
+     * keys exist on disk, generates a new keypair.
      */
-    @Scheduled(fixedDelayString = "${stirling.security.jwt.activeKeyReloadMs:300000}")
-    public void reloadActiveKeyFromDb() {
-        if (!clusterEnabled || !isKeystoreEnabled()) {
-            return;
-        }
+    private void loadExistingKeysFromDisk() {
         try {
-            Optional<JwtSigningKeyEntity> newestOpt =
-                    keyRepository.findFirstByOrderByCreatedAtDesc();
-            if (newestOpt.isEmpty()) {
-                return;
-            }
-            JwtSigningKeyEntity newest = newestOpt.get();
-            JwtVerificationKey current = activeKey;
-            if (current != null && newest.getKeyId().equals(current.getKeyId())) {
-                return;
-            }
-            JwtVerificationKey adopted =
-                    new JwtVerificationKey(newest.getKeyId(), newest.getVerifyingKey());
-            verifyingKeyCache.put(newest.getKeyId(), adopted);
-            activeKey = adopted;
-            log.info(
-                    "Adopted newest JWT signing key {} from the shared DB as active",
-                    newest.getKeyId());
-        } catch (Exception e) {
-            log.warn("Could not reload active JWT key from the shared DB: {}", e.getMessage());
-        }
-    }
+            Path keyDirectory = Path.of(InstallationPathConfig.getPrivateKeyPath());
 
-    /** Load every signing key from the shared DB into the caches; most recent becomes active. */
-    private void loadKeysFromDb() {
-        List<JwtSigningKeyEntity> keys = keyRepository.findAllByOrderByCreatedAtDesc();
-        if (keys.isEmpty()) {
-            log.info("No JWT keys in the database, generating a new keypair");
+            if (!Files.exists(keyDirectory)) {
+                log.info("No existing keys found, generating new keypair");
+                generateAndStoreKeypair();
+                return;
+            }
+
+            List<Path> keyFiles;
+            try (var stream = Files.list(keyDirectory)) {
+                keyFiles =
+                        stream.filter(path -> path.toString().endsWith(KEY_SUFFIX))
+                                .sorted(
+                                        (a, b) ->
+                                                b.getFileName().compareTo(a.getFileName())) // Most
+                                // recent
+                                // first
+                                .toList();
+            }
+
+            if (keyFiles.isEmpty()) {
+                log.info("No existing keys found in directory, generating new keypair");
+                generateAndStoreKeypair();
+                return;
+            }
+
+            log.info("Loading {} existing JWT keys from disk", keyFiles.size());
+            int loadedCount = 0;
+
+            for (Path keyFile : keyFiles) {
+                try {
+                    String keyId = keyFile.getFileName().toString().replace(KEY_SUFFIX, "");
+
+                    // Load private key first
+                    PrivateKey privateKey = loadPrivateKey(keyId);
+
+                    // Try to load public key, or generate it from private key if missing
+                    // (migration)
+                    String encodedPublicKey;
+                    try {
+                        encodedPublicKey = loadPublicKey(keyId);
+                    } catch (IOException e) {
+                        // Public key file doesn't exist - generate it from private key (migration)
+                        log.info("Migrating legacy key: generating public key file for {}", keyId);
+                        KeyPair keyPair = reconstructKeyPair(privateKey);
+
+                        // Save the public key file
+                        Path publicKeyFile = keyDirectory.resolve(keyId + PUB_KEY_SUFFIX);
+                        encodedPublicKey = encodePublicKey(keyPair.getPublic());
+                        Files.writeString(publicKeyFile, encodedPublicKey);
+                        publicKeyFile.toFile().setReadable(true, true);
+                        publicKeyFile.toFile().setWritable(true, true);
+                        publicKeyFile.toFile().setExecutable(false, false);
+
+                        log.info("Successfully migrated key: {}", keyId);
+                    }
+
+                    // Create verification key and add to cache
+                    JwtVerificationKey verifyingKey =
+                            new JwtVerificationKey(keyId, encodedPublicKey);
+                    verifyingKeyCache.put(keyId, verifyingKey);
+                    loadedCount++;
+
+                    // Set the most recent key as active (first in sorted list)
+                    if (activeKey == null) {
+                        activeKey = verifyingKey;
+                        log.info("Set active JWT signing key: {}", keyId);
+                    } else {
+                        log.debug(
+                                "Loaded historical JWT key: {} (created: {})",
+                                keyId,
+                                verifyingKey.getCreatedAt());
+                    }
+                } catch (Exception e) {
+                    log.warn(
+                            "Failed to load key: {}, skipping. Error: {}",
+                            keyFile.getFileName(),
+                            e.getMessage());
+                }
+            }
+
+            if (loadedCount == 0) {
+                log.warn("No valid keys could be loaded from disk, generating new keypair");
+                generateAndStoreKeypair();
+            } else {
+                log.info(
+                        "Successfully loaded {} JWT keys, active key: {}",
+                        loadedCount,
+                        activeKey.getKeyId());
+            }
+
+        } catch (IOException e) {
+            log.error("Failed to load keys from disk, generating new keypair", e);
             generateAndStoreKeypair();
-            return;
         }
-        for (JwtSigningKeyEntity key : keys) {
-            verifyingKeyCache.put(
-                    key.getKeyId(), new JwtVerificationKey(key.getKeyId(), key.getVerifyingKey()));
-        }
-        activeKey = new JwtVerificationKey(keys.get(0).getKeyId(), keys.get(0).getVerifyingKey());
-        log.info("Loaded {} JWT key(s) from DB, active key: {}", keys.size(), activeKey.getKeyId());
     }
 
+    @Transactional
     private JwtVerificationKey generateAndStoreKeypair() {
+        JwtVerificationKey verifyingKey = null;
+
         try {
             KeyPair keyPair = generateRSAKeypair();
             String keyId = generateKeyId();
-            String verifyingKey = encodePublicKey(keyPair.getPublic());
-            String signingKey =
-                    Base64.getEncoder().encodeToString(keyPair.getPrivate().getEncoded());
 
-            // Converter encrypts the private key at rest with the shared credential-encryption key.
-            keyRepository.save(new JwtSigningKeyEntity(keyId, verifyingKey, signingKey));
-
-            keyPairCache.put(keyId, keyPair);
-            JwtVerificationKey verificationKey = new JwtVerificationKey(keyId, verifyingKey);
-            verifyingKeyCache.put(keyId, verificationKey);
-            activeKey = verificationKey;
+            storeKeyPair(keyId, keyPair);
+            verifyingKey = new JwtVerificationKey(keyId, encodePublicKey(keyPair.getPublic()));
+            verifyingKeyCache.put(keyId, verifyingKey);
+            activeKey = verifyingKey;
             log.info("Generated and stored new JWT keypair: {}", keyId);
-            return verificationKey;
-        } catch (RuntimeException e) {
+        } catch (IOException e) {
             log.error("Failed to generate and store keypair", e);
-            return null;
         }
+
+        return verifyingKey;
     }
 
     @Override
@@ -169,29 +207,25 @@ public class KeyPersistenceService implements KeyPersistenceServiceInterface {
 
     @Override
     public Optional<KeyPair> getKeyPair(String keyId) {
-        if (!isKeystoreEnabled() || keyId == null) {
+        if (!isKeystoreEnabled()) {
             return Optional.empty();
         }
-        KeyPair cached = keyPairCache.get(keyId);
-        if (cached != null) {
-            return Optional.of(cached);
-        }
-        Optional<JwtSigningKeyEntity> entityOpt = keyRepository.findById(keyId);
-        if (entityOpt.isEmpty()) {
-            log.warn("No signing key found in DB for keyId: {}", keyId);
-            return Optional.empty();
-        }
-        JwtSigningKeyEntity entity = entityOpt.get();
+
         try {
-            KeyPair keyPair =
-                    new KeyPair(
-                            decodePublicKey(entity.getVerifyingKey()),
-                            decodePrivateKey(entity.getSigningKey()));
-            keyPairCache.put(keyId, keyPair);
-            verifyingKeyCache.put(keyId, new JwtVerificationKey(keyId, entity.getVerifyingKey()));
-            return Optional.of(keyPair);
-        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
-            log.error("Failed to decode keypair for keyId: {}", keyId, e);
+            JwtVerificationKey verifyingKey =
+                    verifyingKeyCache.get(keyId, JwtVerificationKey.class);
+
+            if (verifyingKey == null) {
+                log.warn("No signing key found in database for keyId: {}", keyId);
+                return Optional.empty();
+            }
+
+            PrivateKey privateKey = loadPrivateKey(keyId);
+            PublicKey publicKey = decodePublicKey(verifyingKey.getVerifyingKey());
+
+            return Optional.of(new KeyPair(publicKey, privateKey));
+        } catch (Exception e) {
+            log.error("Failed to load keypair for keyId: {}", keyId, e);
             return Optional.empty();
         }
     }
@@ -207,69 +241,44 @@ public class KeyPersistenceService implements KeyPersistenceServiceInterface {
     }
 
     @Override
+    @CacheEvict(
+            value = {"verifyingKeys"},
+            key = "#keyId",
+            condition = "#root.target.isKeystoreEnabled()")
     public void removeKey(String keyId) {
-        keyRepository.deleteById(keyId);
         verifyingKeyCache.evict(keyId);
-        keyPairCache.remove(keyId);
     }
 
     @Override
     public List<JwtVerificationKey> getKeysEligibleForCleanup(LocalDateTime cutoffDate) {
-        return keyRepository.findByCreatedAtBefore(cutoffDate).stream()
-                .map(e -> new JwtVerificationKey(e.getKeyId(), e.getVerifyingKey()))
+        CaffeineCache caffeineCache = (CaffeineCache) verifyingKeyCache;
+        com.github.benmanes.caffeine.cache.Cache<Object, Object> nativeCache =
+                caffeineCache.getNativeCache();
+
+        log.debug(
+                "Cache size: {}, Checking {} keys for cleanup",
+                nativeCache.estimatedSize(),
+                nativeCache.asMap().size());
+
+        return nativeCache.asMap().values().stream()
+                .filter(value -> value instanceof JwtVerificationKey)
+                .map(value -> (JwtVerificationKey) value)
+                .filter(
+                        key -> {
+                            boolean eligible = key.getCreatedAt().isBefore(cutoffDate);
+                            log.debug(
+                                    "Key {} created at {}, eligible for cleanup: {}",
+                                    key.getKeyId(),
+                                    key.getCreatedAt(),
+                                    eligible);
+                            return eligible;
+                        })
                 .toList();
     }
 
-    /** Import any pre-existing on-disk keys into the DB once, so upgrades keep sessions valid. */
-    private void importLegacyDiskKeysIfPresent() {
-        if (keyRepository.count() > 0) {
-            return;
-        }
-        Path keyDirectory = Path.of(InstallationPathConfig.getPrivateKeyPath());
-        if (!Files.exists(keyDirectory)) {
-            return;
-        }
-        List<Path> keyFiles;
-        try (var stream = Files.list(keyDirectory)) {
-            keyFiles = stream.filter(p -> p.toString().endsWith(KEY_SUFFIX)).toList();
-        } catch (IOException e) {
-            log.warn("Could not list legacy key directory {}: {}", keyDirectory, e.getMessage());
-            return;
-        }
-        int imported = 0;
-        for (Path keyFile : keyFiles) {
-            String keyId = keyFile.getFileName().toString().replace(KEY_SUFFIX, "");
-            try {
-                PrivateKey privateKey = loadPrivateKey(keyId);
-                String verifyingKey = resolveLegacyPublicKey(keyId, privateKey);
-                String signingKey = Base64.getEncoder().encodeToString(privateKey.getEncoded());
-                keyRepository.save(new JwtSigningKeyEntity(keyId, verifyingKey, signingKey));
-                imported++;
-            } catch (Exception e) {
-                log.warn("Skipping legacy key {}: {}", keyId, e.getMessage());
-            }
-        }
-        if (imported > 0) {
-            log.info("Imported {} legacy JWT key(s) from disk into the shared DB", imported);
-        }
-    }
-
-    private String resolveLegacyPublicKey(String keyId, PrivateKey privateKey)
-            throws NoSuchAlgorithmException, InvalidKeySpecException {
-        try {
-            return loadPublicKey(keyId);
-        } catch (IOException e) {
-            // No .pub file: derive the public key from the RSA private key.
-            return encodePublicKey(reconstructKeyPair(privateKey).getPublic());
-        }
-    }
-
-    // UUID suffix so two nodes booting the same second don't collide on keyId.
     private String generateKeyId() {
         return "jwt-key-"
-                + LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd-HHmmss"))
-                + "-"
-                + UUID.randomUUID().toString().substring(0, 8);
+                + LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd-HHmmss"));
     }
 
     private KeyPair generateRSAKeypair() {
@@ -282,49 +291,125 @@ public class KeyPersistenceService implements KeyPersistenceServiceInterface {
         }
     }
 
+    private void ensurePrivateKeyDirectoryExists() throws IOException {
+        Path keyPath = Path.of(InstallationPathConfig.getPrivateKeyPath());
+
+        if (!Files.exists(keyPath)) {
+            Files.createDirectories(keyPath);
+        }
+    }
+
+    /**
+     * Store both private and public keys to disk.
+     *
+     * <p>Private key stored as: keyId.key
+     *
+     * <p>Public key stored as: keyId.pub
+     */
+    private void storeKeyPair(String keyId, KeyPair keyPair) throws IOException {
+        Path keyDirectory = Path.of(InstallationPathConfig.getPrivateKeyPath());
+
+        // Store private key
+        Path privateKeyFile = keyDirectory.resolve(keyId + KEY_SUFFIX);
+        String encodedPrivateKey =
+                Base64.getEncoder().encodeToString(keyPair.getPrivate().getEncoded());
+        Files.writeString(privateKeyFile, encodedPrivateKey);
+
+        // Set read/write to only the owner (security)
+        privateKeyFile.toFile().setReadable(true, true);
+        privateKeyFile.toFile().setWritable(true, true);
+        privateKeyFile.toFile().setExecutable(false, false);
+
+        // Store public key
+        Path publicKeyFile = keyDirectory.resolve(keyId + PUB_KEY_SUFFIX);
+        String encodedPublicKey =
+                Base64.getEncoder().encodeToString(keyPair.getPublic().getEncoded());
+        Files.writeString(publicKeyFile, encodedPublicKey);
+
+        // Public key can be more permissive but still restrict to owner
+        publicKeyFile.toFile().setReadable(true, true);
+        publicKeyFile.toFile().setWritable(true, true);
+        publicKeyFile.toFile().setExecutable(false, false);
+
+        log.debug(
+                "Stored keypair to disk: {} (private: {}, public: {})",
+                keyId,
+                privateKeyFile.getFileName(),
+                publicKeyFile.getFileName());
+    }
+
     private PrivateKey loadPrivateKey(String keyId)
             throws IOException, NoSuchAlgorithmException, InvalidKeySpecException {
         Path keyFile =
                 Path.of(InstallationPathConfig.getPrivateKeyPath()).resolve(keyId + KEY_SUFFIX);
+
         if (!Files.exists(keyFile)) {
             throw new IOException("Private key not found: " + keyFile);
         }
-        return decodePrivateKey(Files.readString(keyFile));
+
+        String encodedKey = Files.readString(keyFile);
+        byte[] keyBytes = Base64.getDecoder().decode(encodedKey);
+        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
+        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
+
+        return keyFactory.generatePrivate(keySpec);
     }
 
+    /**
+     * Load public key from disk.
+     *
+     * @param keyId the key identifier
+     * @return Base64-encoded public key string
+     * @throws IOException if the public key file is not found
+     */
     private String loadPublicKey(String keyId) throws IOException {
         Path publicKeyFile =
                 Path.of(InstallationPathConfig.getPrivateKeyPath()).resolve(keyId + PUB_KEY_SUFFIX);
+
         if (!Files.exists(publicKeyFile)) {
             throw new IOException("Public key not found: " + publicKeyFile);
         }
+
         return Files.readString(publicKeyFile).trim();
     }
 
+    /**
+     * Reconstruct a KeyPair from a PrivateKey.
+     *
+     * <p>For RSA keys, derives the public key from the private key.
+     *
+     * @param privateKey the RSA private key
+     * @return reconstructed KeyPair
+     * @throws NoSuchAlgorithmException if RSA algorithm is not available
+     * @throws InvalidKeySpecException if the key specification is invalid
+     */
     private KeyPair reconstructKeyPair(PrivateKey privateKey)
             throws NoSuchAlgorithmException, InvalidKeySpecException {
+        // For RSA, we can derive the public key from the private key
         KeyFactory keyFactory = KeyFactory.getInstance("RSA");
+
+        // Get the private key spec
         RSAPrivateCrtKey rsaPrivateKey = (RSAPrivateCrtKey) privateKey;
+
+        // Create public key spec from private key parameters
         RSAPublicKeySpec publicKeySpec =
                 new RSAPublicKeySpec(rsaPrivateKey.getModulus(), rsaPrivateKey.getPublicExponent());
-        return new KeyPair(keyFactory.generatePublic(publicKeySpec), privateKey);
+
+        // Generate public key
+        PublicKey publicKey = keyFactory.generatePublic(publicKeySpec);
+
+        return new KeyPair(publicKey, privateKey);
     }
 
     private String encodePublicKey(PublicKey publicKey) {
         return Base64.getEncoder().encodeToString(publicKey.getEncoded());
     }
 
-    @Override
     public PublicKey decodePublicKey(String encodedKey)
             throws NoSuchAlgorithmException, InvalidKeySpecException {
-        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(Base64.getDecoder().decode(encodedKey));
-        return KeyFactory.getInstance("RSA").generatePublic(keySpec);
-    }
-
-    private PrivateKey decodePrivateKey(String encodedKey)
-            throws NoSuchAlgorithmException, InvalidKeySpecException {
-        PKCS8EncodedKeySpec keySpec =
-                new PKCS8EncodedKeySpec(Base64.getDecoder().decode(encodedKey));
-        return KeyFactory.getInstance("RSA").generatePrivate(keySpec);
+        byte[] keyBytes = Base64.getDecoder().decode(encodedKey);
+        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
+        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
+        return keyFactory.generatePublic(keySpec);
     }
 }

app/proprietary/src/main/java/stirling/software/proprietary/service/AiDocumentHtmlRenderer.java

@@ -0,0 +1,135 @@
+package stirling.software.proprietary.service;
+
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.regex.Pattern;
+
+import org.springframework.core.io.ClassPathResource;
+import org.springframework.stereotype.Component;
+
+import com.hubspot.jinjava.Jinjava;
+import com.hubspot.jinjava.JinjavaConfig;
+
+import stirling.software.proprietary.model.api.ai.create.AiDocument;
+
+/** Renders an {@link AiDocument} to HTML using a Jinja template loaded from the classpath. */
+@Component
+public class AiDocumentHtmlRenderer {
+
+    private static final String TEMPLATE_PATH = "templates/ai/create/document.html.jinja2";
+
+    private static final Pattern SAFE_COLOR = Pattern.compile("^#[0-9a-fA-F]{6}$");
+
+    private final Jinjava jinjava;
+    private final String template;
+
+    public AiDocumentHtmlRenderer() {
+        JinjavaConfig config =
+                JinjavaConfig.newBuilder().withNestedInterpretationEnabled(false).build();
+        this.jinjava = new Jinjava(config);
+        this.template = loadTemplate();
+    }
+
+    public String render(AiDocument doc) {
+        return jinjava.render(template, buildContext(doc));
+    }
+
+    private static Map<String, Object> buildContext(AiDocument doc) {
+        Map<String, Object> context = new LinkedHashMap<>();
+        context.put("title", doc.getTitle());
+        context.put("subtitle", doc.getSubtitle());
+        context.put("reference_number", doc.getReferenceNumber());
+
+        AiDocument.Style style = doc.getStyle();
+        if (style != null) {
+            context.put("style_primary", safeColor(style.getPrimaryColor()));
+            context.put("style_background", safeColor(style.getBackgroundColor()));
+            context.put("style_body", safeColor(style.getBodyTextColor()));
+        }
+
+        List<Map<String, Object>> sections = new ArrayList<>();
+        if (doc.getSections() != null) {
+            for (AiDocument.Section section : doc.getSections()) {
+                if (section != null && section.getType() != null) {
+                    sections.add(buildSection(section));
+                }
+            }
+        }
+        context.put("sections", sections);
+        return context;
+    }
+
+    private static Map<String, Object> buildSection(AiDocument.Section section) {
+        Map<String, Object> node = new LinkedHashMap<>();
+        node.put("type", section.getType());
+        node.put("heading", section.getHeading());
+        switch (section.getType()) {
+            case "text" -> node.put("paragraphs", paragraphs(section.getBody()));
+            case "key_value" -> node.put("pairs", pairs(section.getPairs()));
+            case "line_items" -> {
+                node.put("columns", orEmpty(section.getColumns()));
+                node.put("rows", orEmptyRows(section.getRows()));
+                node.put("total_row", emptyToNull(section.getTotalRow()));
+            }
+            case "bullet_list" -> node.put("items", orEmpty(section.getItems()));
+            case "signature" -> node.put("signatories", orEmpty(section.getSignatories()));
+            default -> {}
+        }
+        return node;
+    }
+
+    private static List<String> paragraphs(String body) {
+        String text = body == null ? "" : body;
+        List<String> out = new ArrayList<>();
+        for (String paragraph : text.split("\n\n")) {
+            out.add(paragraph.replace("\n", " "));
+        }
+        return out;
+    }
+
+    private static List<Map<String, String>> pairs(List<List<String>> pairs) {
+        List<Map<String, String>> out = new ArrayList<>();
+        if (pairs != null) {
+            for (List<String> pair : pairs) {
+                Map<String, String> node = new LinkedHashMap<>();
+                node.put("label", pair.isEmpty() ? "" : pair.get(0));
+                node.put("value", pair.size() < 2 ? "" : pair.get(1));
+                out.add(node);
+            }
+        }
+        return out;
+    }
+
+    private static List<String> orEmpty(List<String> values) {
+        return values == null ? List.of() : values;
+    }
+
+    private static List<List<String>> orEmptyRows(List<List<String>> rows) {
+        return rows == null ? List.of() : rows;
+    }
+
+    private static List<String> emptyToNull(List<String> values) {
+        return values == null || values.isEmpty() ? null : values;
+    }
+
+    private static String safeColor(String value) {
+        if (value == null) {
+            return null;
+        }
+        String trimmed = value.trim();
+        return SAFE_COLOR.matcher(trimmed).matches() ? trimmed : null;
+    }
+
+    private static String loadTemplate() {
+        try {
+            return new ClassPathResource(TEMPLATE_PATH).getContentAsString(StandardCharsets.UTF_8);
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+}

app/proprietary/src/main/resources/templates/ai/create/document.html.jinja2

@@ -1,3 +1,4 @@
+{%- autoescape true -%}
 <!DOCTYPE html>
 <html lang="en">
 <head>
@@ -175,18 +176,18 @@
     color: var(--color-label);
   }
 </style>
-{%- if doc.style %}
+{%- if style_primary or style_background or style_body %}
 <style>
   :root {
-    {%- if doc.style.primary_color %}
-    --color-primary: {{ doc.style.primary_color }};
+    {%- if style_primary %}
+    --color-primary: {{ style_primary }};
     {%- endif %}
-    {%- if doc.style.background_color %}
-    --color-bg: {{ doc.style.background_color }};
+    {%- if style_background %}
+    --color-bg: {{ style_background }};
     {%- endif %}
-    {%- if doc.style.body_text_color %}
-    --color-body: {{ doc.style.body_text_color }};
-    --color-label: {{ doc.style.body_text_color }};
+    {%- if style_body %}
+    --color-body: {{ style_body }};
+    --color-label: {{ style_body }};
     {%- endif %}
   }
 </style>
@@ -195,16 +196,16 @@
 <body>
 
 <div class="doc-header">
-  <div class="doc-title">{{ doc.title }}</div>
-  {%- if doc.subtitle %}
-  <div class="doc-subtitle">{{ doc.subtitle }}</div>
+  <div class="doc-title">{{ title }}</div>
+  {%- if subtitle %}
+  <div class="doc-subtitle">{{ subtitle }}</div>
   {%- endif %}
-  {%- if doc.reference_number %}
-  <div class="doc-reference">{{ doc.reference_number }}</div>
+  {%- if reference_number %}
+  <div class="doc-reference">{{ reference_number }}</div>
   {%- endif %}
 </div>
 
-{%- for section in doc.sections %}
+{%- for section in sections %}
 
 {%- if section.type == "text" %}
 <section>
@@ -212,8 +213,8 @@
   <h2>{{ section.heading }}</h2>
   {%- endif %}
   <div class="text-body">
-    {%- for para in section.body.split('\n\n') %}
-    <p>{{ para | replace('\n', ' ') }}</p>
+    {%- for para in section.paragraphs %}
+    <p>{{ para }}</p>
     {%- endfor %}
   </div>
 </section>
@@ -225,10 +226,10 @@
   {%- endif %}
   <table class="kv-table">
     <tbody>
-      {%- for label, value in section.pairs %}
+      {%- for pair in section.pairs %}
       <tr>
-        <td class="kv-label">{{ label }}</td>
-        <td class="kv-value">{{ value }}</td>
+        <td class="kv-label">{{ pair.label }}</td>
+        <td class="kv-value">{{ pair.value }}</td>
       </tr>
       {%- endfor %}
     </tbody>
@@ -299,3 +300,4 @@
 
 </body>
 </html>
+{%- endautoescape %}

app/proprietary/src/test/java/stirling/software/proprietary/policy/controller/PolicyControllerTest.java

@@ -27,8 +27,6 @@ import org.springframework.http.ResponseEntity;
 import org.springframework.web.server.ResponseStatusException;
 import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
 
-import stirling.software.common.cluster.JobStore;
-import stirling.software.common.cluster.inprocess.InProcessJobStore;
 import stirling.software.common.model.ApplicationProperties;
 import stirling.software.common.model.job.JobResponse;
 import stirling.software.common.service.JobOwnershipService;
@@ -80,7 +78,6 @@ class PolicyControllerTest {
     @Mock private JobOwnershipService jobOwnershipService;
 
     private ApplicationProperties applicationProperties;
-    private final JobStore jobStore = new InProcessJobStore();
     private PolicyController controller;
 
     private final java.util.List<stirling.software.proprietary.policy.trigger.PolicyTrigger>
@@ -109,8 +106,7 @@ class PolicyControllerTest {
                         policyTriggers,
                         applicationProperties,
                         tempFileManager,
-                        jobOwnershipService,
-                        jobStore);
+                        jobOwnershipService);
     }
 
     private static stirling.software.proprietary.policy.trigger.PolicyTrigger trigger(
@@ -268,8 +264,6 @@ class PolicyControllerTest {
         @DisplayName("returns 404 when run is unknown")
         void notFound() {
             when(runRegistry.get("missing")).thenReturn(null);
-            when(jobOwnershipService.extractJobId("missing")).thenReturn("missing");
-            when(jobOwnershipService.createScopedJobKey("missing")).thenReturn("missing");
 
             ResponseEntity<PolicyRunView> response = controller.status("missing");
 

app/proprietary/src/test/java/stirling/software/proprietary/policy/engine/PolicyExecutorTest.java

@@ -166,8 +166,8 @@ class PolicyExecutorTest {
                                 new PipelineStep(
                                         createPdf,
                                         Map.of(
-                                                "htmlContent",
-                                                "<p>hi</p>",
+                                                "document",
+                                                "{\"title\":\"PO\",\"sections\":[]}",
                                                 "filename",
                                                 "purchase-order.pdf"))),
                         PolicyInputs.of(List.of()),

app/proprietary/src/test/java/stirling/software/proprietary/security/service/KeyPairCleanupServiceTest.java

@@ -1,102 +0,0 @@
-package stirling.software.proprietary.security.service;
-
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.Mockito.lenient;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
-import java.util.List;
-import java.util.Optional;
-
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.extension.ExtendWith;
-import org.mockito.Mock;
-import org.mockito.junit.jupiter.MockitoExtension;
-
-import stirling.software.common.cluster.DistributedLock;
-import stirling.software.common.cluster.DistributedLock.LockHandle;
-import stirling.software.common.model.ApplicationProperties;
-import stirling.software.proprietary.security.model.JwtVerificationKey;
-
-/** Cluster safety: pruning JWT keys is single-writer, gated on the shared cleanup lock. */
-@ExtendWith(MockitoExtension.class)
-class KeyPairCleanupServiceTest {
-
-    @Mock private KeyPersistenceService keyPersistenceService;
-    @Mock private ApplicationProperties applicationProperties;
-    @Mock private ApplicationProperties.Security security;
-    @Mock private ApplicationProperties.Security.Jwt jwtProperties;
-    @Mock private DistributedLock distributedLock;
-    @Mock private LockHandle lockHandle;
-
-    private KeyPairCleanupService cleanupService;
-
-    @BeforeEach
-    void setUp() {
-        lenient().when(applicationProperties.getSecurity()).thenReturn(security);
-        lenient().when(security.getJwt()).thenReturn(jwtProperties);
-        lenient().when(jwtProperties.isEnableKeyCleanup()).thenReturn(true);
-        lenient().when(keyPersistenceService.isKeystoreEnabled()).thenReturn(true);
-        cleanupService =
-                new KeyPairCleanupService(
-                        keyPersistenceService, applicationProperties, distributedLock);
-    }
-
-    @Test
-    void skipsPruningWhenAnotherNodeHoldsTheLock() {
-        when(distributedLock.tryAcquire(any(), any())).thenReturn(Optional.empty());
-
-        cleanupService.cleanup();
-
-        // No node-local pruning happened; the lock holder owns this cycle.
-        verify(keyPersistenceService, never()).getKeysEligibleForCleanup(any());
-        verify(keyPersistenceService, never()).refreshActiveKeyPair();
-    }
-
-    @Test
-    void prunesAndRotatesWhenLockAcquiredThenReleasesIt() {
-        when(distributedLock.tryAcquire(any(), any())).thenReturn(Optional.of(lockHandle));
-        when(keyPersistenceService.getKeysEligibleForCleanup(any()))
-                .thenReturn(List.of(new JwtVerificationKey("old-key", "cHVi")));
-
-        cleanupService.cleanup();
-
-        verify(keyPersistenceService).removeKey("old-key");
-        verify(keyPersistenceService).refreshActiveKeyPair();
-        verify(lockHandle).close();
-    }
-
-    @Test
-    void releasesTheLockEvenWhenNoKeysAreEligible() {
-        when(distributedLock.tryAcquire(any(), any())).thenReturn(Optional.of(lockHandle));
-        when(keyPersistenceService.getKeysEligibleForCleanup(any())).thenReturn(List.of());
-
-        cleanupService.cleanup();
-
-        verify(keyPersistenceService, never()).refreshActiveKeyPair();
-        verify(lockHandle).close();
-    }
-
-    @Test
-    void skipsPruningWhenTheLockBackendErrors() {
-        // A Valkey blip at boot must not fail startup: tryAcquire throwing degrades to skip.
-        when(distributedLock.tryAcquire(any(), any()))
-                .thenThrow(new RuntimeException("valkey unreachable"));
-
-        cleanupService.cleanup();
-
-        verify(keyPersistenceService, never()).getKeysEligibleForCleanup(any());
-        verify(keyPersistenceService, never()).refreshActiveKeyPair();
-    }
-
-    @Test
-    void doesNothingWhenCleanupDisabled() {
-        when(jwtProperties.isEnableKeyCleanup()).thenReturn(false);
-
-        cleanupService.cleanup();
-
-        verify(distributedLock, never()).tryAcquire(any(), any());
-    }
-}

app/proprietary/src/test/java/stirling/software/proprietary/security/service/KeyPersistenceServiceInterfaceTest.java

@@ -4,41 +4,45 @@ import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
-import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.lenient;
-import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.mockStatic;
 import static org.mockito.Mockito.when;
 
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.security.KeyPair;
 import java.security.KeyPairGenerator;
 import java.security.NoSuchAlgorithmException;
 import java.util.Base64;
-import java.util.List;
 import java.util.Optional;
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.api.io.TempDir;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.Mock;
+import org.mockito.MockedStatic;
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.springframework.cache.CacheManager;
 import org.springframework.cache.concurrent.ConcurrentMapCacheManager;
 
+import stirling.software.common.configuration.InstallationPathConfig;
 import stirling.software.common.model.ApplicationProperties;
-import stirling.software.proprietary.security.model.JwtSigningKeyEntity;
 import stirling.software.proprietary.security.model.JwtVerificationKey;
-import stirling.software.proprietary.security.repository.JwtSigningKeyRepository;
 
-/** DB-backed keystore: a key present only in the shared DB still resolves (the cross-node case). */
 @ExtendWith(MockitoExtension.class)
 class KeyPersistenceServiceInterfaceTest {
 
     @Mock private ApplicationProperties applicationProperties;
+
     @Mock private ApplicationProperties.Security security;
+
     @Mock private ApplicationProperties.Security.Jwt jwtConfig;
-    @Mock private JwtSigningKeyRepository keyRepository;
+
+    @TempDir Path tempDir;
 
     private KeyPersistenceService keyPersistenceService;
     private KeyPair testKeyPair;
@@ -54,128 +58,175 @@ class KeyPersistenceServiceInterfaceTest {
 
         lenient().when(applicationProperties.getSecurity()).thenReturn(security);
         lenient().when(security.getJwt()).thenReturn(jwtConfig);
-        lenient().when(jwtConfig.isEnableKeystore()).thenReturn(true);
-        lenient().when(keyRepository.save(any())).thenAnswer(inv -> inv.getArgument(0));
-        // clusterEnabled=true so the convergence-reload path is exercised.
-        keyPersistenceService =
-                new KeyPersistenceService(applicationProperties, cacheManager, keyRepository, true);
-    }
-
-    private JwtSigningKeyEntity entityFrom(String keyId) {
-        return new JwtSigningKeyEntity(
-                keyId,
-                Base64.getEncoder().encodeToString(testKeyPair.getPublic().getEncoded()),
-                Base64.getEncoder().encodeToString(testKeyPair.getPrivate().getEncoded()));
+        lenient().when(jwtConfig.isEnableKeystore()).thenReturn(true); // Default value
     }
 
     @ParameterizedTest
     @ValueSource(booleans = {true, false})
     void testKeystoreEnabled(boolean keystoreEnabled) {
         when(jwtConfig.isEnableKeystore()).thenReturn(keystoreEnabled);
-        assertEquals(keystoreEnabled, keyPersistenceService.isKeystoreEnabled());
+
+        try (MockedStatic<InstallationPathConfig> mockedStatic =
+                mockStatic(InstallationPathConfig.class)) {
+            mockedStatic
+                    .when(InstallationPathConfig::getPrivateKeyPath)
+                    .thenReturn(tempDir.toString());
+            keyPersistenceService = new KeyPersistenceService(applicationProperties, cacheManager);
+
+            assertEquals(keystoreEnabled, keyPersistenceService.isKeystoreEnabled());
+        }
     }
 
     @Test
-    void generatesAndPersistsAKeyWhenNoneIsActive() {
-        // getActiveKey with no active key mints one and persists it - no disk involved.
-        JwtVerificationKey active = keyPersistenceService.getActiveKey();
+    void testGetActiveKeypairWhenNoActiveKeyExists() {
+        try (MockedStatic<InstallationPathConfig> mockedStatic =
+                mockStatic(InstallationPathConfig.class)) {
+            mockedStatic
+                    .when(InstallationPathConfig::getPrivateKeyPath)
+                    .thenReturn(tempDir.toString());
+            keyPersistenceService = new KeyPersistenceService(applicationProperties, cacheManager);
+            keyPersistenceService.initializeKeystore();
 
-        assertNotNull(active);
-        assertNotNull(active.getKeyId());
-        assertNotNull(active.getVerifyingKey());
-        verify(keyRepository).save(any(JwtSigningKeyEntity.class));
+            JwtVerificationKey result = keyPersistenceService.getActiveKey();
+
+            assertNotNull(result);
+            assertNotNull(result.getKeyId());
+            assertNotNull(result.getVerifyingKey());
+        }
     }
 
     @Test
-    void loadsTheMostRecentExistingKeyAsActive() {
-        when(keyRepository.count()).thenReturn(1L);
-        when(keyRepository.findAllByOrderByCreatedAtDesc())
-                .thenReturn(List.of(entityFrom("jwt-key-2026-07-13-000000-abcd1234")));
+    void testGetActiveKeyPairWithExistingKey() throws Exception {
+        String keyId = "test-key-2024-01-01-120000";
+        String publicKeyBase64 =
+                Base64.getEncoder().encodeToString(testKeyPair.getPublic().getEncoded());
+        String privateKeyBase64 =
+                Base64.getEncoder().encodeToString(testKeyPair.getPrivate().getEncoded());
 
-        keyPersistenceService.initializeKeystore();
-        JwtVerificationKey active = keyPersistenceService.getActiveKey();
+        JwtVerificationKey existingKey = new JwtVerificationKey(keyId, publicKeyBase64);
 
-        assertEquals("jwt-key-2026-07-13-000000-abcd1234", active.getKeyId());
+        Path keyFile = tempDir.resolve(keyId + ".key");
+        Files.writeString(keyFile, privateKeyBase64);
+
+        try (MockedStatic<InstallationPathConfig> mockedStatic =
+                mockStatic(InstallationPathConfig.class)) {
+            mockedStatic
+                    .when(InstallationPathConfig::getPrivateKeyPath)
+                    .thenReturn(tempDir.toString());
+            keyPersistenceService = new KeyPersistenceService(applicationProperties, cacheManager);
+            keyPersistenceService.initializeKeystore();
+
+            JwtVerificationKey result = keyPersistenceService.getActiveKey();
+
+            assertNotNull(result);
+            assertNotNull(result.getKeyId());
+        }
     }
 
     @Test
-    void getKeyPairResolvesAKeyPresentOnlyInTheSharedDb() {
-        // Never initialised locally: the key lives only in the DB, as if another node minted it.
-        String keyId = "jwt-key-from-another-node";
-        when(keyRepository.findById(keyId)).thenReturn(Optional.of(entityFrom(keyId)));
+    void testGetKeyPair() throws Exception {
+        String keyId = "test-key-123";
+        String publicKeyBase64 =
+                Base64.getEncoder().encodeToString(testKeyPair.getPublic().getEncoded());
+        String privateKeyBase64 =
+                Base64.getEncoder().encodeToString(testKeyPair.getPrivate().getEncoded());
 
-        Optional<KeyPair> result = keyPersistenceService.getKeyPair(keyId);
+        JwtVerificationKey signingKey = new JwtVerificationKey(keyId, publicKeyBase64);
 
-        assertTrue(result.isPresent());
-        assertNotNull(result.get().getPublic());
-        assertNotNull(result.get().getPrivate());
+        Path keyFile = tempDir.resolve(keyId + ".key");
+        Files.writeString(keyFile, privateKeyBase64);
+
+        try (MockedStatic<InstallationPathConfig> mockedStatic =
+                mockStatic(InstallationPathConfig.class)) {
+            mockedStatic
+                    .when(InstallationPathConfig::getPrivateKeyPath)
+                    .thenReturn(tempDir.toString());
+            keyPersistenceService = new KeyPersistenceService(applicationProperties, cacheManager);
+
+            keyPersistenceService
+                    .getClass()
+                    .getDeclaredField("verifyingKeyCache")
+                    .setAccessible(true);
+            var cache = cacheManager.getCache("verifyingKeys");
+            cache.put(keyId, signingKey);
+
+            Optional<KeyPair> result = keyPersistenceService.getKeyPair(keyId);
+
+            assertTrue(result.isPresent());
+            assertNotNull(result.get().getPublic());
+            assertNotNull(result.get().getPrivate());
+        }
     }
 
     @Test
-    void getKeyPairIsEmptyWhenTheKeyIsUnknown() {
-        when(keyRepository.findById("nope")).thenReturn(Optional.empty());
-        assertFalse(keyPersistenceService.getKeyPair("nope").isPresent());
+    void testGetKeyPairNotFound() {
+        String keyId = "non-existent-key";
+
+        try (MockedStatic<InstallationPathConfig> mockedStatic =
+                mockStatic(InstallationPathConfig.class)) {
+            mockedStatic
+                    .when(InstallationPathConfig::getPrivateKeyPath)
+                    .thenReturn(tempDir.toString());
+            keyPersistenceService = new KeyPersistenceService(applicationProperties, cacheManager);
+
+            Optional<KeyPair> result = keyPersistenceService.getKeyPair(keyId);
+
+            assertFalse(result.isPresent());
+        }
     }
 
     @Test
-    void getKeyPairIsEmptyWhenKeystoreDisabled() {
+    void testGetKeyPairWhenKeystoreDisabled() {
         when(jwtConfig.isEnableKeystore()).thenReturn(false);
-        assertFalse(keyPersistenceService.getKeyPair("any-key").isPresent());
+
+        try (MockedStatic<InstallationPathConfig> mockedStatic =
+                mockStatic(InstallationPathConfig.class)) {
+            mockedStatic
+                    .when(InstallationPathConfig::getPrivateKeyPath)
+                    .thenReturn(tempDir.toString());
+            keyPersistenceService = new KeyPersistenceService(applicationProperties, cacheManager);
+
+            Optional<KeyPair> result = keyPersistenceService.getKeyPair("any-key");
+
+            assertFalse(result.isPresent());
+        }
     }
 
     @Test
-    void eligibleForCleanupIsSourcedFromTheDb() {
-        when(keyRepository.findByCreatedAtBefore(any())).thenReturn(List.of(entityFrom("old-key")));
+    void testInitializeKeystoreCreatesDirectory() throws IOException {
+        try (MockedStatic<InstallationPathConfig> mockedStatic =
+                mockStatic(InstallationPathConfig.class)) {
+            mockedStatic
+                    .when(InstallationPathConfig::getPrivateKeyPath)
+                    .thenReturn(tempDir.toString());
+            keyPersistenceService = new KeyPersistenceService(applicationProperties, cacheManager);
+            keyPersistenceService.initializeKeystore();
 
-        List<JwtVerificationKey> stale =
-                keyPersistenceService.getKeysEligibleForCleanup(java.time.LocalDateTime.now());
-
-        assertEquals(1, stale.size());
-        assertEquals("old-key", stale.get(0).getKeyId());
+            assertTrue(Files.exists(tempDir));
+            assertTrue(Files.isDirectory(tempDir));
+        }
     }
 
     @Test
-    void reloadAdoptsTheNewestKeyAPeerMinted() {
-        // Boot with our own key active, then a peer mints a newer one in the shared DB.
-        when(keyRepository.count()).thenReturn(1L);
-        when(keyRepository.findAllByOrderByCreatedAtDesc())
-                .thenReturn(List.of(entityFrom("jwt-key-local-old")));
-        keyPersistenceService.initializeKeystore();
-        assertEquals("jwt-key-local-old", keyPersistenceService.getActiveKey().getKeyId());
+    void testLoadExistingKeypairWithMissingPrivateKeyFile() throws Exception {
+        String keyId = "test-key-missing-file";
+        String publicKeyBase64 =
+                Base64.getEncoder().encodeToString(testKeyPair.getPublic().getEncoded());
 
-        when(keyRepository.findFirstByOrderByCreatedAtDesc())
-                .thenReturn(Optional.of(entityFrom("jwt-key-peer-new")));
+        JwtVerificationKey existingKey = new JwtVerificationKey(keyId, publicKeyBase64);
 
-        keyPersistenceService.reloadActiveKeyFromDb();
+        try (MockedStatic<InstallationPathConfig> mockedStatic =
+                mockStatic(InstallationPathConfig.class)) {
+            mockedStatic
+                    .when(InstallationPathConfig::getPrivateKeyPath)
+                    .thenReturn(tempDir.toString());
+            keyPersistenceService = new KeyPersistenceService(applicationProperties, cacheManager);
+            keyPersistenceService.initializeKeystore();
 
-        // Converged: this node now signs with the peer's newer key.
-        assertEquals("jwt-key-peer-new", keyPersistenceService.getActiveKey().getKeyId());
-    }
-
-    @Test
-    void reloadDoesNothingOffCluster() {
-        KeyPersistenceService singleNode =
-                new KeyPersistenceService(
-                        applicationProperties, cacheManager, keyRepository, false);
-
-        singleNode.reloadActiveKeyFromDb();
-
-        // Off-cluster the DB is never consulted for convergence.
-        verify(keyRepository, org.mockito.Mockito.never()).findFirstByOrderByCreatedAtDesc();
-    }
-
-    @Test
-    void reloadIsANoOpWhenAlreadyHoldingTheNewestKey() {
-        when(keyRepository.count()).thenReturn(1L);
-        when(keyRepository.findAllByOrderByCreatedAtDesc())
-                .thenReturn(List.of(entityFrom("jwt-key-current")));
-        keyPersistenceService.initializeKeystore();
-
-        when(keyRepository.findFirstByOrderByCreatedAtDesc())
-                .thenReturn(Optional.of(entityFrom("jwt-key-current")));
-
-        keyPersistenceService.reloadActiveKeyFromDb();
-
-        assertEquals("jwt-key-current", keyPersistenceService.getActiveKey().getKeyId());
+            JwtVerificationKey result = keyPersistenceService.getActiveKey();
+            assertNotNull(result);
+            assertNotNull(result.getKeyId());
+            assertNotNull(result.getVerifyingKey());
+        }
     }
 }

app/proprietary/src/test/java/stirling/software/proprietary/service/AiDocumentHtmlRendererTest.java

@@ -0,0 +1,139 @@
+package stirling.software.proprietary.service;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.util.List;
+
+import org.junit.jupiter.api.Test;
+
+import stirling.software.proprietary.model.api.ai.create.AiDocument;
+
+class AiDocumentHtmlRendererTest {
+
+    private final AiDocumentHtmlRenderer renderer = new AiDocumentHtmlRenderer();
+
+    private static AiDocument.Section section(String type) {
+        AiDocument.Section s = new AiDocument.Section();
+        s.setType(type);
+        return s;
+    }
+
+    private static AiDocument document(String title, List<AiDocument.Section> sections) {
+        AiDocument doc = new AiDocument();
+        doc.setTitle(title);
+        doc.setSections(sections);
+        return doc;
+    }
+
+    @Test
+    void rendersAllSectionTypes() {
+        AiDocument.Section text = section("text");
+        text.setBody("Some prose text.");
+        AiDocument.Section kv = section("key_value");
+        kv.setPairs(List.of(List.of("Key", "Value")));
+        AiDocument.Section items = section("line_items");
+        items.setColumns(List.of("A", "B"));
+        items.setRows(List.of(List.of("1", "2")));
+        AiDocument.Section bullets = section("bullet_list");
+        bullets.setItems(List.of("item one"));
+        AiDocument.Section sign = section("signature");
+        sign.setSignatories(List.of("Alice"));
+
+        String html = renderer.render(document("All", List.of(text, kv, items, bullets, sign)));
+
+        assertTrue(html.contains("<!DOCTYPE html>"));
+        assertTrue(html.contains("Some prose text."));
+        assertTrue(html.contains("Key") && html.contains("Value"));
+        assertTrue(html.contains("<th>"));
+        assertTrue(html.contains("item one"));
+        assertTrue(html.contains("Alice"));
+    }
+
+    @Test
+    void rendersMarkupCharactersAsText() {
+        AiDocument.Section text = section("text");
+        text.setBody("a <b>x</b> & y");
+
+        String html = renderer.render(document("Doc", List.of(text)));
+
+        assertFalse(html.contains("<b>"));
+        assertTrue(html.contains("&lt;b&gt;"));
+    }
+
+    @Test
+    void totalRowRenderedWhenPresent() {
+        AiDocument.Section items = section("line_items");
+        items.setColumns(List.of("Item", "Total"));
+        items.setRows(List.of(List.of("Widget", "$10")));
+        items.setTotalRow(List.of("Total", "$10"));
+
+        assertTrue(
+                renderer.render(document("Table", List.of(items)))
+                        .contains("<tr class=\"total-row\">"));
+    }
+
+    @Test
+    void totalRowAbsentWhenNotProvided() {
+        AiDocument.Section items = section("line_items");
+        items.setColumns(List.of("Item"));
+        items.setRows(List.of(List.of("Widget")));
+
+        assertFalse(
+                renderer.render(document("Table", List.of(items)))
+                        .contains("<tr class=\"total-row\">"));
+    }
+
+    @Test
+    void rendersSubtitleAndReference() {
+        AiDocument doc = document("My Doc", List.of());
+        doc.setSubtitle("Subtitle Here");
+        doc.setReferenceNumber("REF-42");
+
+        String html = renderer.render(doc);
+
+        assertTrue(html.contains("Subtitle Here"));
+        assertTrue(html.contains("REF-42"));
+    }
+
+    @Test
+    void appliesHexColourOverride() {
+        AiDocument doc = document("Styled", List.of());
+        AiDocument.Style style = new AiDocument.Style();
+        style.setPrimaryColor("#ff00ff");
+        style.setBackgroundColor("#111111");
+        doc.setStyle(style);
+
+        String html = renderer.render(doc);
+
+        assertTrue(html.contains("--color-primary: #ff00ff"));
+        assertTrue(html.contains("--color-bg: #111111"));
+    }
+
+    @Test
+    void ignoresColourWithDisallowedCharacters() {
+        AiDocument doc = document("Styled", List.of());
+        AiDocument.Style style = new AiDocument.Style();
+        style.setPrimaryColor("rgb(255, 0, 0)");
+        doc.setStyle(style);
+
+        String html = renderer.render(doc);
+
+        assertFalse(html.contains("rgb("));
+        assertTrue(html.contains("<!DOCTYPE html>"));
+    }
+
+    @Test
+    void ignoresNonHexColour() {
+        AiDocument doc = document("Styled", List.of());
+        AiDocument.Style style = new AiDocument.Style();
+        style.setPrimaryColor("magenta");
+        style.setBackgroundColor("#fff");
+        doc.setStyle(style);
+
+        String html = renderer.render(doc);
+
+        assertFalse(html.contains("--color-primary: magenta"));
+        assertFalse(html.contains("--color-bg: #fff;"));
+    }
+}

build.gradle

@@ -36,6 +36,8 @@ ext {
     okhttpBomVersion = "5.3.2"
     gsonVersion = "2.14.0"
     guavaVersion = "33.6.0-jre"
+    jinjavaVersion = "2.8.3"
+    jackson2Version = "2.21.2"
     bucket4jVersion = "8.19.0"
     archunitVersion = "1.4.2"
     batikVersion = "1.19"
@@ -222,6 +224,13 @@ subprojects {
         resolutionStrategy.force "org.apache.commons:commons-lang3:${commonsLang3}"
         // CVE-2024-47554: commons-io DoS prevention
         resolutionStrategy.force "commons-io:commons-io:${commonsIoVersion}"
+        // Jackson 2 is transitive-only here (jinjava, opensaml, jjwt request older versions);
+        // pin the family to a current release and keep modules aligned.
+        resolutionStrategy.force "com.fasterxml.jackson.core:jackson-core:${jackson2Version}"
+        resolutionStrategy.force "com.fasterxml.jackson.core:jackson-databind:${jackson2Version}"
+        resolutionStrategy.force "com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:${jackson2Version}"
+        resolutionStrategy.force "com.fasterxml.jackson.datatype:jackson-datatype-jdk8:${jackson2Version}"
+        resolutionStrategy.force "com.fasterxml.jackson.datatype:jackson-datatype-jsr310:${jackson2Version}"
         // Keep BouncyCastle modules aligned to avoid runtime linkage errors
         resolutionStrategy.force "org.bouncycastle:bcprov-jdk18on:${bouncycastleVersion}"
         resolutionStrategy.force "org.bouncycastle:bcpkix-jdk18on:${bouncycastleVersion}"

engine/pyproject.toml

@@ -5,7 +5,6 @@ description = "AI Document Engine"
 requires-python = ">=3.13"
 dependencies = [
     "fastapi>=0.116.0",
-    "jinja2>=3.1.0",
     "pgvector>=0.3.6",
     "psycopg[binary,pool]>=3.2",
     "pydantic>=2.0.0",

engine/src/stirling/agents/pdf_create/agent.py

@@ -10,7 +10,7 @@ Flow:
   4. SectionWriterAgents (smart_model) run in parallel via asyncio.gather.
      Each returns a WrittenSections with fully populated DocumentSection objects.
   5. The assembler collects sections in plan order → GeneratedDocument.
-  6. Jinja renders the document to HTML. The LLM never writes HTML.
+  6. The assembled document is emitted as structured fields. The LLM never writes HTML.
 
 The planner is split into two calls (meta then sections) so each LLM output schema
 stays small enough for grammar compilation on all model tiers including Haiku.
@@ -22,9 +22,7 @@ import asyncio
 import logging
 import re
 from dataclasses import dataclass
-from pathlib import Path
 
-from jinja2 import Environment, FileSystemLoader
 from pydantic_ai import Agent
 from pydantic_ai.output import NativeOutput
 
@@ -51,8 +49,6 @@ from stirling.services import AppRuntime
 
 logger = logging.getLogger(__name__)
 
-_TEMPLATES_DIR = Path(__file__).parent / "templates"
-
 # ── Token budget ──────────────────────────────────────────────────────────────────────────────────
 
 # Conservative per-section token estimates mapped from planner-assigned depth.
@@ -166,10 +162,13 @@ Analyse the user's request and produce a DocumentMeta with:
   document, if the user provides one. Leave empty if the user provides no such context.
 
 - style_primary_color: accent and heading colour. Set ONLY when the user explicitly names a
-  colour or colour scheme (e.g. "make it red", "use navy blue"). Use CSS named colours
-  (e.g. "magenta", "navy", "crimson") or hex values. Leave null if no colour is stated.
-- style_background_color: page background colour. Set only if explicitly requested.
-- style_body_text_color: body text colour. Set only if explicitly requested.
+  colour or colour scheme (e.g. "make it red", "use navy blue"). Express it as a 6-digit hex
+  code in #RRGGBB format (map any named colour to its hex value yourself, e.g. "navy" →
+  "#000080"). No other format is accepted. Leave null if no colour is stated.
+- style_background_color: page background colour, same #RRGGBB format. Set only if explicitly
+  requested.
+- style_body_text_color: body text colour, same #RRGGBB format. Set only if explicitly
+  requested.
 
 - cannot_do_reason: set this ONLY when the request is not asking to create a document at all
   (e.g. a question, a greeting, an edit request to an existing document). Never set it
@@ -299,15 +298,6 @@ def _build_writer_prompt(plan: DocumentPlan, chunk: _Chunk) -> str:
 # ── Helpers ───────────────────────────────────────────────────────────────────────────────────────
 
 
-def _build_jinja_env() -> Environment:
-    return Environment(
-        loader=FileSystemLoader(str(_TEMPLATES_DIR)),
-        autoescape=True,
-        trim_blocks=True,
-        lstrip_blocks=True,
-    )
-
-
 def _safe_filename(title: str) -> str:
     slug = re.sub(r"[^\w\s-]", "", title.lower())
     slug = re.sub(r"[\s_-]+", "-", slug).strip("-")
@@ -320,7 +310,6 @@ def _safe_filename(title: str) -> str:
 class PdfCreateAgent:
     def __init__(self, runtime: AppRuntime) -> None:
         self.runtime = runtime
-        self._jinja_env = _build_jinja_env()
 
         self._meta_planner: Agent[None, DocumentMeta] = Agent(
             model=runtime.smart_model,
@@ -401,14 +390,12 @@ class PdfCreateAgent:
             sections=all_sections,
         )
 
-        # ── Phase 6: render ────────────────────────────────────────────────────
-        logger.info("[pdf-create] phase 6/6: rendering HTML")
-        html = self._render(doc)
+        # ── Phase 6: emit ──────────────────────────────────────────────────────
         filename = _safe_filename(plan.title)
         logger.info(
-            "[pdf-create] done — filename=%r html_bytes=%d",
+            "[pdf-create] done — filename=%r sections=%d",
             filename,
-            len(html),
+            len(all_sections),
         )
 
         return EditPlanResponse(
@@ -417,7 +404,7 @@ class PdfCreateAgent:
                 ToolOperationStep(
                     tool=AgentToolId.CREATE_PDF_FROM_HTML_AGENT,
                     parameters=CreatePdfFromHtmlAgentParams(
-                        html_content=html,
+                        document=doc.model_dump_json(),
                         filename=filename,
                     ),
                 )
@@ -437,7 +424,3 @@ class PdfCreateAgent:
             len(result.output.sections),
         )
         return result.output
-
-    def _render(self, doc: GeneratedDocument) -> str:
-        template = self._jinja_env.get_template("document.html.jinja2")
-        return template.render(doc=doc)

engine/src/stirling/contracts/pdf_create.py

@@ -1,14 +1,14 @@
 """Contracts for the PDF Create Agent.
 
 The agent accepts a natural-language prompt and returns a single
-CREATE_PDF_FROM_HTML_AGENT plan step carrying the rendered HTML.
+CREATE_PDF_FROM_HTML_AGENT plan step carrying the assembled document.
 
 Pipeline:
   1. PlannerAgent (smart_model) → DocumentPlan: structured skeleton, no body text.
   2. Python chunks the plan by token budget.
   3. SectionWriterAgents (smart_model, parallel) → WrittenSections per chunk.
   4. Assembler collects sections in plan order → GeneratedDocument.
-  5. Jinja renders GeneratedDocument → HTML. The LLM never writes HTML.
+  5. The document is emitted as structured fields. The LLM never writes HTML.
 """
 
 from __future__ import annotations
@@ -81,14 +81,12 @@ type DocumentSection = Annotated[
 ]
 
 
-# Named colour or hex only — anything else is dropped so a colour can't inject CSS into the
-# <style> block (which would let WeasyPrint fetch an attacker-controlled url() → SSRF).
-_SAFE_COLOR_RE = re.compile(r"^#[0-9a-fA-F]{3,8}$|^[a-zA-Z]{1,30}$")
+# Colours must be a 6-digit hex code (#RRGGBB); anything else is dropped to None.
+_SAFE_COLOR_RE = re.compile(r"^#[0-9a-fA-F]{6}$")
 
 
 class DocumentStyle(ApiModel):
-    """Document colours, inferred by the meta planner and rendered into the engine's Jinja
-    template (never sent to Java). Unsafe colours are dropped to ``None``."""
+    """Document colours, inferred by the meta planner. Non-hex values are dropped to ``None``."""
 
     primary_color: str | None = Field(default=None)
     background_color: str | None = Field(default=None)
@@ -103,7 +101,7 @@ class DocumentStyle(ApiModel):
 
 
 class GeneratedDocument(ApiModel):
-    """The full document model passed to Jinja for HTML rendering."""
+    """The full document model emitted for rendering."""
 
     title: str
     subtitle: str | None = None

engine/src/stirling/models/agent_tool_models.py

@@ -29,7 +29,7 @@ class PdfCommentAgentParams(ApiModel):
 
 
 class CreatePdfFromHtmlAgentParams(ApiModel):
-    html_content: str
+    document: str
     filename: str = Field(pattern=r"^.+\.pdf$")
 
 

engine/src/stirling/services/tracking.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import logging
 from collections import OrderedDict
 from collections.abc import Mapping
 from contextvars import ContextVar
@@ -29,6 +30,8 @@ from posthog.client import Client as PostHogClient
 from stirling.config import AppSettings
 from stirling.models import UserId
 
+logger = logging.getLogger(__name__)
+
 # Per-request user ID, set by middleware from the X-User-Id header.
 # When not set, PostHog generates a random ID and marks the event as personless.
 current_user_id: ContextVar[UserId | None] = ContextVar("current_user_id", default=None)
@@ -127,17 +130,27 @@ class PostHogSpanProcessor(SpanProcessor):
         pass
 
     def on_end(self, span: ReadableSpan) -> None:
-        attrs = dict(span.attributes or {})
-        if attrs.get(GEN_AI_OPERATION_NAME) != GenAiOperationNameValues.CHAT.value:
-            return
+        # OpenTelemetry calls this synchronously from ``Span.end()`` with no
+        # exception handling of its own (see SynchronousMultiSpanProcessor),
+        # and pydantic-ai ends the model span inside a ``with`` block. So any
+        # exception raised here propagates straight into the model request and
+        # breaks the agent run. A telemetry defect must never take down the
+        # assistant or silently halt emission for every subsequent span, so we
+        # translate and capture defensively and swallow (but log) any failure.
+        try:
+            attrs = dict(span.attributes or {})
+            if attrs.get(GEN_AI_OPERATION_NAME) != GenAiOperationNameValues.CHAT.value:
+                return
 
-        properties = self._build_generation_properties(span, attrs)
-        self._maybe_emit_trace_event(span, attrs, properties)
-        self._client.capture(
-            distinct_id=current_user_id.get(),
-            event="$ai_generation",
-            properties=properties,
-        )
+            properties = self._build_generation_properties(span, attrs)
+            self._maybe_emit_trace_event(span, attrs, properties)
+            self._client.capture(
+                distinct_id=current_user_id.get(),
+                event="$ai_generation",
+                properties=properties,
+            )
+        except Exception:
+            logger.exception("Failed to emit $ai_generation for span; dropping this event")
 
     def _build_generation_properties(self, span: ReadableSpan, attrs: Mapping[str, Any]) -> dict[str, object]:
         """Build the $ai_generation event properties from span data."""
@@ -231,6 +244,19 @@ class PostHogSpanProcessor(SpanProcessor):
         return True
 
 
+def _log_delivery_error(error: Exception, _batch: Any) -> None:
+    """PostHog client ``on_error`` callback.
+
+    The PostHog SDK delivers events from a background consumer thread and its
+    ``capture`` method is decorated ``@no_throw``, so without this hook a
+    failing (or dead) delivery pipeline is completely invisible: events simply
+    stop arriving with nothing in our logs. Surfacing the error lets a halt be
+    noticed in minutes instead of sitting silently at zero. Kept lightweight so
+    a transient upload error can't spam at capture volume.
+    """
+    logger.warning("PostHog event delivery failed: %s", error)
+
+
 def setup_posthog_tracking(settings: AppSettings) -> TracerProvider | None:
     """Configure OpenTelemetry with a PostHog span processor for LLM analytics.
 
@@ -238,11 +264,17 @@ def setup_posthog_tracking(settings: AppSettings) -> TracerProvider | None:
     or None when tracking is disabled.
     """
     if not settings.posthog_enabled or not settings.posthog_api_key:
+        logger.info("PostHog LLM tracking is disabled (posthog_enabled=%s)", settings.posthog_enabled)
         return None
 
-    client = PostHogClient(project_api_key=settings.posthog_api_key, host=settings.posthog_host)
+    client = PostHogClient(
+        project_api_key=settings.posthog_api_key,
+        host=settings.posthog_host,
+        on_error=_log_delivery_error,
+    )
     processor = PostHogSpanProcessor(client)
 
     provider = TracerProvider()
     provider.add_span_processor(processor)
+    logger.info("PostHog LLM tracking enabled (host=%s)", settings.posthog_host)
     return provider

engine/tests/agents/test_pdf_create.py

@@ -2,7 +2,7 @@
 
 Coverage:
 1. Section model validation (each section type round-trips correctly)
-2. Jinja rendering (_render produces valid HTML for each section type)
+2. orchestrate() emits the assembled document as structured JSON
 3. _safe_filename produces clean slugs
 4. _make_chunks groups sections correctly by token budget
 5. orchestrate() produces the correct EditPlanResponse via planner + writer mocks
@@ -12,6 +12,8 @@ Coverage:
 
 from __future__ import annotations
 
+import json
+
 import pytest
 from conftest import build_app_settings
 from pydantic_ai.models.test import TestModel
@@ -64,37 +66,6 @@ def agent(runtime: AppRuntime) -> PdfCreateAgent:
 # ── Helpers ───────────────────────────────────────────────────────────────────────────────────────
 
 
-def _invoice_doc() -> GeneratedDocument:
-    return GeneratedDocument(
-        title="Invoice",
-        subtitle="Acme Corp",
-        reference_number="Invoice #INV-001",
-        sections=[
-            KeyValueSection(
-                heading="Details",
-                pairs=[("Date", "2026-05-06"), ("Due", "2026-06-06"), ("Currency", "USD")],
-            ),
-            LineItemsSection(
-                heading="Line Items",
-                columns=["Description", "Qty", "Unit Price", "Total"],
-                rows=[
-                    ["Consulting services", "10", "$500.00", "$5,000.00"],
-                    ["Expenses", "1", "$200.00", "$200.00"],
-                ],
-                total_row=["Total", "", "", "$5,200.00"],
-            ),
-            TextSection(
-                heading="Payment Terms",
-                body="Payment is due within 30 days.\n\nPlease reference the invoice number.",
-            ),
-            SignatureSection(
-                heading="Authorised By",
-                signatories=["Jane Smith, CEO", "Bob Jones, CFO"],
-            ),
-        ],
-    )
-
-
 def _simple_meta() -> DocumentMeta:
     return DocumentMeta(
         title="Invoice",
@@ -199,82 +170,6 @@ def test_generated_document_optional_fields() -> None:
     assert doc.reference_number is None
 
 
-# ── Jinja rendering ───────────────────────────────────────────────────────────────────────────────
-
-
-def test_render_produces_html(agent: PdfCreateAgent) -> None:
-    doc = _invoice_doc()
-    html = agent._render(doc)
-    assert "<!DOCTYPE html>" in html
-    assert "Invoice" in html
-    assert "INV-001" in html
-
-
-def test_render_includes_all_section_types(agent: PdfCreateAgent) -> None:
-    doc = GeneratedDocument(
-        title="All Sections",
-        sections=[
-            TextSection(body="Some prose text."),
-            KeyValueSection(pairs=[("Key", "Value")]),
-            LineItemsSection(columns=["A", "B"], rows=[["1", "2"]]),
-            BulletListSection(items=["item one"]),
-            SignatureSection(signatories=["Alice"]),
-        ],
-    )
-    html = agent._render(doc)
-    assert "Some prose text." in html
-    assert "Key" in html and "Value" in html
-    assert "<th>" in html
-    assert "item one" in html
-    assert "Alice" in html
-
-
-def test_render_escapes_html_in_content(agent: PdfCreateAgent) -> None:
-    doc = GeneratedDocument(
-        title="XSS Test",
-        sections=[TextSection(body="<script>alert('xss')</script>")],
-    )
-    html = agent._render(doc)
-    assert "<script>" not in html
-    assert "&lt;script&gt;" in html
-
-
-def test_render_total_row_present(agent: PdfCreateAgent) -> None:
-    doc = GeneratedDocument(
-        title="Table",
-        sections=[
-            LineItemsSection(
-                columns=["Item", "Total"],
-                rows=[["Widget", "$10"]],
-                total_row=["Total", "$10"],
-            )
-        ],
-    )
-    html = agent._render(doc)
-    assert "total-row" in html
-
-
-def test_render_no_total_row_skips_tfoot(agent: PdfCreateAgent) -> None:
-    doc = GeneratedDocument(
-        title="Table",
-        sections=[LineItemsSection(columns=["Item"], rows=[["Widget"]])],
-    )
-    html = agent._render(doc)
-    assert "<tfoot>" not in html
-
-
-def test_render_subtitle_and_reference(agent: PdfCreateAgent) -> None:
-    doc = GeneratedDocument(
-        title="My Doc",
-        subtitle="Subtitle Here",
-        reference_number="REF-42",
-        sections=[TextSection(body="Content.")],
-    )
-    html = agent._render(doc)
-    assert "Subtitle Here" in html
-    assert "REF-42" in html
-
-
 # ── _safe_filename ────────────────────────────────────────────────────────────────────────────────
 
 
@@ -396,8 +291,9 @@ async def test_orchestrate_returns_plan_step(agent: PdfCreateAgent) -> None:
     assert step.tool == AgentToolId.CREATE_PDF_FROM_HTML_AGENT
     assert isinstance(step.parameters, CreatePdfFromHtmlAgentParams)
     assert step.parameters.filename.endswith(".pdf")
-    assert "<!DOCTYPE html>" in step.parameters.html_content
-    assert "Invoice" in step.parameters.html_content
+    parsed = json.loads(step.parameters.document)
+    assert parsed["title"] == "Invoice"
+    assert parsed["sections"]
 
 
 @pytest.mark.anyio
@@ -468,9 +364,9 @@ async def test_orchestrate_assembles_multiple_chunks(agent: PdfCreateAgent) -> N
         result = await agent.orchestrate(_orchestrator_request("Create a multi-chunk doc"))
 
     assert isinstance(result, EditPlanResponse)
-    html = result.steps[0].parameters.html_content  # type: ignore[union-attr]
-    assert "Introduction text." in html
-    assert "Details" in html
+    document = result.steps[0].parameters.document  # type: ignore[union-attr]
+    assert "Introduction text." in document
+    assert "Details" in document
 
 
 # ── Style inference ───────────────────────────────────────────────────────────────────────────────
@@ -482,7 +378,7 @@ async def test_orchestrate_applies_planner_inferred_style(agent: PdfCreateAgent)
     meta = DocumentMeta(
         title="Styled Doc",
         tone_brief="Professional.",
-        style_primary_color="magenta",
+        style_primary_color="#ff00ff",
     )
     sections = _simple_sections()
     written = _written_sections()
@@ -499,50 +395,35 @@ async def test_orchestrate_applies_planner_inferred_style(agent: PdfCreateAgent)
         result = await agent.orchestrate(_orchestrator_request("Make an invoice, magenta styling"))
 
     assert isinstance(result, EditPlanResponse)
-    html = result.steps[0].parameters.html_content  # type: ignore[union-attr]
-    assert "magenta" in html
+    document = result.steps[0].parameters.document  # type: ignore[union-attr]
+    assert json.loads(document)["style"]["primaryColor"] == "#ff00ff"
 
 
-def test_render_applies_style(agent: PdfCreateAgent) -> None:
-    """DocumentStyle fields are injected as CSS custom properties in the rendered HTML."""
-    doc = GeneratedDocument(
-        title="Styled",
-        sections=[TextSection(body="Content.")],
-        style=DocumentStyle(primary_color="magenta", background_color="#111111"),
-    )
-    html = agent._render(doc)
-    assert "--color-primary: magenta" in html
-    assert "--color-bg: #111111" in html
-
-
-def test_document_style_drops_unsafe_colors() -> None:
-    """Unsafe colours (not named/hex) are dropped, closing the <style> url() injection."""
-    safe = DocumentStyle(primary_color="navy", background_color="#1e3a5f", body_text_color="#fff")
+def test_document_style_keeps_only_six_digit_hex() -> None:
+    """Only #RRGGBB hex is kept; named colours and other formats drop to None."""
+    safe = DocumentStyle(primary_color="#1e3a5f", background_color="#ffffff", body_text_color="#1A1A1A")
     assert (safe.primary_color, safe.background_color, safe.body_text_color) == (
-        "navy",
         "#1e3a5f",
-        "#fff",
+        "#ffffff",
+        "#1A1A1A",
     )
 
-    unsafe = DocumentStyle(
-        primary_color="red; background: url(http://evil.test/steal)",
-        background_color="expression(alert(1))",
-        body_text_color="navy; }",
-    )
-    assert unsafe.primary_color is None
-    assert unsafe.background_color is None
-    assert unsafe.body_text_color is None
+    assert DocumentStyle(primary_color="navy").primary_color is None
+    assert DocumentStyle(primary_color="#fff").primary_color is None
+    assert DocumentStyle(primary_color="#1e3a5f00").primary_color is None
+    assert DocumentStyle(primary_color="rgb(255, 0, 0)").primary_color is None
+    assert DocumentStyle(background_color="teal darken-2").background_color is None
     # A trailing newline must not slip a value through (fullmatch, not $-before-newline).
-    assert DocumentStyle(primary_color="navy\n").primary_color is None
+    assert DocumentStyle(primary_color="#1e3a5f\n").primary_color is None
 
 
 @pytest.mark.anyio
-async def test_orchestrate_drops_unsafe_planner_color(agent: PdfCreateAgent) -> None:
-    """An unsafe colour inferred by the meta planner never reaches the rendered HTML."""
+async def test_orchestrate_drops_non_hex_planner_colour(agent: PdfCreateAgent) -> None:
+    """A non-hex colour inferred by the meta planner is dropped before the document is emitted."""
     meta = DocumentMeta(
         title="Doc",
         tone_brief="Professional.",
-        style_primary_color="blue; background: url(http://evil.test/)",
+        style_primary_color="rgb(0, 0, 255)",
     )
     sections = _simple_sections()
     written = _written_sections()
@@ -559,6 +440,6 @@ async def test_orchestrate_drops_unsafe_planner_color(agent: PdfCreateAgent) ->
         result = await agent.orchestrate(_orchestrator_request("make it blue"))
 
     assert isinstance(result, EditPlanResponse)
-    html = result.steps[0].parameters.html_content  # type: ignore[union-attr]
-    assert "evil.test" not in html
-    assert "url(" not in html
+    document = result.steps[0].parameters.document  # type: ignore[union-attr]
+    assert "rgb(" not in document
+    assert json.loads(document)["style"]["primaryColor"] is None

engine/tests/test_tracking.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any
+
+from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import (
+    GEN_AI_INPUT_MESSAGES,
+    GEN_AI_OPERATION_NAME,
+    GEN_AI_SYSTEM,
+    GenAiOperationNameValues,
+)
+
+from stirling.config import AppSettings
+from stirling.services.tracking import PostHogSpanProcessor, current_user_id, setup_posthog_tracking
+
+
+@dataclass
+class _FakeContext:
+    trace_id: int
+    span_id: int
+
+
+@dataclass
+class _FakeParent:
+    span_id: int
+
+
+class _FakeSpan:
+    """Minimal stand-in for the OpenTelemetry ReadableSpan fields the processor reads."""
+
+    def __init__(self, attributes: dict[str, Any], trace_id: int = 0x1234, span_id: int = 0x9) -> None:
+        self.attributes = attributes
+        self.context = _FakeContext(trace_id=trace_id, span_id=span_id)
+        self.parent = _FakeParent(span_id=0x1)
+        self.start_time = 1_000_000_000
+        self.end_time = 2_000_000_000
+
+
+class _RecordingClient:
+    def __init__(self) -> None:
+        self.events: list[tuple[str, dict[str, Any]]] = []
+
+    def capture(self, *, distinct_id: str | None, event: str, properties: dict[str, Any]) -> None:
+        self.events.append((event, properties))
+
+
+class _RaisingClient:
+    def capture(self, **_kwargs: Any) -> None:
+        raise RuntimeError("boom")
+
+
+def _processor(client: object) -> PostHogSpanProcessor:
+    # The processor only calls ``capture`` on its client; the test doubles model
+    # that surface without being real PostHog clients (which start threads).
+    return PostHogSpanProcessor(client)  # type: ignore[arg-type]
+
+
+def _on_end(processor: PostHogSpanProcessor, span: object) -> None:
+    processor.on_end(span)  # type: ignore[arg-type]
+
+
+def _chat_attrs() -> dict[str, Any]:
+    return {
+        GEN_AI_OPERATION_NAME: GenAiOperationNameValues.CHAT.value,
+        GEN_AI_SYSTEM: "anthropic",
+        GEN_AI_INPUT_MESSAGES: json.dumps([{"role": "user", "parts": [{"type": "text", "content": "hello"}]}]),
+    }
+
+
+def test_chat_span_emits_generation_and_trace() -> None:
+    client = _RecordingClient()
+
+    _on_end(_processor(client), _FakeSpan(_chat_attrs()))
+
+    emitted = [event for event, _ in client.events]
+    assert "$ai_generation" in emitted
+    assert "$ai_trace" in emitted
+
+
+def test_trace_event_deduplicated_per_trace() -> None:
+    client = _RecordingClient()
+    processor = _processor(client)
+
+    _on_end(processor, _FakeSpan(_chat_attrs(), trace_id=0x42, span_id=0x1))
+    _on_end(processor, _FakeSpan(_chat_attrs(), trace_id=0x42, span_id=0x2))
+
+    trace_events = [event for event, _ in client.events if event == "$ai_trace"]
+    generation_events = [event for event, _ in client.events if event == "$ai_generation"]
+    assert len(trace_events) == 1
+    assert len(generation_events) == 2
+
+
+def test_non_chat_span_is_ignored() -> None:
+    client = _RecordingClient()
+
+    _on_end(_processor(client), _FakeSpan({GEN_AI_OPERATION_NAME: "embeddings"}))
+
+    assert client.events == []
+
+
+def test_on_end_never_raises_when_delivery_fails() -> None:
+    # A telemetry failure must not propagate into Span.end() and break the
+    # model request, nor wedge emission for every subsequent span.
+    _on_end(_processor(_RaisingClient()), _FakeSpan(_chat_attrs()))  # must not raise
+
+
+def test_on_end_never_raises_on_malformed_span() -> None:
+    processor = _processor(_RecordingClient())
+
+    broken = _FakeSpan(_chat_attrs())
+    broken.context = None  # type: ignore[assignment]
+
+    _on_end(processor, broken)  # must not raise
+
+
+def test_distinct_id_pulled_from_context_var() -> None:
+    client = _RecordingClient()
+    processor = _processor(client)
+
+    token = current_user_id.set("user-123")  # type: ignore[arg-type]
+    try:
+        _on_end(processor, _FakeSpan(_chat_attrs()))
+    finally:
+        current_user_id.reset(token)
+
+    assert client.events, "expected at least one captured event"
+
+
+def test_setup_returns_none_when_disabled(app_settings: AppSettings) -> None:
+    # conftest builds settings with posthog disabled; setup must no-op.
+    assert setup_posthog_tracking(app_settings) is None

engine/uv.lock

@@ -604,7 +604,6 @@ version = "0.1.0"
 source = { editable = "." }
 dependencies = [
     { name = "fastapi" },
-    { name = "jinja2" },
     { name = "opentelemetry-sdk" },
     { name = "pgvector" },
     { name = "posthog" },
@@ -631,7 +630,6 @@ dev = [
 [package.metadata]
 requires-dist = [
     { name = "fastapi", specifier = ">=0.116.0" },
-    { name = "jinja2", specifier = ">=3.1.0" },
     { name = "opentelemetry-sdk", specifier = ">=1.39.0" },
     { name = "pgvector", specifier = ">=0.3.6" },
     { name = "posthog", specifier = ">=3.0.0" },

frontend/editor/src/core/hooks/tools/shared/toolAutomation.test.ts

@@ -13,7 +13,6 @@ import {
 import {
   deserializeToolStep,
   getExecutableTools,
-  serializeStepFromEndpoint,
   serializeToolStep,
   stepRequiresUpload,
   type WorkingToolStep,
@@ -199,41 +198,6 @@ describe("serialize/deserialize round-trip", () => {
   });
 });
 
-describe("serializeStepFromEndpoint", () => {
-  test("maps a wizard step's UI params to the backend contract, filling defaults", () => {
-    // The shape the policy setup wizard holds: an endpoint plus UI-shaped params
-    // (redact's `wordsToRedact`), with several fields left to their defaults.
-    const api = serializeStepFromEndpoint(
-      "/api/v1/security/auto-redact",
-      { mode: "automatic", useRegex: true, wordsToRedact: ["ssn", "card"] },
-      dynamicRegistry,
-    );
-
-    expect(api.operation).toBe("/api/v1/security/auto-redact");
-    // wordsToRedact -> listOfText (the field the backend actually reads), and the
-    // frontend-only `mode` is dropped.
-    expect(api.parameters).toMatchObject({ listOfText: "ssn\ncard" });
-    expect(api.parameters).not.toHaveProperty("wordsToRedact");
-    expect(api.parameters).not.toHaveProperty("mode");
-    // Fields the wizard never set still get their defaults so the body is complete.
-    expect(api.parameters).toHaveProperty("wholeWordSearch");
-    expect(api.parameters).toHaveProperty("customPadding");
-  });
-
-  test("passes an unmapped endpoint's params through unchanged", () => {
-    expect(
-      serializeStepFromEndpoint(
-        "/api/v1/unknown/thing",
-        { keep: true },
-        dynamicRegistry,
-      ),
-    ).toEqual({
-      operation: "/api/v1/unknown/thing",
-      parameters: { keep: true },
-    });
-  });
-});
-
 describe("stepRequiresUpload", () => {
   const step = (params: Record<string, unknown>): WorkingToolStep => ({
     toolId: "compress" as ToolId,

frontend/editor/src/core/hooks/tools/shared/toolAutomation.ts

@@ -197,31 +197,6 @@ export function serializeToolStep(
   return { operation, parameters };
 }
 
-/**
- * Serialize a step held as an endpoint path plus frontend-shaped params - the form the policy setup
- * wizard keeps, where params match the tool's UI shape (e.g. redact's `wordsToRedact`) rather than
- * the backend contract - into the backend step contract, mapping params through the tool's
- * `toApiParams` (merged over its defaults, so fields the wizard never set still get their defaults).
- * The endpoint maps to a tool by path, so this works for dynamic-endpoint tools whose config
- * endpoint is a function. Endpoints that map to no known tool pass through unchanged.
- */
-export function serializeStepFromEndpoint(
-  operation: string,
-  params: ErasedToolParams,
-  registry: Partial<ToolRegistry>,
-): ToolApiStep {
-  const match = findToolByEndpoint({ operation, parameters: params }, registry);
-  const config = match?.[1].operationConfig;
-  if (!config) return { operation, parameters: params };
-  const merged = { ...(config.defaultParameters ?? {}), ...params };
-  return {
-    operation: resolveEndpoint(config, merged) ?? operation,
-    parameters: config.toApiParams
-      ? (config.toApiParams(merged) as Record<string, unknown>)
-      : {},
-  };
-}
-
 /**
  * Find the registry tool for a stored step's endpoint: exact match for static endpoints, else
  * membership in a dynamic tool's declared `endpoints` set (replaying its function can't recover a

frontend/editor/src/core/hooks/tools/shared/toolOperationDescriptor.test.ts

@@ -0,0 +1,37 @@
+import { describe, expect, test } from "vitest";
+import { describeToolOperation } from "@app/hooks/tools/shared/toolOperationDescriptor";
+
+interface Params {
+  a: number;
+}
+
+// A minimal config that type-checks against the flatten endpoint's model.
+const CONFIG = {
+  endpoint: "/api/v1/misc/flatten" as const,
+  defaultParameters: { a: 1 } satisfies Params,
+  toApiParams: (p: Params) => ({ renderDpi: p.a }),
+  fromApiParams: (api: { renderDpi?: number }) => ({ a: api.renderDpi ?? 0 }),
+};
+
+describe("describeToolOperation", () => {
+  test("wraps the config's mappers and endpoint into a descriptor", () => {
+    const d = describeToolOperation("/api/v1/misc/flatten", CONFIG);
+    expect(d.endpoint).toBe("/api/v1/misc/flatten");
+    expect(d.toApi({ a: 200 })).toEqual({ renderDpi: 200 });
+  });
+
+  test("fromApi merges the mapped values over the defaults", () => {
+    const d = describeToolOperation("/api/v1/misc/flatten", CONFIG);
+    expect(d.fromApi({ renderDpi: 72 })).toEqual({ a: 72 });
+  });
+
+  test("throws when the config lacks a mapper", () => {
+    expect(() =>
+      describeToolOperation("/api/v1/misc/flatten", {
+        endpoint: "/api/v1/misc/flatten" as const,
+        defaultParameters: { a: 1 },
+        toApiParams: (p: Params) => ({ renderDpi: p.a }),
+      }),
+    ).toThrow(/mappers/);
+  });
+});

frontend/editor/src/core/hooks/tools/shared/toolOperationDescriptor.ts

@@ -0,0 +1,59 @@
+/**
+ * Typed wrapper over a tool's `toApiParams`/`fromApiParams` mappers, binding one endpoint to safe
+ * frontend<->backend parameter conversion.
+ */
+
+import type { ToolApiParams, ToolEndpoint } from "@app/types/toolApiTypes";
+
+export interface ToolOperationDescriptor<E extends ToolEndpoint, TParams> {
+  readonly endpoint: E;
+  readonly defaultParameters: TParams;
+  toApi(params: TParams): ToolApiParams[E];
+  /** Backend model -> full frontend params (defaults merged under the mapped values). */
+  fromApi(api: ToolApiParams[E]): TParams;
+}
+
+/**
+ * Structural subset of a tool's config. `CE` is the config's declared endpoint type, inferred from
+ * the `endpoint` field: the literal for static tools, or the whole `ToolEndpoint` union for
+ * dynamic-endpoint tools (whose endpoint is a function typed against the union).
+ */
+export interface BidirectionalToolConfig<TParams, CE extends ToolEndpoint> {
+  endpoint: CE | null | ((params: TParams) => CE | null);
+  defaultParameters?: TParams;
+  toApiParams?(params: TParams): ToolApiParams[CE];
+  fromApiParams?(api: ToolApiParams[CE]): Partial<TParams>;
+}
+
+/**
+ * Pin a config to `endpoint` (passed explicitly, since dynamic-endpoint tools declare `endpoint` as
+ * a function). `E extends CE` rejects pairing a static tool's config with the wrong endpoint, while
+ * allowing a dynamic tool whose `CE` is the full union. Throws when mappers or defaults are missing.
+ */
+export function describeToolOperation<
+  E extends CE,
+  CE extends ToolEndpoint,
+  TParams,
+>(
+  endpoint: E,
+  config: BidirectionalToolConfig<TParams, CE>,
+): ToolOperationDescriptor<E, TParams> {
+  const { toApiParams, fromApiParams, defaultParameters } = config;
+  if (!toApiParams || !fromApiParams || defaultParameters === undefined) {
+    throw new Error(
+      `describeToolOperation: "${endpoint}" is missing mappers or defaults`,
+    );
+  }
+  return {
+    endpoint,
+    defaultParameters,
+    // A dynamic tool's mapper is typed against the union; narrow to this endpoint (sound - the
+    // runtime mapper produces this endpoint's model).
+    toApi: (params) => toApiParams(params) as ToolApiParams[E],
+    fromApi: (api) =>
+      ({
+        ...defaultParameters,
+        ...fromApiParams(api as ToolApiParams[CE]),
+      }) as TParams,
+  };
+}

frontend/editor/src/portal/api/policies.ts

@@ -13,6 +13,8 @@ import type { TFunction } from "i18next";
 import { apiClient } from "@portal/api/http";
 import { fromWirePolicy, toWirePolicy } from "@app/policies/codec";
 import { runsToActivity, runsToStats } from "@app/policies/runs";
+import { policyStep, type PolicyToolStep } from "@app/policies/operations";
+import type { ToolEndpoint } from "@app/types/toolApiTypes";
 import type {
   PolicyDecodedState,
   PolicyRunView,
@@ -66,7 +68,7 @@ export interface PolicyConfigDef {
   rules: string[];
   scopeLabel: string;
   fields: PolicyField[];
-  defaultOperations: WirePipelineStep[];
+  defaultOperations: PolicyToolStep[];
 }
 
 export interface PolicyState {
@@ -128,20 +130,11 @@ export interface CatalogueEntry {
 }
 
 /* ──────────────────────────────────────────────────────────────────────── */
-/*  Tool → endpoint registry                                                  */
+/*  Endpoint display labels                                                   */
 /* ──────────────────────────────────────────────────────────────────────── */
 
-export const TOOL_ENDPOINTS: Record<string, string> = {
-  redact: "/api/v1/security/auto-redact",
-  sanitize: "/api/v1/security/sanitize-pdf",
-  watermark: "/api/v1/security/add-watermark",
-  ocr: "/api/v1/misc/ocr-pdf",
-  flatten: "/api/v1/misc/flatten",
-  compress: "/api/v1/misc/compress-pdf",
-};
-
-/** Values are i18n keys — render with t(). */
-export const ENDPOINT_LABELS: Record<string, string> = {
+/** i18n keys keyed by {@link ToolEndpoint}; labels stored steps in the detail view. */
+export const ENDPOINT_LABELS: Partial<Record<ToolEndpoint, string>> = {
   "/api/v1/security/auto-redact": "portal.policies.endpoints.autoRedact",
   "/api/v1/security/sanitize-pdf": "portal.policies.endpoints.sanitizePdf",
   "/api/v1/security/add-watermark": "portal.policies.endpoints.addWatermark",
@@ -154,7 +147,8 @@ export function humanizeEndpoint(
   path: string,
   t: (key: string) => string,
 ): string {
-  if (ENDPOINT_LABELS[path]) return t(ENDPOINT_LABELS[path]);
+  const label = ENDPOINT_LABELS[path as ToolEndpoint];
+  if (label) return t(label);
   const last = path.split("/").filter(Boolean).pop() ?? path;
   return last
     .replace(/-/g, " ")
@@ -230,10 +224,7 @@ export const POLICY_CONFIG: Record<string, PolicyConfigDef> = {
       "portal.policies.config.ingestion.rules.3",
     ],
     scopeLabel: "portal.policies.config.scopeAll",
-    defaultOperations: [
-      { operation: TOOL_ENDPOINTS.ocr, parameters: {} },
-      { operation: TOOL_ENDPOINTS.flatten, parameters: {} },
-    ],
+    defaultOperations: [policyStep("ocr"), policyStep("flatten")],
     fields: [
       {
         label: "portal.policies.config.ingestion.fields.minConfidence",
@@ -260,32 +251,16 @@ export const POLICY_CONFIG: Record<string, PolicyConfigDef> = {
     ],
     scopeLabel: "portal.policies.config.scopeAll",
     defaultOperations: [
-      {
-        operation: TOOL_ENDPOINTS.redact,
-        parameters: {
-          mode: "automatic",
-          useRegex: true,
-          convertPDFToImage: true,
-          wordsToRedact: DEFAULT_PII_PATTERNS,
-        },
-      },
-      {
-        operation: TOOL_ENDPOINTS.sanitize,
-        parameters: {
-          removeJavaScript: true,
-          removeEmbeddedFiles: false,
-          removeMetadata: false,
-          removeLinks: false,
-          removeFonts: false,
-        },
-      },
-      {
-        operation: TOOL_ENDPOINTS.watermark,
-        // convertPDFToImage bakes the watermark in so it can't be stripped
-        parameters: {
-          convertPDFToImage: true,
-        },
-      },
+      // Flatten to image so redactions can't be lifted off.
+      policyStep("redact", {
+        useRegex: true,
+        convertPDFToImage: true,
+        wordsToRedact: DEFAULT_PII_PATTERNS,
+      }),
+      // JavaScript removal only; the tool enables removeEmbeddedFiles by default, so turn it off.
+      policyStep("sanitize", { removeEmbeddedFiles: false }),
+      // Bake in via image so it can't be stripped.
+      policyStep("watermark", { convertPDFToImage: true }),
     ],
     fields: [],
   },
@@ -297,10 +272,7 @@ export const POLICY_CONFIG: Record<string, PolicyConfigDef> = {
       "portal.policies.config.compliance.rules.2",
     ],
     scopeLabel: "portal.policies.config.scopeAll",
-    defaultOperations: [
-      { operation: TOOL_ENDPOINTS.sanitize, parameters: {} },
-      { operation: TOOL_ENDPOINTS.flatten, parameters: {} },
-    ],
+    defaultOperations: [policyStep("sanitize"), policyStep("flatten")],
     fields: [
       {
         label: "portal.policies.config.compliance.fields.frameworks",
@@ -343,7 +315,7 @@ export const POLICY_CONFIG: Record<string, PolicyConfigDef> = {
       "portal.policies.config.routing.rules.2",
     ],
     scopeLabel: "portal.policies.config.scopeAll",
-    defaultOperations: [{ operation: TOOL_ENDPOINTS.compress, parameters: {} }],
+    defaultOperations: [policyStep("compress")],
     fields: [
       {
         label: "portal.policies.config.routing.fields.destination",
@@ -374,7 +346,7 @@ export const POLICY_CONFIG: Record<string, PolicyConfigDef> = {
       "portal.policies.config.retention.rules.2",
     ],
     scopeLabel: "portal.policies.config.scopeAll",
-    defaultOperations: [{ operation: TOOL_ENDPOINTS.compress, parameters: {} }],
+    defaultOperations: [policyStep("compress")],
     fields: [
       {
         label: "portal.policies.config.retention.fields.keepFor",

frontend/editor/src/portal/components/pipelines/PipelineStepSettings.test.tsx

@@ -0,0 +1,52 @@
+import { describe, expect, it, vi } from "vitest";
+import { render, screen } from "@testing-library/react";
+import { MantineProvider } from "@mantine/core";
+import { Tooltip } from "@app/components/shared/Tooltip";
+import type { ToolRegistry } from "@app/data/toolsTaxonomy";
+import type { WorkingToolStep } from "@app/hooks/tools/shared/toolAutomation";
+import { PipelineStepSettings } from "@portal/components/pipelines/PipelineStepSettings";
+
+vi.mock("react-i18next", () => ({
+  useTranslation: () => ({
+    t: (key: string, fallback?: string) => fallback ?? key,
+  }),
+}));
+
+// A stand-in tool-settings UI that uses the shared editor Tooltip. The Tooltip
+// pulls in the Preferences + Sidebar contexts, which the portal does not mount
+// app-wide — so this reproduces the "usePreferences must be used within a
+// PreferencesProvider" crash unless PipelineStepSettings supplies them.
+function TooltipSettings() {
+  return (
+    <Tooltip content="help">
+      <button type="button">field</button>
+    </Tooltip>
+  );
+}
+
+const step = {
+  support: "editable",
+  toolId: "compress",
+  params: {},
+} as unknown as WorkingToolStep;
+
+const registry = {
+  compress: { automationSettings: TooltipSettings },
+} as unknown as Partial<ToolRegistry>;
+
+describe("PipelineStepSettings", () => {
+  it("renders reused editor tool settings (which use the shared Tooltip) without app-wide Preferences/Sidebar providers", () => {
+    expect(() =>
+      render(
+        <MantineProvider>
+          <PipelineStepSettings
+            step={step}
+            registry={registry}
+            onChange={() => {}}
+          />
+        </MantineProvider>,
+      ),
+    ).not.toThrow();
+    expect(screen.getByText("field")).toBeInTheDocument();
+  });
+});

frontend/editor/src/portal/components/pipelines/PipelineStepSettings.tsx

@@ -1,6 +1,8 @@
 import { Suspense } from "react";
 import { useTranslation } from "react-i18next";
 import { Banner } from "@app/ui";
+import { PreferencesProvider } from "@app/contexts/PreferencesContext";
+import { SidebarProvider } from "@app/contexts/SidebarContext";
 import { type ToolRegistry } from "@app/data/toolsTaxonomy";
 import { type ErasedToolParams } from "@app/hooks/tools/shared/toolOperationTypes";
 import { type WorkingToolStep } from "@app/hooks/tools/shared/toolAutomation";
@@ -46,14 +48,18 @@ export function PipelineStepSettings({
   }
 
   return (
-    <Suspense fallback={null}>
-      <Settings
-        parameters={step.params}
-        onParameterChange={(key, value) =>
-          onChange({ ...step.params, [key]: value })
-        }
-        disabled={false}
-      />
-    </Suspense>
+    <PreferencesProvider>
+      <SidebarProvider>
+        <Suspense fallback={null}>
+          <Settings
+            parameters={step.params}
+            onParameterChange={(key, value) =>
+              onChange({ ...step.params, [key]: value })
+            }
+            disabled={false}
+          />
+        </Suspense>
+      </SidebarProvider>
+    </PreferencesProvider>
   );
 }

frontend/editor/src/portal/components/policies/PolicySetupWizard.test.tsx

@@ -0,0 +1,134 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import {
+  fireEvent,
+  render as baseRender,
+  screen,
+  waitFor,
+} from "@testing-library/react";
+import { MantineProvider } from "@mantine/core";
+import { PolicySetupWizard } from "@portal/components/policies/PolicySetupWizard";
+import {
+  POLICY_CATEGORIES,
+  POLICY_CONFIG,
+  type CatalogueEntry,
+  type DecoratedPolicy,
+  type PolicySetupResult,
+  type PipelineStep,
+} from "@portal/api/policies";
+
+const render = (ui: Parameters<typeof baseRender>[0]) =>
+  baseRender(ui, { wrapper: MantineProvider });
+
+// Deterministic i18n: return the fallback when given, else the key. initReactI18next is stubbed
+// because the import graph pulls core/i18n.ts, which registers it as a plugin.
+vi.mock("react-i18next", () => ({
+  useTranslation: () => ({
+    // Second arg is a string fallback in some call sites and an interpolation object in others;
+    // only treat a string as the fallback.
+    t: (key: string, fallback?: unknown) =>
+      typeof fallback === "string" ? fallback : key,
+    i18n: { changeLanguage: vi.fn() },
+  }),
+  initReactI18next: { type: "3rdParty", init: vi.fn() },
+}));
+
+const fetchSources = vi.fn();
+vi.mock("@portal/api/sources", () => ({
+  fetchSources: () => fetchSources(),
+}));
+
+const CONTINUE = "portal.policies.wizard.actions.continue";
+const SAVE_CHANGES = "portal.policies.wizard.actions.saveChanges";
+const ENABLE = "portal.policies.wizard.actions.enablePolicy";
+
+const security = POLICY_CATEGORIES.find((c) => c.id === "security")!;
+const securityConfig = POLICY_CONFIG.security;
+
+function editEntry(steps: PipelineStep[]): CatalogueEntry {
+  const policy: DecoratedPolicy = {
+    category: security,
+    config: securityConfig,
+    state: {
+      configured: true,
+      status: "active",
+      sources: ["editor"],
+      scopeTypes: [],
+      reviewerEmail: "",
+      fieldValues: {},
+      runOn: "upload",
+      outputMode: "new_version",
+      outputName: "",
+      outputNamePosition: "suffix",
+      maxRetries: 0,
+      retryDelayMinutes: 0,
+      backendId: "pol-1",
+      isDefault: true,
+    },
+    steps,
+    stats: { enforced: 0, dataProcessed: "-", activeFor: "-" },
+    activity: [],
+  };
+  return { category: security, config: securityConfig, policy };
+}
+
+/** Advance the wizard from the workflow tab to the settings tab and submit. */
+async function submitWizard(saveLabel: string) {
+  fireEvent.click(await screen.findByRole("button", { name: CONTINUE }));
+  fireEvent.click(await screen.findByRole("button", { name: saveLabel }));
+}
+
+describe("PolicySetupWizard", () => {
+  beforeEach(() => {
+    fetchSources.mockResolvedValue({ sources: [] });
+  });
+
+  it("round-trips a saved step's backend params on edit", async () => {
+    const onSubmit = vi.fn().mockResolvedValue(undefined);
+    const entry = editEntry([
+      {
+        operation: "/api/v1/security/auto-redact",
+        parameters: { listOfText: "foo\nbar", useRegex: true },
+      },
+    ]);
+
+    render(
+      <PolicySetupWizard entry={entry} onClose={vi.fn()} onSubmit={onSubmit} />,
+    );
+    await submitWizard(SAVE_CHANGES);
+
+    await waitFor(() => expect(onSubmit).toHaveBeenCalled());
+    const result = onSubmit.mock.calls[0][1] as PolicySetupResult;
+    // Only the saved tool is enabled on edit, and its patterns survive the wire -> UI -> wire trip.
+    expect(result.steps).toEqual([
+      expect.objectContaining({
+        operation: "/api/v1/security/auto-redact",
+        parameters: expect.objectContaining({ listOfText: "foo\nbar" }),
+      }),
+    ]);
+  });
+
+  it("seeds the preset chain for a new policy (redact + sanitize on, watermark off)", async () => {
+    const onSubmit = vi.fn().mockResolvedValue(undefined);
+    const entry: CatalogueEntry = {
+      category: security,
+      config: securityConfig,
+      policy: null,
+    };
+
+    render(
+      <PolicySetupWizard entry={entry} onClose={vi.fn()} onSubmit={onSubmit} />,
+    );
+    await submitWizard(ENABLE);
+
+    await waitFor(() => expect(onSubmit).toHaveBeenCalled());
+    const result = onSubmit.mock.calls[0][1] as PolicySetupResult;
+    const endpoints = result.steps.map((s) => s.operation);
+    expect(endpoints).toEqual([
+      "/api/v1/security/auto-redact",
+      "/api/v1/security/sanitize-pdf",
+    ]);
+    // Redact carries the preset PII patterns as the backend's listOfText.
+    const redact = result.steps[0].parameters as { listOfText?: string };
+    expect(redact.listOfText).toBeTruthy();
+  });
+});

frontend/editor/src/portal/components/policies/PolicySetupWizard.tsx

@@ -13,23 +13,24 @@ import {
 } from "@app/ui";
 import { SettingsRow } from "@app/ui/SettingsRow";
 import {
-  TOOL_ENDPOINTS,
   humanizeEndpoint,
   type CatalogueEntry,
   type PipelineStep,
   type PolicySetupResult,
 } from "@portal/api/policies";
-import type { ToolRegistry, ToolRegistryEntry } from "@app/data/toolsTaxonomy";
 import {
-  deserializeToolStep,
-  serializeStepFromEndpoint,
-} from "@app/hooks/tools/shared/toolAutomation";
+  policyEndpoint,
+  policyStepFromWire,
+  policyStepToWire,
+  type PolicyParams,
+  type PolicyToolId,
+  type PolicyToolStep,
+} from "@app/policies/operations";
 import { fetchSources } from "@portal/api/sources";
 import { useAsync } from "@portal/hooks/useAsync";
 import { PolicyFieldRow } from "@portal/components/policies/PolicyFieldRow";
 import { policyIcon } from "@portal/components/policies/policyIcons";
 import { sourceTypeMeta } from "@portal/components/sources/sourceTypes";
-import { useToolRegistry } from "@app/contexts/ToolRegistryContext";
 import { PolicyRedactConfig } from "@app/components/policies/PolicyRedactConfig";
 import { PolicyWatermarkConfig } from "@app/components/policies/PolicyWatermarkConfig";
 import "@portal/views/Policies.css";
@@ -47,12 +48,8 @@ interface PolicySetupWizardProps {
 
 type Step = "workflow" | "settings";
 
-/** A configurable tool in the workflow step: whether it runs + its params. */
-interface ToolState {
-  operation: string;
-  enabled: boolean;
-  parameters: Record<string, unknown>;
-}
+/** A policy step plus whether it runs. */
+type ToolState = PolicyToolStep & { enabled: boolean };
 
 /** Resolve each field's effective value: saved override, else definition default. */
 function resolveFieldValues(
@@ -69,9 +66,8 @@ function resolveFieldValues(
  * round-trips); otherwise the category preset's default chain. Each preset step
  * starts enabled — the user toggles tools off in the workflow.
  */
-// Temporary: tracks which tools start disabled until the tool registry lands in
-// the portal and can drive this via registry metadata or a defaultEnabled flag.
-const DISABLED_BY_DEFAULT = new Set(["/api/v1/security/add-watermark"]);
+// Temporary until the catalogue carries a defaultEnabled flag.
+const DISABLED_BY_DEFAULT = new Set<PolicyToolId>(["watermark"]);
 
 /**
  * Policy-facing framing for each capability a policy can include. Labels and
@@ -81,43 +77,43 @@ const DISABLED_BY_DEFAULT = new Set(["/api/v1/security/add-watermark"]);
  * the humanised endpoint name with no description.
  */
 const CAPABILITY_META: Record<
-  string,
+  PolicyToolId,
   { labelKey: string; labelEn: string; descKey: string; descEn: string }
 > = {
-  [TOOL_ENDPOINTS.redact]: {
+  redact: {
     labelKey: "portal.policies.wizard.capability.redact.label",
     labelEn: "Redact sensitive information",
     descKey: "portal.policies.wizard.capability.redact.desc",
     descEn:
       "Finds and blacks out sensitive details — like Social Security and card numbers — so they can't be read.",
   },
-  [TOOL_ENDPOINTS.sanitize]: {
+  sanitize: {
     labelKey: "portal.policies.wizard.capability.sanitize.label",
     labelEn: "Strip active content",
     descKey: "portal.policies.wizard.capability.sanitize.desc",
     descEn:
       "Removes hidden JavaScript so nothing can run automatically when the document is opened.",
   },
-  [TOOL_ENDPOINTS.watermark]: {
+  watermark: {
     labelKey: "portal.policies.wizard.capability.watermark.label",
     labelEn: "Apply a watermark",
     descKey: "portal.policies.wizard.capability.watermark.desc",
     descEn: "Stamps a visible mark (e.g. “Confidential”) across every page.",
   },
-  [TOOL_ENDPOINTS.ocr]: {
+  ocr: {
     labelKey: "portal.policies.wizard.capability.ocr.label",
     labelEn: "Make text searchable",
     descKey: "portal.policies.wizard.capability.ocr.desc",
     descEn: "Runs OCR so scanned pages become selectable, searchable text.",
   },
-  [TOOL_ENDPOINTS.flatten]: {
+  flatten: {
     labelKey: "portal.policies.wizard.capability.flatten.label",
     labelEn: "Flatten the document",
     descKey: "portal.policies.wizard.capability.flatten.desc",
     descEn:
       "Merges form fields and annotations into the page so they can't be edited.",
   },
-  [TOOL_ENDPOINTS.compress]: {
+  compress: {
     labelKey: "portal.policies.wizard.capability.compress.label",
     labelEn: "Reduce file size",
     descKey: "portal.policies.wizard.capability.compress.desc",
@@ -125,29 +121,24 @@ const CAPABILITY_META: Record<
   },
 };
 
-function seedTools(
-  entry: CatalogueEntry,
-  registry: Partial<ToolRegistry>,
-): ToolState[] {
+function seedTools(entry: CatalogueEntry): ToolState[] {
   const savedSteps = entry.policy?.steps ?? [];
-  const savedByOp = new Map(savedSteps.map((s) => [s.operation, s]));
-  // Always use defaultOperations as the canonical list so tools added after a
-  // policy was first saved still appear when editing.
-  return entry.config.defaultOperations.map((s) => {
-    const saved = savedByOp.get(s.operation);
+  const savedByTool = new Map<PolicyToolId, PolicyToolStep>();
+  for (const wire of savedSteps) {
+    const step = policyStepFromWire(wire);
+    if (step) savedByTool.set(step.toolId, step);
+  }
+  // defaultOperations is the canonical list (so tools added later still show on edit); a saved
+  // step's params win over the preset.
+  return entry.config.defaultOperations.map((preset) => {
+    const saved = savedByTool.get(preset.toolId);
     return {
-      operation: s.operation,
+      ...(saved ?? preset),
       enabled: saved
         ? true
         : savedSteps.length > 0
           ? false
-          : !DISABLED_BY_DEFAULT.has(s.operation),
-      // Saved steps are in the backend contract shape; map them back to the UI
-      // shape the config controls edit (e.g. `listOfText` -> `wordsToRedact`).
-      // Presets are already authored in the UI shape, so use them as-is.
-      parameters: saved
-        ? deserializeToolStep(saved, registry).params
-        : s.parameters,
+          : !DISABLED_BY_DEFAULT.has(preset.toolId),
     };
   });
 }
@@ -185,26 +176,12 @@ function PolicySetupWizardBody({
   onSubmit: (entry: CatalogueEntry, result: PolicySetupResult) => Promise<void>;
 }) {
   const { t } = useTranslation();
-  const { allTools: toolRegistry } = useToolRegistry();
-
-  // Portal tool operations are endpoint paths (/api/v1/…), not short registry IDs.
-  // Build a reverse map so we can look up icons and display names by endpoint.
-  const registryByEndpoint = useMemo(() => {
-    const map = new Map<string, ToolRegistryEntry>();
-    for (const entry of Object.values(toolRegistry)) {
-      const ep = (entry as ToolRegistryEntry).operationConfig?.endpoint;
-      if (typeof ep === "string") map.set(ep, entry as ToolRegistryEntry);
-    }
-    return map;
-  }, [toolRegistry]);
 
   const { category, config, policy } = entry;
   const isEdit = policy != null;
 
   const [step, setStep] = useState<Step>("workflow");
-  const [tools, setTools] = useState<ToolState[]>(() =>
-    seedTools(entry, toolRegistry),
-  );
+  const [tools, setTools] = useState<ToolState[]>(() => seedTools(entry));
   const [fieldValues, setFieldValues] = useState(() =>
     resolveFieldValues(entry),
   );
@@ -213,22 +190,11 @@ function PolicySetupWizardBody({
   );
 
   const sourcesAsync = useAsync(() => fetchSources(), []);
-  const availableSources = useMemo(() => {
-    const backendSources = (sourcesAsync.data?.sources ?? []).filter(
-      (s) => s.status !== "disabled",
-    );
-    const editorSource = {
-      id: "editor",
-      name: t("portal.sources.types.editor.label"),
-      type: "editor",
-      status: "active" as const,
-      referenceCount: 0,
-      referencingPolicies: [],
-      config: [],
-      docsTotal: null,
-    };
-    return [editorSource, ...backendSources];
-  }, [sourcesAsync.data, t]);
+  const availableSources = useMemo(
+    () =>
+      (sourcesAsync.data?.sources ?? []).filter((s) => s.status !== "disabled"),
+    [sourcesAsync.data],
+  );
   // Document-type scoping has no UI; preserve any saved scope on edit and
   // default new policies to all document types.
   const [scopeTypes] = useState<string[]>(policy?.state.scopeTypes ?? []);
@@ -256,9 +222,20 @@ function PolicySetupWizardBody({
 
   const enabledTools = useMemo(() => tools.filter((tl) => tl.enabled), [tools]);
 
-  function patchTool(operation: string, patch: Partial<ToolState>) {
+  function setToolEnabled(toolId: PolicyToolId, enabled: boolean) {
     setTools((prev) =>
-      prev.map((tl) => (tl.operation === operation ? { ...tl, ...patch } : tl)),
+      prev.map((tl) => (tl.toolId === toolId ? { ...tl, enabled } : tl)),
+    );
+  }
+
+  function setToolParams<Id extends PolicyToolId>(
+    toolId: Id,
+    params: PolicyParams<Id>,
+  ) {
+    setTools((prev) =>
+      prev.map((tl) =>
+        tl.toolId === toolId ? ({ ...tl, params } as ToolState) : tl,
+      ),
     );
   }
 
@@ -277,11 +254,8 @@ function PolicySetupWizardBody({
     }
     setError(null);
     setSubmitting(true);
-    // Map each tool's UI-shaped params (e.g. redact's `wordsToRedact`) into the
-    // backend step contract (e.g. `listOfText`) via its `toApiParams`; saving the
-    // UI shape verbatim would drop those fields and the step would run with none.
     const steps: PipelineStep[] = enabledTools.map((tl) =>
-      serializeStepFromEndpoint(tl.operation, tl.parameters, toolRegistry),
+      policyStepToWire(tl),
     );
     try {
       await onSubmit(entry, {
@@ -386,20 +360,16 @@ function PolicySetupWizardBody({
           <Card padding="none">
             <div className="portal-policies__capabilities">
               {tools.map((tl) => {
-                const meta = CAPABILITY_META[tl.operation];
+                const meta = CAPABILITY_META[tl.toolId];
                 const label = meta
                   ? t(meta.labelKey, meta.labelEn)
-                  : (registryByEndpoint.get(tl.operation)?.name ??
-                    humanizeEndpoint(tl.operation, t));
+                  : humanizeEndpoint(policyEndpoint(tl.toolId), t);
                 const description = meta
                   ? t(meta.descKey, meta.descEn)
                   : undefined;
-                const hasConfig =
-                  tl.operation === TOOL_ENDPOINTS.redact ||
-                  tl.operation === TOOL_ENDPOINTS.watermark;
                 return (
                   <div
-                    key={tl.operation}
+                    key={tl.toolId}
                     className="portal-policies__capability"
                     data-on={tl.enabled || undefined}
                   >
@@ -411,27 +381,27 @@ function PolicySetupWizardBody({
                           size="sm"
                           checked={tl.enabled}
                           onChange={(checked) =>
-                            patchTool(tl.operation, { enabled: checked })
+                            setToolEnabled(tl.toolId, checked)
                           }
                           label=""
                         />
                       }
                     />
-                    {tl.enabled && hasConfig && (
+                    {tl.enabled && (
                       <div className="portal-policies__capability-config">
-                        {tl.operation === TOOL_ENDPOINTS.redact && (
+                        {tl.toolId === "redact" && (
                           <PolicyRedactConfig
-                            parameters={tl.parameters}
-                            onChange={(parameters) =>
-                              patchTool(tl.operation, { parameters })
+                            parameters={tl.params}
+                            onChange={(params) =>
+                              setToolParams("redact", params)
                             }
                           />
                         )}
-                        {tl.operation === TOOL_ENDPOINTS.watermark && (
+                        {tl.toolId === "watermark" && (
                           <PolicyWatermarkConfig
-                            parameters={tl.parameters}
-                            onChange={(parameters) =>
-                              patchTool(tl.operation, { parameters })
+                            parameters={tl.params}
+                            onChange={(params) =>
+                              setToolParams("watermark", params)
                             }
                           />
                         )}
@@ -475,9 +445,8 @@ function PolicySetupWizardBody({
               {t("portal.policies.wizard.sources.loading")}
             </p>
           ) : (
-            // The editor is always an available source (unconditionally prepended
-            // to availableSources), so the list is never empty — no "no sources"
-            // state exists.
+            // The backend always returns the editor as a virtual source, so the
+            // loaded list is never empty - no "no sources" state exists.
             <div className="portal-policies__sources">
               {availableSources.map((src) => (
                 // A selectable multi-line tile (icon + name + type + check).

frontend/editor/src/portal/mocks/policies.ts

@@ -4,13 +4,33 @@
  * only builds seed data for the MSW handlers and tests.
  */
 
-import type { PolicyRunView, WirePolicy } from "@app/policies/types";
-import { POLICY_CONFIG } from "@portal/api/policies";
+import type {
+  PolicyRunView,
+  WirePipelineStep,
+  WirePolicy,
+} from "@app/policies/types";
 
 /* ──────────────────────────────────────────────────────────────────────── */
 /*  Seed data — real backend wire format                                      */
 /* ──────────────────────────────────────────────────────────────────────── */
 
+// Literal wire steps (not derived from the catalogue) so this fixtures module stays independent of
+// @portal/api/policies and its heavy tool-operation import graph.
+const SECURITY_STEPS: WirePipelineStep[] = [
+  {
+    operation: "/api/v1/security/auto-redact",
+    parameters: {
+      listOfText: "",
+      useRegex: true,
+      convertPDFToImage: true,
+    },
+  },
+  {
+    operation: "/api/v1/security/sanitize-pdf",
+    parameters: { removeJavaScript: true },
+  },
+];
+
 export function seedPolicies(): WirePolicy[] {
   return [
     {
@@ -19,7 +39,7 @@ export function seedPolicies(): WirePolicy[] {
       owner: "security@acme.com",
       enabled: true,
       trigger: null,
-      steps: POLICY_CONFIG.security.defaultOperations,
+      steps: SECURITY_STEPS,
       output: {
         type: "inline",
         options: {

frontend/editor/src/proprietary/components/policies/PolicyPiiField.tsx

@@ -1,6 +1,7 @@
 import { MultiSelect } from "@app/ui/MultiSelect";
 import { useTranslation } from "react-i18next";
 import { PII_PRESETS } from "@app/data/policyDefinitions";
+import type { RedactParameters } from "@app/hooks/tools/redact/useRedactParameters";
 
 /** The set of preset regexes — used to separate preset words from custom ones. */
 export const PRESET_PATTERNS = new Set(PII_PRESETS.map((p) => p.pattern));
@@ -8,8 +9,8 @@ const PATTERN_BY_VALUE = new Map(PII_PRESETS.map((p) => [p.value, p.pattern]));
 const VALUE_BY_PATTERN = new Map(PII_PRESETS.map((p) => [p.pattern, p.value]));
 
 interface PolicyPiiFieldProps {
-  parameters: Record<string, unknown>;
-  onChange: (parameters: Record<string, unknown>) => void;
+  parameters: RedactParameters;
+  onChange: (parameters: RedactParameters) => void;
   disabled?: boolean;
 }
 
@@ -26,9 +27,7 @@ export function PolicyPiiField({
   disabled,
 }: PolicyPiiFieldProps) {
   const { t } = useTranslation();
-  const words = Array.isArray(parameters.wordsToRedact)
-    ? (parameters.wordsToRedact as string[])
-    : [];
+  const words = parameters.wordsToRedact;
   const selected = words
     .map((w) => VALUE_BY_PATTERN.get(w))
     .filter((v): v is string => Boolean(v));

frontend/editor/src/proprietary/components/policies/PolicyRedactConfig.tsx

@@ -1,9 +1,10 @@
 import { useEffect } from "react";
 import { PolicyPiiField } from "@app/components/policies/PolicyPiiField";
+import type { RedactParameters } from "@app/hooks/tools/redact/useRedactParameters";
 
 interface PolicyRedactConfigProps {
-  parameters: Record<string, unknown>;
-  onChange: (parameters: Record<string, unknown>) => void;
+  parameters: RedactParameters;
+  onChange: (parameters: RedactParameters) => void;
   disabled?: boolean;
 }
 

frontend/editor/src/proprietary/components/policies/PolicyWatermarkConfig.tsx

@@ -3,8 +3,8 @@ import AddWatermarkSingleStepSettings from "@app/components/tools/addWatermark/A
 import type { AddWatermarkParameters } from "@app/hooks/tools/addWatermark/useAddWatermarkParameters";
 
 interface PolicyWatermarkConfigProps {
-  parameters: Record<string, unknown>;
-  onChange: (parameters: Record<string, unknown>) => void;
+  parameters: AddWatermarkParameters;
+  onChange: (parameters: AddWatermarkParameters) => void;
   disabled?: boolean;
 }
 
@@ -20,7 +20,7 @@ export function PolicyWatermarkConfig({
   disabled,
 }: PolicyWatermarkConfigProps) {
   useEffect(() => {
-    const patch: Record<string, unknown> = {};
+    const patch: Partial<AddWatermarkParameters> = {};
     if (parameters.convertPDFToImage !== true) patch.convertPDFToImage = true;
     // Policies only support text watermarks.
     if (parameters.watermarkType !== "text") patch.watermarkType = "text";
@@ -29,7 +29,7 @@ export function PolicyWatermarkConfig({
 
   return (
     <AddWatermarkSingleStepSettings
-      parameters={parameters as unknown as AddWatermarkParameters}
+      parameters={parameters}
       onParameterChange={(key, value) =>
         onChange({ ...parameters, [key]: value })
       }

frontend/editor/src/proprietary/policies/operations.test.ts

@@ -0,0 +1,112 @@
+import { describe, expect, test } from "vitest";
+import {
+  POLICY_OPERATIONS,
+  policyEndpoint,
+  policyStep,
+  policyStepFromWire,
+  policyStepToWire,
+  policyToolIdForEndpoint,
+  type PolicyToolId,
+} from "@app/policies/operations";
+
+const ALL_TOOL_IDS = Object.keys(POLICY_OPERATIONS) as PolicyToolId[];
+
+describe("POLICY_OPERATIONS", () => {
+  test("every category operation is a typed descriptor with a known endpoint", () => {
+    // The catalogue uses these six across all categories; each must be wired.
+    expect(ALL_TOOL_IDS.sort()).toEqual([
+      "compress",
+      "flatten",
+      "ocr",
+      "redact",
+      "sanitize",
+      "watermark",
+    ]);
+    for (const id of ALL_TOOL_IDS) {
+      expect(POLICY_OPERATIONS[id].endpoint).toBe(policyEndpoint(id));
+      expect(typeof POLICY_OPERATIONS[id].toApi).toBe("function");
+      expect(typeof POLICY_OPERATIONS[id].fromApi).toBe("function");
+    }
+  });
+
+  test("policyEndpoint returns the pinned endpoint literal", () => {
+    expect(policyEndpoint("redact")).toBe("/api/v1/security/auto-redact");
+    expect(policyEndpoint("sanitize")).toBe("/api/v1/security/sanitize-pdf");
+    expect(policyEndpoint("watermark")).toBe("/api/v1/security/add-watermark");
+    expect(policyEndpoint("ocr")).toBe("/api/v1/misc/ocr-pdf");
+    expect(policyEndpoint("flatten")).toBe("/api/v1/misc/flatten");
+    expect(policyEndpoint("compress")).toBe("/api/v1/misc/compress-pdf");
+  });
+
+  test("policyToolIdForEndpoint maps endpoints back, and rejects non-policy ones", () => {
+    for (const id of ALL_TOOL_IDS) {
+      expect(policyToolIdForEndpoint(policyEndpoint(id))).toBe(id);
+    }
+    expect(policyToolIdForEndpoint("/api/v1/misc/repair")).toBeNull();
+    expect(policyToolIdForEndpoint("not-an-endpoint")).toBeNull();
+  });
+});
+
+describe("policyStep", () => {
+  test("merges partial params over the tool's defaults", () => {
+    const step = policyStep("redact", {
+      useRegex: true,
+      wordsToRedact: ["ssn", "card"],
+    });
+    expect(step.toolId).toBe("redact");
+    // Overrides applied...
+    expect(step.params.useRegex).toBe(true);
+    expect(step.params.wordsToRedact).toEqual(["ssn", "card"]);
+    // ...and untouched fields fall back to the tool's defaults.
+    expect(step.params.mode).toBe("automatic");
+    expect(step.params.redactColor).toBe("#000000");
+  });
+});
+
+describe("wire conversion", () => {
+  test("redact maps frontend params to the backend request model (wordsToRedact -> listOfText)", () => {
+    const wire = policyStepToWire(
+      policyStep("redact", {
+        useRegex: true,
+        convertPDFToImage: true,
+        wordsToRedact: ["ssn", "card"],
+      }),
+    );
+    expect(wire.operation).toBe("/api/v1/security/auto-redact");
+    // The backend field the endpoint actually reads, and no frontend-only `mode`/`wordsToRedact`.
+    expect(wire.parameters).toMatchObject({
+      listOfText: "ssn\ncard",
+      useRegex: true,
+      convertPDFToImage: true,
+    });
+    expect(wire.parameters).not.toHaveProperty("wordsToRedact");
+    expect(wire.parameters).not.toHaveProperty("mode");
+  });
+
+  test("every policy operation round-trips through wire and back", () => {
+    for (const id of ALL_TOOL_IDS) {
+      const step = policyStep(id);
+      const back = policyStepFromWire(policyStepToWire(step));
+      expect(back?.toolId).toBe(id);
+    }
+  });
+
+  test("redact round-trip preserves the configured patterns", () => {
+    const step = policyStep("redact", {
+      useRegex: true,
+      wordsToRedact: ["ssn", "card"],
+    });
+    const back = policyStepFromWire(policyStepToWire(step));
+    expect(back?.toolId).toBe("redact");
+    if (back?.toolId === "redact") {
+      expect(back.params.wordsToRedact).toEqual(["ssn", "card"]);
+      expect(back.params.useRegex).toBe(true);
+    }
+  });
+
+  test("a non-policy endpoint decodes to null", () => {
+    expect(
+      policyStepFromWire({ operation: "/api/v1/misc/repair", parameters: {} }),
+    ).toBeNull();
+  });
+});

frontend/editor/src/proprietary/policies/operations.ts

@@ -0,0 +1,131 @@
+/**
+ * The tool operations the Policies feature can run, each a typed {@link ToolOperationDescriptor}.
+ * Source of truth for the catalogue, wizard, and wire conversion. Add a tool here to use it in a
+ * policy - the catalogue can't reference an untyped operation.
+ */
+
+import { describeToolOperation } from "@app/hooks/tools/shared/toolOperationDescriptor";
+import { redactOperationConfig } from "@app/hooks/tools/redact/useRedactOperation";
+import { sanitizeOperationConfig } from "@app/hooks/tools/sanitize/useSanitizeOperation";
+import { addWatermarkOperationConfig } from "@app/hooks/tools/addWatermark/useAddWatermarkOperation";
+import { ocrOperationConfig } from "@app/hooks/tools/ocr/useOCROperation";
+import { flattenOperationConfig } from "@app/hooks/tools/flatten/useFlattenOperation";
+import { compressOperationConfig } from "@app/hooks/tools/compress/useCompressOperation";
+import type { ToolOperationDescriptor } from "@app/hooks/tools/shared/toolOperationDescriptor";
+import type { ToolApiParams, ToolEndpoint } from "@app/types/toolApiTypes";
+import type { WirePipelineStep } from "@app/policies/types";
+
+export const POLICY_OPERATIONS = {
+  redact: describeToolOperation(
+    "/api/v1/security/auto-redact",
+    redactOperationConfig,
+  ),
+  sanitize: describeToolOperation(
+    "/api/v1/security/sanitize-pdf",
+    sanitizeOperationConfig,
+  ),
+  watermark: describeToolOperation(
+    "/api/v1/security/add-watermark",
+    addWatermarkOperationConfig,
+  ),
+  ocr: describeToolOperation("/api/v1/misc/ocr-pdf", ocrOperationConfig),
+  flatten: describeToolOperation(
+    "/api/v1/misc/flatten",
+    flattenOperationConfig,
+  ),
+  compress: describeToolOperation(
+    "/api/v1/misc/compress-pdf",
+    compressOperationConfig,
+  ),
+} as const;
+
+export type PolicyToolId = keyof typeof POLICY_OPERATIONS;
+
+export type PolicyParams<Id extends PolicyToolId> =
+  (typeof POLICY_OPERATIONS)[Id] extends ToolOperationDescriptor<
+    ToolEndpoint,
+    infer P
+  >
+    ? P
+    : never;
+
+/** Discriminated on `toolId` so `params` matches the tool. */
+export type PolicyToolStep = {
+  [Id in PolicyToolId]: { toolId: Id; params: PolicyParams<Id> };
+}[PolicyToolId];
+
+export type PolicyToolStepOf<Id extends PolicyToolId> = Extract<
+  PolicyToolStep,
+  { toolId: Id }
+>;
+
+const POLICY_TOOL_IDS = Object.keys(POLICY_OPERATIONS) as PolicyToolId[];
+
+const TOOL_ID_BY_ENDPOINT = new Map<string, PolicyToolId>(
+  POLICY_TOOL_IDS.map((id) => [POLICY_OPERATIONS[id].endpoint, id]),
+);
+
+export function policyEndpoint(toolId: PolicyToolId): ToolEndpoint {
+  return POLICY_OPERATIONS[toolId].endpoint;
+}
+
+/** Tool id for an endpoint path, or null if it isn't a policy tool. */
+export function policyToolIdForEndpoint(endpoint: string): PolicyToolId | null {
+  return TOOL_ID_BY_ENDPOINT.get(endpoint) ?? null;
+}
+
+/** A step for `toolId`, partial params merged over the tool's defaults. */
+export function policyStep<Id extends PolicyToolId>(
+  toolId: Id,
+  params: Partial<PolicyParams<Id>> = {},
+): PolicyToolStepOf<Id> {
+  const defaults = POLICY_OPERATIONS[toolId].defaultParameters as object;
+  return {
+    toolId,
+    params: { ...defaults, ...(params as object) },
+  } as PolicyToolStepOf<Id>;
+}
+
+export function policyStepToWire(step: PolicyToolStep): WirePipelineStep {
+  return serializeStep(step);
+}
+
+// Generic over the id so `params` stays correlated with the descriptor; TS can't do that through
+// the union, so `op` is widened here (a contained cast at the wire boundary).
+function serializeStep<Id extends PolicyToolId>(step: {
+  toolId: Id;
+  params: PolicyParams<Id>;
+}): WirePipelineStep {
+  const op = POLICY_OPERATIONS[step.toolId] as ToolOperationDescriptor<
+    ToolEndpoint,
+    PolicyParams<Id>
+  >;
+  return {
+    operation: op.endpoint,
+    parameters: op.toApi(step.params) as Record<string, unknown>,
+  };
+}
+
+/** Wire step -> typed policy step, or null if the endpoint isn't a policy tool. */
+export function policyStepFromWire(
+  wire: WirePipelineStep,
+): PolicyToolStep | null {
+  const toolId = policyToolIdForEndpoint(wire.operation);
+  if (!toolId) return null;
+  return deserializeStep(toolId, wire.parameters);
+}
+
+function deserializeStep<Id extends PolicyToolId>(
+  toolId: Id,
+  parameters: Record<string, unknown>,
+): PolicyToolStepOf<Id> {
+  const op = POLICY_OPERATIONS[toolId] as ToolOperationDescriptor<
+    ToolEndpoint,
+    PolicyParams<Id>
+  >;
+  // Wire params are untyped JSON; this is the one point they enter the typed model.
+  const params = op.fromApi(
+    parameters as unknown as ToolApiParams[ToolEndpoint],
+  );
+  return { toolId, params } as unknown as PolicyToolStepOf<Id>;
+}