diff --git a/app/common/src/main/java/stirling/software/common/service/AutomationRunContext.java b/app/common/src/main/java/stirling/software/common/service/AutomationRunContext.java
new file mode 100644
index 0000000000..9ef4cb0511
--- /dev/null
+++ b/app/common/src/main/java/stirling/software/common/service/AutomationRunContext.java
@@ -0,0 +1,56 @@
+package stirling.software.common.service;
+
+/**
+ * Thread-scoped correlation id for one automation run — a single pipeline, policy, or AI-workflow
+ * execution over its input file(s).
+ *
+ *
Automations dispatch each tool step as a separate internal loopback POST via {@link
+ * InternalApiClient}. The orchestrator opens a run scope around its dispatch loop; {@code
+ * InternalApiClient} reads {@link #current()} and stamps it on every sub-step request as {@link
+ * #RUN_ID_HEADER}. The SaaS PAYG interceptor uses that header so all sub-steps of ONE run group
+ * into a single charge, while two separate runs that happen to touch identical bytes stay
+ * distinct charges (the old content+time-window grouping merged them).
+ *
+ *
Sub-steps dispatch synchronously on the orchestrator's own thread (loopback {@code
+ * RestTemplate}), so this ThreadLocal is visible to {@code InternalApiClient}. The id then crosses
+ * to the receiving request thread via the HTTP header — never via this ThreadLocal.
+ *
+ *
No-op when the id is absent (a standalone tool call): the interceptor treats a missing run id
+ * as "its own charge", which is exactly what a one-off call should be.
+ */
+public final class AutomationRunContext {
+
+ /** Header carrying the run id on internal sub-step dispatches. */
+ public static final String RUN_ID_HEADER = "X-Stirling-Run-Id";
+
+ private static final ThreadLocal CURRENT = new ThreadLocal<>();
+
+ private AutomationRunContext() {}
+
+ /**
+ * Opens a run scope on the current thread. Returns an {@link AutoCloseable} that restores the
+ * previously-active id (nesting-safe) — use in try-with-resources around the dispatch loop.
+ */
+ public static Scope open(String runId) {
+ String previous = CURRENT.get();
+ CURRENT.set(runId);
+ return () -> {
+ if (previous == null) {
+ CURRENT.remove();
+ } else {
+ CURRENT.set(previous);
+ }
+ };
+ }
+
+ /** The run id active on this thread, or {@code null} when not inside a run scope. */
+ public static String current() {
+ return CURRENT.get();
+ }
+
+ /** AutoCloseable whose {@link #close()} declares no checked exception. */
+ public interface Scope extends AutoCloseable {
+ @Override
+ void close();
+ }
+}
diff --git a/app/common/src/main/java/stirling/software/common/service/InternalApiClient.java b/app/common/src/main/java/stirling/software/common/service/InternalApiClient.java
index 8df2d5e410..39e27457b3 100644
--- a/app/common/src/main/java/stirling/software/common/service/InternalApiClient.java
+++ b/app/common/src/main/java/stirling/software/common/service/InternalApiClient.java
@@ -111,6 +111,14 @@ public class InternalApiClient {
// step inside a policy run must bill as AUTOMATION, not AI). Set unconditionally because
// every caller of this dispatcher is an automation surface by design.
headers.add(AUTOMATION_HEADER, "true");
+ // Propagate the current automation run id (set by the orchestrator around its dispatch
+ // loop) so the PAYG interceptor groups every sub-step of this one run into a single charge,
+ // and never merges two separate runs that happen to touch identical bytes. Absent → the
+ // receiving call is treated as standalone. See AutomationRunContext.
+ String runId = AutomationRunContext.current();
+ if (runId != null && !runId.isEmpty()) {
+ headers.add(AutomationRunContext.RUN_ID_HEADER, runId);
+ }
// A no-file ai/tools call (e.g. create-pdf-from-html-agent) sends only string params, so
// without this RestTemplate would use urlencoded instead of the multipart the controller
diff --git a/app/core/src/main/java/stirling/software/SPDF/controller/api/pipeline/PipelineProcessor.java b/app/core/src/main/java/stirling/software/SPDF/controller/api/pipeline/PipelineProcessor.java
index c57a43f27f..c4f1c76f5e 100644
--- a/app/core/src/main/java/stirling/software/SPDF/controller/api/pipeline/PipelineProcessor.java
+++ b/app/core/src/main/java/stirling/software/SPDF/controller/api/pipeline/PipelineProcessor.java
@@ -10,6 +10,7 @@ import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Map.Entry;
+import java.util.UUID;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
@@ -27,6 +28,7 @@ import stirling.software.SPDF.model.PipelineConfig;
import stirling.software.SPDF.model.PipelineOperation;
import stirling.software.SPDF.model.PipelineResult;
import stirling.software.SPDF.service.ApiDocService;
+import stirling.software.common.service.AutomationRunContext;
import stirling.software.common.service.InternalApiClient;
import stirling.software.common.util.TempFileManager;
import stirling.software.common.util.ZipExtractionUtils;
@@ -71,6 +73,17 @@ public class PipelineProcessor {
PipelineResult runPipelineAgainstFiles(List outputFiles, PipelineConfig config)
throws Exception {
+ // One pipeline execution = one automation run. Scope a run id so every tool sub-step
+ // dispatched via InternalApiClient groups into a single charge on the SaaS billing side
+ // (see AutomationRunContext); pipeline steps run synchronously on this thread.
+ try (AutomationRunContext.Scope ignored =
+ AutomationRunContext.open(UUID.randomUUID().toString())) {
+ return runPipelineAgainstFilesInternal(outputFiles, config);
+ }
+ }
+
+ private PipelineResult runPipelineAgainstFilesInternal(
+ List outputFiles, PipelineConfig config) throws Exception {
PipelineResult result = new PipelineResult();
ByteArrayOutputStream logStream = new ByteArrayOutputStream();
diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/policy/engine/PolicyEngine.java b/app/proprietary/src/main/java/stirling/software/proprietary/policy/engine/PolicyEngine.java
index 98a6f90096..cef0f29e6f 100644
--- a/app/proprietary/src/main/java/stirling/software/proprietary/policy/engine/PolicyEngine.java
+++ b/app/proprietary/src/main/java/stirling/software/proprietary/policy/engine/PolicyEngine.java
@@ -21,6 +21,7 @@ import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import stirling.software.common.model.job.ResultFile;
+import stirling.software.common.service.AutomationRunContext;
import stirling.software.common.service.FileStorage;
import stirling.software.common.service.InternalApiTimeoutException;
import stirling.software.common.service.JobOwnershipService;
@@ -199,63 +200,78 @@ public class PolicyEngine {
PolicyProgressListener listener,
CompletableFuture completion) {
String runId = run.getRunId();
- try {
- run.markRunning();
- PolicyExecutionResult result =
- stepExecutor.execute(run.getDefinition(), inputs, listener);
- OutputSpec output = run.getDefinition().output();
- List outputs =
- sinkFor(output)
- .deliver(
- new OutputDelivery(runId, run.getPolicyId()),
- result.files(),
- output);
- taskManager.setMultipleFileResults(runId, outputs);
- taskManager.setComplete(runId);
- run.complete(outputs);
- } catch (PolicyInputRequiredException e) {
- // Expected path: suspend rather than fail. Persist intermediates as fileIds so the run
- // can resume after this worker thread is gone.
- WaitState wait = suspend(e);
- run.waitForInput(wait);
- taskManager.addNote(runId, "Waiting for input: " + e.getMessage());
- } catch (InternalApiTimeoutException e) {
- String message = toolTimeoutMessage(e);
- log.error(
- "Policy run {} timed out on {}: {}",
- runId,
- e.getEndpointPath(),
- e.getMessage());
- run.fail(message);
- taskManager.setError(runId, message);
- } catch (RestClientResponseException e) {
- // A downstream tool call returned an error status. When it's a structured entitlement
- // response (401/402 with a JSON `error` sentinel), surface that code onto the run so
- // the
- // client can react — e.g. pop the usage-limit modal — instead of only seeing a generic
- // failure. We don't interpret the code here (that would couple this module to the saas
- // billing layer); we just pass it through for the client to map. Other statuses fall
- // through to the generic failure below.
- String code = DownstreamEntitlementError.extractCode(e);
- if (code != null) {
- log.info("Policy run {} blocked by downstream entitlement gate ({})", runId, code);
- String message = "Usage limit reached";
- run.failWithCode(message, code, DownstreamEntitlementError.extractSubscribed(e));
- taskManager.setError(runId, message);
- } else {
- String message = "Policy run failed: " + e.getMessage();
- log.error("Policy run {} failed (downstream HTTP error)", runId, e);
+ // One policy run = one automation run. Scope the run id on this worker thread (the async
+ // hop already happened) so every tool sub-step dispatched via InternalApiClient groups into
+ // a single charge, and two separate policy runs on the same document stay distinct charges.
+ try (AutomationRunContext.Scope runScope = AutomationRunContext.open(runId)) {
+ try {
+ run.markRunning();
+ PolicyExecutionResult result =
+ stepExecutor.execute(run.getDefinition(), inputs, listener);
+ OutputSpec output = run.getDefinition().output();
+ List outputs =
+ sinkFor(output)
+ .deliver(
+ new OutputDelivery(runId, run.getPolicyId()),
+ result.files(),
+ output);
+ taskManager.setMultipleFileResults(runId, outputs);
+ taskManager.setComplete(runId);
+ run.complete(outputs);
+ } catch (PolicyInputRequiredException e) {
+ // Expected path: suspend rather than fail. Persist intermediates as fileIds so the
+ // run
+ // can resume after this worker thread is gone.
+ WaitState wait = suspend(e);
+ run.waitForInput(wait);
+ taskManager.addNote(runId, "Waiting for input: " + e.getMessage());
+ } catch (InternalApiTimeoutException e) {
+ String message = toolTimeoutMessage(e);
+ log.error(
+ "Policy run {} timed out on {}: {}",
+ runId,
+ e.getEndpointPath(),
+ e.getMessage());
run.fail(message);
taskManager.setError(runId, message);
+ } catch (RestClientResponseException e) {
+ // A downstream tool call returned an error status. When it's a structured
+ // entitlement
+ // response (401/402 with a JSON `error` sentinel), surface that code onto the run
+ // so
+ // the
+ // client can react — e.g. pop the usage-limit modal — instead of only seeing a
+ // generic
+ // failure. We don't interpret the code here (that would couple this module to the
+ // saas
+ // billing layer); we just pass it through for the client to map. Other statuses
+ // fall
+ // through to the generic failure below.
+ String code = DownstreamEntitlementError.extractCode(e);
+ if (code != null) {
+ log.info(
+ "Policy run {} blocked by downstream entitlement gate ({})",
+ runId,
+ code);
+ String message = "Usage limit reached";
+ run.failWithCode(
+ message, code, DownstreamEntitlementError.extractSubscribed(e));
+ taskManager.setError(runId, message);
+ } else {
+ String message = "Policy run failed: " + e.getMessage();
+ log.error("Policy run {} failed (downstream HTTP error)", runId, e);
+ run.fail(message);
+ taskManager.setError(runId, message);
+ }
+ } catch (Exception e) {
+ String message = "Policy run failed: " + e.getMessage();
+ log.error("Policy run {} failed", runId, e);
+ run.fail(message);
+ taskManager.setError(runId, message);
+ } finally {
+ // Always resolve so stream/await callers unblock.
+ completion.complete(run);
}
- } catch (Exception e) {
- String message = "Policy run failed: " + e.getMessage();
- log.error("Policy run {} failed", runId, e);
- run.fail(message);
- taskManager.setError(runId, message);
- } finally {
- // Always resolve so stream/await callers unblock.
- completion.complete(run);
}
}
diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java
index 0e9460eca7..74e68d4602 100644
--- a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java
+++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java
@@ -7,6 +7,7 @@ import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
+import java.util.UUID;
import java.util.stream.Collectors;
import org.apache.commons.io.FilenameUtils;
@@ -27,6 +28,7 @@ import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import stirling.software.common.model.ApplicationProperties;
+import stirling.software.common.service.AutomationRunContext;
import stirling.software.common.service.CustomPDFDocumentFactory;
import stirling.software.common.service.FileStorage;
import stirling.software.common.service.InternalApiTimeoutException;
@@ -156,33 +158,41 @@ public class AiWorkflowService {
throws IOException {
validateRequest(request);
- // Key by opaque file id, not filename. Filenames aren't guaranteed unique across an
- // upload (users can rotate the same 'scan.pdf' twice), and the engine identifies files
- // by id in every response shape that asks Java to look a file up again.
- Map filesById = new LinkedHashMap<>();
- List files = new ArrayList<>();
- for (AiWorkflowFileInput fileInput : request.getFileInputs()) {
- MultipartFile multipartFile = fileInput.getFileInput();
- AiFile aiFile =
- new AiFile(
- fileIdStrategy.idFor(multipartFile),
- multipartFile.getOriginalFilename());
- filesById.put(aiFile.getId(), multipartFile);
- files.add(aiFile);
- }
+ // One AI orchestration = one automation run. Scope a run id (on whichever thread runs
+ // orchestrate — request thread for sync, stream-executor for streaming) so every tool
+ // sub-step it dispatches via PolicyExecutor → InternalApiClient groups into one charge.
+ try (AutomationRunContext.Scope ignored =
+ AutomationRunContext.open(UUID.randomUUID().toString())) {
- WorkflowTurnRequest initialRequest = new WorkflowTurnRequest();
- initialRequest.setUserMessage(request.getUserMessage().trim());
- initialRequest.setFiles(files);
- initialRequest.setConversationHistory(new ArrayList<>(request.getConversationHistory()));
- initialRequest.setEnabledEndpoints(endpointResolver.getEnabledEndpointUrls());
- listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.ANALYZING));
+ // Key by opaque file id, not filename. Filenames aren't guaranteed unique across an
+ // upload (users can rotate the same 'scan.pdf' twice), and the engine identifies files
+ // by id in every response shape that asks Java to look a file up again.
+ Map filesById = new LinkedHashMap<>();
+ List files = new ArrayList<>();
+ for (AiWorkflowFileInput fileInput : request.getFileInputs()) {
+ MultipartFile multipartFile = fileInput.getFileInput();
+ AiFile aiFile =
+ new AiFile(
+ fileIdStrategy.idFor(multipartFile),
+ multipartFile.getOriginalFilename());
+ filesById.put(aiFile.getId(), multipartFile);
+ files.add(aiFile);
+ }
- WorkflowState state = new WorkflowState.Pending(initialRequest);
- while (state instanceof WorkflowState.Pending pending) {
- state = advance(pending.request(), filesById, listener);
+ WorkflowTurnRequest initialRequest = new WorkflowTurnRequest();
+ initialRequest.setUserMessage(request.getUserMessage().trim());
+ initialRequest.setFiles(files);
+ initialRequest.setConversationHistory(
+ new ArrayList<>(request.getConversationHistory()));
+ initialRequest.setEnabledEndpoints(endpointResolver.getEnabledEndpointUrls());
+ listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.ANALYZING));
+
+ WorkflowState state = new WorkflowState.Pending(initialRequest);
+ while (state instanceof WorkflowState.Pending pending) {
+ state = advance(pending.request(), filesById, listener);
+ }
+ return ((WorkflowState.Terminal) state).response();
}
- return ((WorkflowState.Terminal) state).response();
}
private WorkflowState advance(
diff --git a/app/saas/src/main/java/stirling/software/saas/ai/controller/AiCreateController.java b/app/saas/src/main/java/stirling/software/saas/ai/controller/AiCreateController.java
index 371060c032..28bbfbdea0 100644
--- a/app/saas/src/main/java/stirling/software/saas/ai/controller/AiCreateController.java
+++ b/app/saas/src/main/java/stirling/software/saas/ai/controller/AiCreateController.java
@@ -117,7 +117,8 @@ public class AiCreateController {
user.getTeam().getId(),
source,
ProcessType.SINGLE_TOOL,
- BillingCategory.AI);
+ BillingCategory.AI,
+ null);
jobChargeService.chargeStandalone(ctx, 1);
} catch (RuntimeException e) {
log.warn(
diff --git a/app/saas/src/main/java/stirling/software/saas/payg/api/PaygWalletController.java b/app/saas/src/main/java/stirling/software/saas/payg/api/PaygWalletController.java
index 90616bea55..18da81ae2b 100644
--- a/app/saas/src/main/java/stirling/software/saas/payg/api/PaygWalletController.java
+++ b/app/saas/src/main/java/stirling/software/saas/payg/api/PaygWalletController.java
@@ -173,7 +173,9 @@ public class PaygWalletController {
int spend = clampToInt(snap.periodSpendUnits());
Integer limit = snap.periodCapUnits() != null ? clampToInt(snap.periodCapUnits()) : null;
- CategoryBreakdown breakdown = buildBreakdown(teamId, snap.periodStart(), snap.periodEnd());
+ BreakdownPair breakdowns = buildBreakdowns(teamId, snap.periodStart(), snap.periodEnd());
+ UsageAnalytics analytics =
+ buildUsageAnalytics(teamId, snap.periodStart(), snap.periodEnd());
// Estimated bill = paid (Stripe-metered) docs this period × rate — the free portion was
// already netted out at charge time, so this is the metered total, not spend − grant.
@@ -203,30 +205,63 @@ public class PaygWalletController {
noCap,
billing.subscriptionId(),
spend,
- breakdown,
+ breakdowns.units(),
members,
- buildActivity(teamId));
+ buildActivity(teamId),
+ breakdowns.docs(),
+ analytics.docsProcessed(),
+ analytics.uniquePdfs(),
+ analytics.sizeMultiplierPdfs());
return ResponseEntity.ok(body);
}
- private CategoryBreakdown buildBreakdown(
+ /** Per-category size-scaled units + input-file counts for the same window. */
+ private record BreakdownPair(CategoryBreakdown units, CategoryBreakdown docs) {}
+
+ /** Period usage analytics: total input files, unique PDFs, and size-multiplier files. */
+ private record UsageAnalytics(int docsProcessed, int uniquePdfs, int sizeMultiplierPdfs) {}
+
+ private BreakdownPair buildBreakdowns(
Long teamId, LocalDateTime periodStart, LocalDateTime periodEnd) {
- Map byCategory = new HashMap<>();
+ Map units = new HashMap<>();
+ Map docs = new HashMap<>();
for (Object[] row :
- ledgerRepo.sumPeriodAmountByCategory(
+ ledgerRepo.sumPeriodByCategoryWithDocs(
teamId, LedgerEntryType.DEBIT, periodStart, periodEnd)) {
- if (row.length >= 2
- && row[0] instanceof BillingCategory cat
- && row[1] instanceof Number n) {
- byCategory.put(cat, n.longValue());
+ if (row.length >= 3 && row[0] instanceof BillingCategory cat) {
+ if (row[1] instanceof Number u) {
+ units.put(cat, u.longValue());
+ }
+ if (row[2] instanceof Number d) {
+ docs.put(cat, d.longValue());
+ }
}
}
+ return new BreakdownPair(categoryBreakdown(units), categoryBreakdown(docs));
+ }
+
+ private static CategoryBreakdown categoryBreakdown(Map byCategory) {
return new CategoryBreakdown(
clampToInt(byCategory.getOrDefault(BillingCategory.API, 0L)),
clampToInt(byCategory.getOrDefault(BillingCategory.AI, 0L)),
clampToInt(byCategory.getOrDefault(BillingCategory.AUTOMATION, 0L)));
}
+ private UsageAnalytics buildUsageAnalytics(
+ Long teamId, LocalDateTime periodStart, LocalDateTime periodEnd) {
+ List