PAYG B-2: shadow-mode filter + interceptor (engine activation) (#6519)

## What this PR does

Wires the **B-1 shadow charging engine** into real HTTP request flow.
After this lands, flipping an internal team to ``PAYG_SHADOW`` via SQL
begins populating ``payg_shadow_charge`` automatically — with **zero
impact** on the legacy credit deduction path.

**This is the load-bearing PR for shadow mode.** Without it, B-1's
engine sits idle — nothing in the codebase calls
``JobChargeService.openProcess()`` from a real HTTP request.

Stacks on top of #6477 (PR B-1).

## Components

| Class | Role |
|---|---|
| ``PaygResponseBodyWrapperFilter`` | Servlet filter, installs tee'ing
response wrapper. Defers wrapper close to ``AsyncListener`` for
``DeferredResult`` / ``CompletableFuture`` controllers so the lifetime
spans the async window. |
| ``PaygResponseBodyWrapper`` | ``HttpServletResponseWrapper`` —
in-memory ``ByteArrayOutputStream`` up to 10 MiB; spills to ``TempFile``
above. ``materialisedPath()`` always returns a uniform ``Path``
interface. |
| ``PaygChargeInterceptor`` | ``AsyncHandlerInterceptor`` mirroring
``UnifiedCreditInterceptor`` shape. ``preHandle`` gates on
``@AutoJobPostMapping``, materialises multipart inputs, calls
``JobChargeService.openProcess``. ``afterCompletion`` branches on HTTP
status. |
| ``PaygOutputExtractor`` | Pulls PDFs out of the response body. Direct
``application/pdf`` returns body verbatim; ``application/zip`` iterates
entries and keeps each ``.pdf`` entry whose first bytes match the
``%PDF-`` magic. |
| ``PaygWebMvcConfig`` | Registers filter at end of Spring filter chain
(after security); interceptor after ``UnifiedCreditInterceptor``. |
| ``PaygFilterProperties`` | ``payg.filter.enabled`` master switch +
in-memory threshold + optional max-bytes ceiling. |

## Status branching in afterCompletion

| HTTP status | Action |
|---|---|
| **2xx** | Append OK step; extract PDFs from response;
``JobService.recordOutput`` per PDF |
| **4xx** | Append FAILED step with ``errorCode``. No refund — customer
paid for the attempt. No OUTPUT recording. |
| **5xx + OPENED** (first-step) |
``JobChargeService.markFirstStepFailed`` → shadow row flipped to
``REFUNDED``, process CLOSED. Refund counter incremented. |
| **5xx + JOINED** (mid-chain) | ``JobChargeService.decrementStepCount``
— step slot returned without resetting ``lastStepAt`` (workflow window
stays active for retry). |

## New ``JobChargeService`` methods

- **``markFirstStepFailed(jobId, reason)``** — flips shadow row to
``REFUNDED`` with ``refundedAt`` + ``refundReason``, closes the process.
Idempotent. Mimics the eventual Stripe
``meter_event_adjustment(cancel)`` flow that real-mode will invoke at
the same callsite. **Refund implies close** so a same-input retry can't
lineage-join into a refunded chain for free work.
- **``decrementStepCount(jobId)``** — defensive floor at 1; never drives
count negative.

## Schema

- Backend: ``V13__payg_shadow_charge_status.sql`` adds ``status``
(``CHARGED`` | ``REFUNDED``) + ``refunded_at`` + ``refund_reason``.
``DEFAULT 'CHARGED'`` so existing B-1 rows stay correct without
backfill.
- Supabase: matching migration in
[Stirling-PDF-SaaS#payg-shadow-charge-status](https://github.com/Stirling-Tools/Stirling-PDF-SaaS/tree/payg-shadow-charge-status)

## Fail-open semantics in shadow

Any ``RuntimeException`` in ``preHandle`` / ``afterCompletion`` is
logged at WARN, increments ``payg.filter.errors``, and lets the
customer's tool call proceed unbilled. This **reverses to fail-closed**
when ``wallet_policy.engine = PAYG`` (real charging) — that reversal
lives inside ``JobChargeService`` and ships with the cap evaluator PR
(PR-C1 in PAYG_DESIGN.md).

## Observability

Micrometer metrics:
- ``payg.filter.errors`` Counter — internal failures (preHandle +
afterCompletion). Alert source.
- ``payg.filter.calls`` Counter, tagged ``disposition`` (``OPENED`` |
``JOINED`` | ``SHORT_CIRCUIT``)
- ``payg.filter.refunds`` Counter — first-step 5xx refunds
- ``payg.filter.duration`` Timer — preHandle + afterCompletion
wall-clock per request

## Test coverage (38 tests across 4 classes)

- **PaygResponseBodyWrapperTest** (12 tests) — in-memory, spill,
threshold crossing mid-chunk, writer vs outputStream exclusivity,
``resetBuffer`` with and without spill, close idempotency, single-byte
writes across threshold.
- **PaygOutputExtractorTest** (7 tests) — direct PDF, parametrised
content type, ZIP with mixed entries + magic-byte gate, corrupt ZIP
fail-open, empty ZIP.
- **PaygChargeInterceptorTest** (13 tests) — all preHandle
short-circuits, OPENED disposition stash, fail-open on chargeService
exception, 2xx recordOutputs path, 5xx OPENED → markFirstStepFailed, 5xx
JOINED → decrementStepCount, 4xx FAILED step append, max-bytes ceiling
skip, PIPELINE header detection.
- **JobChargeServiceTest extended** (+6 tests) — markFirstStepFailed
happy path, idempotency, missing-shadow-row case, long-reason trim;
decrementStepCount happy path, floor-at-1 defence, missing-job no-op.

## What's NOT in this PR (deliberate)

- **No SpringBootTest layer.** The saas module doesn't have bootstrap
test infrastructure (Supabase JWT config + H2 schema harness).
Integration confidence comes from the manual staging deploy + SQL-flip
of an internal team. Bootstrap-test infra is a focused follow-up if
needed.
- **No saas-mode Behave / docker-compose.** Per design §17 — deferred.
Existing ``testing/cucumber/`` infrastructure doesn't yet have a
saas-profile compose target; that's its own PR when warranted.
- **No CreditService wire-in** (per design §13 decision). Per-row
comparison data moves to the reconciliation report PR (PR-S2).
``legacy_credits_charged`` + ``diff_pct`` columns stay at 0 in shadow
rows.
- **No reconciliation report endpoint.** Direct SQL queries against
``payg_shadow_charge`` cover the data-access need until patterns emerge.

## Rollback levers

| Symptom | Lever |
|---|---|
| Some / all tool calls breaking due to filter |
``payg.filter.enabled=false`` + restart (~20s) |
| Shadow rows look wrong for a specific team | ``UPDATE wallet_policy
SET engine = 'LEGACY' WHERE team_id = ?`` |
| Mass shadow weirdness | ``UPDATE wallet_policy SET engine = 'LEGACY'``
|
| Memory exhaustion from response tee | Lower
``payg.filter.response.in-memory-threshold-bytes`` |

## Test plan

- [ ] CI green (build + tests)
- [ ] Aikido / Snyk / SonarCloud clean
- [ ] Manual: deploy to staging
- [ ] Manual: flip one internal team via ``UPDATE wallet_policy SET
engine = 'PAYG_SHADOW' WHERE team_id = ?``
- [ ] Manual: hit ``/api/v1/security/add-password`` with that team's
JWT; verify a ``payg_shadow_charge`` row appears with
``status='CHARGED'``
- [ ] Manual: trigger a 503 (e.g. via temporary backend kill
mid-request); verify the resulting row is ``status='REFUNDED'`` + the
process is ``CLOSED``
- [ ] Manual: hit ``/api/v1/general/split`` with a multi-page PDF;
verify one OUTPUT signature per inner PDF appears in
``job_artifact_hash``
- [ ] Manual: chain ``add-password`` → ``compress`` on the output;
verify the second call JOINS the first process (no new shadow row) and
the inner output OUTPUT signature is what drove the lineage join

## Stacks on / references

- Stacks on: #6477 (B-1 — shadow charging engine)
- Schema mirror: Stirling-PDF-SaaS#payg-shadow-charge-status branch
- Design doc: ``notes/PAYG_FILTER_DESIGN.md`` (all 19 decisions DECIDED)
This commit is contained in:
ConnorYoh
2026-06-04 15:07:58 +00:00
committed by GitHub
parent 3807cdfbc6
commit 22dacbed01
23 changed files with 2384 additions and 51 deletions
@@ -2,8 +2,11 @@ package stirling.software.saas.payg.charge;
import java.io.IOException;
import java.nio.file.Path;
import java.time.LocalDateTime;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.UUID;
import org.springframework.context.annotation.Profile;
import org.springframework.stereotype.Service;
@@ -17,10 +20,14 @@ import stirling.software.saas.payg.docs.DocumentMetrics;
import stirling.software.saas.payg.job.JobContext;
import stirling.software.saas.payg.job.JobService;
import stirling.software.saas.payg.job.JoinOrOpenResult;
import stirling.software.saas.payg.job.ProcessingJob;
import stirling.software.saas.payg.model.JobSource;
import stirling.software.saas.payg.model.JobStatus;
import stirling.software.saas.payg.model.ShadowChargeStatus;
import stirling.software.saas.payg.policy.PricingPolicy;
import stirling.software.saas.payg.policy.PricingPolicyService;
import stirling.software.saas.payg.repository.PaygShadowChargeRepository;
import stirling.software.saas.payg.repository.ProcessingJobRepository;
import stirling.software.saas.payg.shadow.PaygShadowCharge;
/**
@@ -47,16 +54,19 @@ public class JobChargeService {
private final PricingPolicyService policyService;
private final DocumentClassifier classifier;
private final PaygShadowChargeRepository shadowRepository;
private final ProcessingJobRepository jobRepository;
public JobChargeService(
JobService jobService,
PricingPolicyService policyService,
DocumentClassifier classifier,
PaygShadowChargeRepository shadowRepository) {
PaygShadowChargeRepository shadowRepository,
ProcessingJobRepository jobRepository) {
this.jobService = Objects.requireNonNull(jobService, "jobService");
this.policyService = Objects.requireNonNull(policyService, "policyService");
this.classifier = Objects.requireNonNull(classifier, "classifier");
this.shadowRepository = Objects.requireNonNull(shadowRepository, "shadowRepository");
this.jobRepository = Objects.requireNonNull(jobRepository, "jobRepository");
}
/**
@@ -117,10 +127,14 @@ public class JobChargeService {
private int computeUnits(List<JobInput> inputs, PricingPolicy policy) {
List<MultipartFile> multiparts = inputs.stream().map(JobInput::multipart).toList();
// Reuse the temp file the caller already wrote (in PaygChargeInterceptor.preHandle for
// lineage hashing) instead of materialising the same bytes a second time inside the
// classifier. Saves one write + one read per PDF input.
List<Path> paths = inputs.stream().map(JobInput::path).toList();
DocumentMetrics metrics =
multiparts.size() == 1
? classifier.classify(multiparts.get(0), policy)
: classifier.classify(multiparts, policy);
? classifier.classify(multiparts.get(0), paths.get(0), policy)
: classifier.classify(multiparts, paths, policy);
// Apply the policy-level minChargeUnits floor per design § 3.4. The classifier returns
// raw docUnits with a "non-empty input → ≥1" floor; the charge formula's
// max(min_charge_units, docUnits) layers on top.
@@ -138,6 +152,83 @@ public class JobChargeService {
// CreditService in the follow-up PR. Until then, diff stays at 0.
row.setLegacyCreditsCharged(0);
row.setDiffPct(0);
row.setStatus(ShadowChargeStatus.CHARGED);
shadowRepository.save(row);
}
/**
* First-step failure on a freshly-opened process: mimic a successful Stripe
* meter_event_adjustment(cancel) by flipping the shadow row to {@link
* ShadowChargeStatus#REFUNDED}, and close the process so a same-input retry can't lineage-join
* into a refunded chain for free work.
*
* <p>Idempotent: re-invoking on an already-REFUNDED row or already-CLOSED process is a silent
* no-op.
*/
@Transactional
public void markFirstStepFailed(UUID jobId, String refundReason) {
Objects.requireNonNull(jobId, "jobId");
LocalDateTime now = LocalDateTime.now();
Optional<PaygShadowCharge> rowOpt = shadowRepository.findFirstByJobIdOrderByIdAsc(jobId);
if (rowOpt.isEmpty()) {
log.debug(
"markFirstStepFailed: no shadow row for job {} (PAYG not active for team?)",
jobId);
} else {
PaygShadowCharge row = rowOpt.get();
if (row.getStatus() != ShadowChargeStatus.REFUNDED) {
row.setStatus(ShadowChargeStatus.REFUNDED);
row.setRefundedAt(now);
row.setRefundReason(trimReason(refundReason));
shadowRepository.save(row);
}
}
ProcessingJob job = jobRepository.findById(jobId).orElse(null);
if (job == null) {
log.warn("markFirstStepFailed: no ProcessingJob with id {}", jobId);
return;
}
if (job.getStatus() == JobStatus.OPEN) {
job.setStatus(JobStatus.CLOSED);
job.setClosedAt(now);
jobRepository.save(job);
}
}
/**
* Mid-chain 5xx on a JOINED step: return the step slot. The {@code lastStepAt} timestamp stays
* advanced (workflow window intentionally remains active for the next retry). No shadow-row
* change — only OPENED-disposition calls wrote a row.
*
* <p>Defensive lower bound: never drives {@code stepCount} below 1 (which would imply we'd
* decremented an already-decremented slot, or were called on a fresh process).
*/
@Transactional
public void decrementStepCount(UUID jobId) {
Objects.requireNonNull(jobId, "jobId");
ProcessingJob job = jobRepository.findById(jobId).orElse(null);
if (job == null) {
log.warn("decrementStepCount: no ProcessingJob with id {}", jobId);
return;
}
int current = job.getStepCount() == null ? 0 : job.getStepCount();
if (current <= 1) {
log.debug(
"decrementStepCount: stepCount already at {} for job {}; no-op",
current,
jobId);
return;
}
job.setStepCount(current - 1);
jobRepository.save(job);
}
private static String trimReason(String reason) {
if (reason == null) {
return null;
}
return reason.length() > 128 ? reason.substring(0, 128) : reason;
}
}
@@ -4,6 +4,7 @@ import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.List;
import java.util.Objects;
@@ -48,10 +49,16 @@ public class DefaultDocumentClassifier implements DocumentClassifier {
@Override
public DocumentMetrics classify(MultipartFile file, PricingPolicy policy) {
return classify(file, null, policy);
}
@Override
public DocumentMetrics classify(
MultipartFile file, Path materialisedPath, PricingPolicy policy) {
Objects.requireNonNull(file, "file");
Objects.requireNonNull(policy, "policy");
FileFacts facts = inspect(file);
FileFacts facts = inspect(file, materialisedPath);
long rawUnits = computeRawUnits(facts.pages, facts.bytes, policy);
// toIntExact: fail loud on overflow rather than silently wrapping a billing number.
int units =
@@ -64,19 +71,35 @@ public class DefaultDocumentClassifier implements DocumentClassifier {
@Override
public DocumentMetrics classify(List<MultipartFile> files, PricingPolicy policy) {
return classify(files, null, policy);
}
@Override
public DocumentMetrics classify(
List<MultipartFile> files, List<Path> materialisedPaths, PricingPolicy policy) {
Objects.requireNonNull(files, "files");
Objects.requireNonNull(policy, "policy");
if (files.isEmpty()) {
throw new IllegalArgumentException("files must not be empty");
}
if (materialisedPaths != null && materialisedPaths.size() != files.size()) {
throw new IllegalArgumentException(
"materialisedPaths size ("
+ materialisedPaths.size()
+ ") must equal files size ("
+ files.size()
+ ")");
}
int totalPages = 0;
long totalBytes = 0;
long rawUnitsSum = 0;
String firstContentType = null;
for (MultipartFile file : files) {
FileFacts facts = inspect(file);
for (int i = 0; i < files.size(); i++) {
MultipartFile file = files.get(i);
Path path = materialisedPaths == null ? null : materialisedPaths.get(i);
FileFacts facts = inspect(file, path);
// Sum the *raw* (unclamped) per-file units so the group cap below can actually bind.
// Per-file clamping in this loop would make the group cap a no-op.
rawUnitsSum =
@@ -103,11 +126,17 @@ public class DefaultDocumentClassifier implements DocumentClassifier {
totalUnits);
}
private FileFacts inspect(MultipartFile file) {
private FileFacts inspect(MultipartFile file, Path materialisedPath) {
long bytes = file.getSize();
String contentType =
file.getContentType() != null ? file.getContentType() : DEFAULT_CONTENT_TYPE;
int pages = isPdf(contentType, file.getOriginalFilename()) ? readPageCount(file) : 0;
int pages = 0;
if (isPdf(contentType, file.getOriginalFilename())) {
pages =
materialisedPath != null
? readPageCountFromPath(materialisedPath, file.getOriginalFilename())
: readPageCount(file);
}
return new FileFacts(pages, bytes, contentType);
}
@@ -141,9 +170,7 @@ public class DefaultDocumentClassifier implements DocumentClassifier {
OutputStream out = Files.newOutputStream(temp.getPath())) {
in.transferTo(out);
}
try (PdfDocument doc = PdfDocument.open(temp.getPath())) {
return doc.pageCount();
}
return readPageCountFromPath(temp.getPath(), file.getOriginalFilename());
} catch (IOException | RuntimeException e) {
log.debug(
"Could not read PDF page count for {} ({}); falling back to bytes-only units",
@@ -153,6 +180,23 @@ public class DefaultDocumentClassifier implements DocumentClassifier {
}
}
/**
* Page-count read against an already-materialised file. Used by callers that already wrote the
* bytes to disk (the PAYG interceptor materialises every input for the lineage hash) so we
* avoid a second copy.
*/
private int readPageCountFromPath(Path path, String displayName) {
try (PdfDocument doc = PdfDocument.open(path)) {
return doc.pageCount();
} catch (RuntimeException e) {
log.debug(
"Could not read PDF page count for {} ({}); falling back to bytes-only units",
displayName,
e.getClass().getSimpleName());
return 0;
}
}
private static int saturatedAdd(int a, int b) {
long sum = (long) a + b;
if (sum > Integer.MAX_VALUE) {
@@ -1,5 +1,6 @@
package stirling.software.saas.payg.docs;
import java.nio.file.Path;
import java.util.List;
import org.springframework.web.multipart.MultipartFile;
@@ -11,6 +12,17 @@ import stirling.software.saas.payg.policy.PricingPolicy;
*
* <p>Returns {@code docUnits} with an absolute floor of 1 for non-empty input. {@code
* policy.minChargeUnits} is applied at charge time, not here.
*
* <p>Each overload comes in two flavours:
*
* <ul>
* <li>{@code classify(MultipartFile, ...)} — classifier reads bytes via {@code getInputStream()}
* and writes its own temp file to feed jpdfium (for PDFs). Use when no on-disk copy exists.
* <li>{@code classify(MultipartFile, Path, ...)} — caller has already materialised the bytes to
* {@code Path}; classifier reads page count directly from there without re-writing. Hot-path
* callers (the PAYG interceptor) should use this form since they materialise inputs anyway
* for the lineage hash.
* </ul>
*/
public interface DocumentClassifier {
@@ -22,4 +34,20 @@ public interface DocumentClassifier {
* units, capped at {@code fileUnitCap × files.size()} and floored at 1.
*/
DocumentMetrics classify(List<MultipartFile> files, PricingPolicy policy);
/**
* Same as {@link #classify(MultipartFile, PricingPolicy)} but uses {@code materialisedPath} for
* PDF page-count extraction, avoiding a second copy of the upload bytes to disk. Callers that
* hold an on-disk copy (e.g. the PAYG interceptor materialises every input for the lineage
* hash) should prefer this form.
*/
DocumentMetrics classify(MultipartFile file, Path materialisedPath, PricingPolicy policy);
/**
* Multi-file variant of {@link #classify(MultipartFile, Path, PricingPolicy)}. {@code
* materialisedPaths} must align positionally with {@code files} — entry {@code i} is the
* on-disk copy of {@code files.get(i)}.
*/
DocumentMetrics classify(
List<MultipartFile> files, List<Path> materialisedPaths, PricingPolicy policy);
}
@@ -0,0 +1,420 @@
package stirling.software.saas.payg.filter;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.UUID;
import org.springframework.context.annotation.Profile;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.stereotype.Component;
import org.springframework.util.MultiValueMap;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.web.multipart.MultipartHttpServletRequest;
import org.springframework.web.servlet.AsyncHandlerInterceptor;
import io.micrometer.core.instrument.Counter;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.Timer;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import stirling.software.common.annotations.AutoJobPostMapping;
import stirling.software.common.util.TempFile;
import stirling.software.common.util.TempFileManager;
import stirling.software.proprietary.security.database.repository.UserRepository;
import stirling.software.proprietary.security.model.ApiKeyAuthenticationToken;
import stirling.software.proprietary.security.model.User;
import stirling.software.saas.payg.charge.ChargeContext;
import stirling.software.saas.payg.charge.ChargeOutcome;
import stirling.software.saas.payg.charge.JobChargeService;
import stirling.software.saas.payg.charge.JobInput;
import stirling.software.saas.payg.job.JobService;
import stirling.software.saas.payg.model.JobSource;
import stirling.software.saas.payg.model.JobStepStatus;
import stirling.software.saas.payg.model.ProcessType;
import stirling.software.saas.util.AuthenticationUtils;
/**
* The hot-path PAYG interceptor. Mirrors the {@code UnifiedCreditInterceptor} shape: registered
* after it in {@code PaygWebMvcConfig} so legacy credit-rejection short-circuits before we waste
* work hashing inputs.
*
* <p>{@code preHandle}: gates on {@code @AutoJobPostMapping}, reads the parsed multipart parts,
* materialises each input to a {@code TempFile}, and asks {@link JobChargeService#openProcess} to
* open (or join) a process. The resulting {@link ChargeOutcome} plus input temp-files are stashed
* as request attributes for {@code afterCompletion}.
*
* <p>{@code afterCompletion}: branches on HTTP status — 2xx hashes the response body for OUTPUT
* lineage; 4xx records a step append for audit; 5xx triggers refund-and-close (OPENED) or
* step-quota return (JOINED). Closes all input temp files and the response wrapper at the end.
*
* <p>Fail-open everywhere: any unexpected {@link RuntimeException} is swallowed, logged at WARN,
* and counted on {@code payg.filter.errors}. The customer's tool call always proceeds.
*
* <p>Async controllers ({@code DeferredResult}, {@code CompletableFuture}) are handled
* transparently — {@link AsyncHandlerInterceptor#afterConcurrentHandlingStarted} is a no-op; the
* normal {@code afterCompletion} fires when the async work resolves.
*/
@Slf4j
@Component
@Profile("saas")
public class PaygChargeInterceptor implements AsyncHandlerInterceptor {
static final String ATTR_JOB_ID = PaygChargeInterceptor.class.getName() + ".JOB_ID";
static final String ATTR_DISPOSITION = PaygChargeInterceptor.class.getName() + ".DISPOSITION";
static final String ATTR_INPUT_TEMP_FILES =
PaygChargeInterceptor.class.getName() + ".INPUT_TEMP_FILES";
static final String ATTR_INPUT_BYTES = PaygChargeInterceptor.class.getName() + ".INPUT_BYTES";
static final String ATTR_FAILED = PaygChargeInterceptor.class.getName() + ".FAILED";
static final String ATTR_TOOL_ID = PaygChargeInterceptor.class.getName() + ".TOOL_ID";
private static final String AUTOMATION_HEADER = "X-Stirling-Automation";
private final JobChargeService chargeService;
private final JobService jobService;
private final UserRepository userRepository;
private final TempFileManager tempFileManager;
private final PaygOutputExtractor outputExtractor;
private final PaygFilterProperties properties;
private final Counter errorsCounter;
private final Counter callsOpened;
private final Counter callsJoined;
private final Counter callsShortCircuit;
private final Counter refundsCounter;
private final Timer durationTimer;
public PaygChargeInterceptor(
JobChargeService chargeService,
JobService jobService,
UserRepository userRepository,
TempFileManager tempFileManager,
PaygOutputExtractor outputExtractor,
PaygFilterProperties properties,
MeterRegistry meterRegistry) {
this.chargeService = chargeService;
this.jobService = jobService;
this.userRepository = userRepository;
this.tempFileManager = tempFileManager;
this.outputExtractor = outputExtractor;
this.properties = properties;
this.errorsCounter =
Counter.builder("payg.filter.errors")
.description("PAYG interceptor / filter internal failures")
.register(meterRegistry);
this.callsOpened =
Counter.builder("payg.filter.calls")
.tag("disposition", "OPENED")
.register(meterRegistry);
this.callsJoined =
Counter.builder("payg.filter.calls")
.tag("disposition", "JOINED")
.register(meterRegistry);
this.callsShortCircuit =
Counter.builder("payg.filter.calls")
.tag("disposition", "SHORT_CIRCUIT")
.register(meterRegistry);
this.refundsCounter =
Counter.builder("payg.filter.refunds")
.description("First-step 5xx refunds applied to shadow rows")
.register(meterRegistry);
this.durationTimer =
Timer.builder("payg.filter.duration")
.description("preHandle + afterCompletion wall-clock per request")
.register(meterRegistry);
}
@Override
public boolean preHandle(
HttpServletRequest request, HttpServletResponse response, Object handler) {
Timer.Sample sample = Timer.start();
try {
if (!properties.isEnabled()) {
return true;
}
if (!(handler instanceof HandlerMethod hm)
|| hm.getMethodAnnotation(AutoJobPostMapping.class) == null) {
callsShortCircuit.increment();
return true;
}
try {
doPreHandle(request);
} catch (RuntimeException e) {
log.warn("PAYG preHandle failed; passing through unbilled", e);
errorsCounter.increment();
request.setAttribute(ATTR_FAILED, Boolean.TRUE);
cleanupInputs(request);
}
return true;
} finally {
sample.stop(durationTimer);
}
}
private void doPreHandle(HttpServletRequest request) {
Authentication auth = SecurityContextHolder.getContext().getAuthentication();
User currentUser = resolveUser(auth);
if (currentUser == null) {
callsShortCircuit.increment();
return;
}
if (!(request instanceof MultipartHttpServletRequest mreq)) {
callsShortCircuit.increment();
return;
}
MultiValueMap<String, MultipartFile> map = mreq.getMultiFileMap();
List<MultipartFile> nonEmpty = new ArrayList<>();
for (List<MultipartFile> bucket : map.values()) {
for (MultipartFile mp : bucket) {
if (mp.getSize() > 0) {
nonEmpty.add(mp);
}
}
}
if (nonEmpty.isEmpty()) {
callsShortCircuit.increment();
return;
}
List<TempFile> tempFiles = new ArrayList<>(nonEmpty.size());
List<JobInput> inputs = new ArrayList<>(nonEmpty.size());
long totalInputBytes = 0L;
try {
for (MultipartFile mp : nonEmpty) {
TempFile tf = tempFileManager.createManagedTempFile(".upload");
tempFiles.add(tf);
try (InputStream in = mp.getInputStream();
OutputStream out = Files.newOutputStream(tf.getPath())) {
in.transferTo(out);
}
inputs.add(new JobInput(mp, tf.getPath()));
totalInputBytes += mp.getSize();
}
} catch (IOException e) {
for (TempFile tf : tempFiles) {
tf.close();
}
throw new RuntimeException("Failed to materialise multipart input", e);
}
// Publish as an immutable list. preHandle fully populates `tempFiles` before this
// setAttribute, and cleanupInputs / afterCompletion are read-only consumers — exposing
// an unmodifiable view (a) makes the safe-publication guarantee from the container's
// attribute-map synchronization unambiguous to readers on different threads (sync vs
// async-dispatch), and (b) hard-fails any future caller that tries to mutate it.
request.setAttribute(ATTR_INPUT_TEMP_FILES, Collections.unmodifiableList(tempFiles));
request.setAttribute(ATTR_INPUT_BYTES, totalInputBytes);
request.setAttribute(ATTR_TOOL_ID, request.getRequestURI());
ChargeContext ctx =
new ChargeContext(
currentUser.getId(),
currentUser.getTeam() == null ? null : currentUser.getTeam().getId(),
determineSource(request, auth),
ProcessType.SINGLE_TOOL);
ChargeOutcome outcome;
try {
outcome = chargeService.openProcess(ctx, inputs);
} catch (IOException e) {
throw new RuntimeException("openProcess IO failure", e);
}
request.setAttribute(ATTR_JOB_ID, outcome.processId());
request.setAttribute(ATTR_DISPOSITION, outcome.disposition());
if (outcome.disposition() == ChargeOutcome.Disposition.OPENED) {
callsOpened.increment();
} else {
callsJoined.increment();
}
}
@Override
public void afterCompletion(
HttpServletRequest request,
HttpServletResponse response,
Object handler,
Exception ex) {
Timer.Sample sample = Timer.start();
try {
if (Boolean.TRUE.equals(request.getAttribute(ATTR_FAILED))) {
cleanupInputs(request);
closeWrapper(request);
return;
}
UUID jobId = (UUID) request.getAttribute(ATTR_JOB_ID);
if (jobId == null) {
closeWrapper(request);
return;
}
try {
doAfterCompletion(request, response, jobId);
} catch (RuntimeException e) {
log.warn(
"PAYG afterCompletion failed for job {}; lineage may be incomplete",
jobId,
e);
errorsCounter.increment();
} finally {
cleanupInputs(request);
closeWrapper(request);
}
} finally {
sample.stop(durationTimer);
}
}
private void doAfterCompletion(
HttpServletRequest request, HttpServletResponse response, UUID jobId) {
int status = response.getStatus();
ChargeOutcome.Disposition disposition =
(ChargeOutcome.Disposition) request.getAttribute(ATTR_DISPOSITION);
String toolId =
Optional.ofNullable((String) request.getAttribute(ATTR_TOOL_ID)).orElse("unknown");
Long inputBytes = (Long) request.getAttribute(ATTR_INPUT_BYTES);
// Step audit row — appended for every disposition + every outcome class. Done first so a
// refund-and-close still has the failure recorded against the now-CLOSED process.
JobStepStatus stepStatus = status < 400 ? JobStepStatus.OK : JobStepStatus.FAILED;
String errorCode = status >= 400 ? String.valueOf(status) : null;
try {
jobService.appendStep(jobId, toolId, stepStatus, null, inputBytes, errorCode);
} catch (RuntimeException e) {
log.debug("appendStep failed for job {}: {}", jobId, e.getMessage());
}
if (status >= 500) {
if (disposition == ChargeOutcome.Disposition.OPENED) {
chargeService.markFirstStepFailed(jobId, "first-step-5xx:" + status);
refundsCounter.increment();
} else {
chargeService.decrementStepCount(jobId);
}
return;
}
if (status >= 400) {
// 4xx: customer paid for the attempt. No OUTPUT recording, no refund.
return;
}
recordOutputs(request, response, jobId);
}
private void recordOutputs(
HttpServletRequest request, HttpServletResponse response, UUID jobId) {
PaygResponseBodyWrapper wrapper =
(PaygResponseBodyWrapper)
request.getAttribute(PaygResponseBodyWrapperFilter.REQUEST_ATTRIBUTE);
if (wrapper == null) {
return;
}
Long maxBytes = properties.getResponse().getMaxBytes();
if (maxBytes != null && wrapper.bytesWritten() > maxBytes) {
log.debug(
"Response size {} exceeds payg.filter.response.max-bytes={}; skipping OUTPUT recording",
wrapper.bytesWritten(),
maxBytes);
return;
}
Path bodyPath;
try {
bodyPath = wrapper.materialisedPath();
} catch (IOException e) {
log.debug("materialisedPath failed for job {}: {}", jobId, e.getMessage());
return;
}
if (bodyPath == null) {
return;
}
List<PaygOutputExtractor.ExtractedPdf> pdfs =
outputExtractor.extract(response.getContentType(), bodyPath);
try {
for (PaygOutputExtractor.ExtractedPdf pdf : pdfs) {
try {
jobService.recordOutput(jobId, pdf.path());
} catch (IOException e) {
log.debug(
"recordOutput failed for job {} path {}: {}",
jobId,
pdf.path(),
e.getMessage());
}
}
} finally {
for (PaygOutputExtractor.ExtractedPdf pdf : pdfs) {
pdf.close();
}
}
}
@Override
public void afterConcurrentHandlingStarted(
HttpServletRequest request, HttpServletResponse response, Object handler) {
// Async handoff: don't touch state. afterCompletion will fire when the async work resolves.
}
@SuppressWarnings("unchecked")
private void cleanupInputs(HttpServletRequest request) {
Object raw = request.getAttribute(ATTR_INPUT_TEMP_FILES);
if (raw instanceof List<?>) {
for (Object o : (List<Object>) raw) {
if (o instanceof TempFile tf) {
tf.close();
}
}
}
request.removeAttribute(ATTR_INPUT_TEMP_FILES);
}
private void closeWrapper(HttpServletRequest request) {
Object wrapper = request.getAttribute(PaygResponseBodyWrapperFilter.REQUEST_ATTRIBUTE);
if (wrapper instanceof PaygResponseBodyWrapper w) {
w.close();
}
}
private User resolveUser(Authentication auth) {
if (auth == null || !auth.isAuthenticated()) {
return null;
}
if (auth instanceof ApiKeyAuthenticationToken && auth.getPrincipal() instanceof User u) {
return u;
}
try {
String supabaseId = AuthenticationUtils.extractSupabaseId(auth);
if (supabaseId == null) {
return null;
}
UUID supabaseUuid = UUID.fromString(supabaseId);
return userRepository.findBySupabaseId(supabaseUuid).orElse(null);
} catch (RuntimeException e) {
log.debug("PAYG resolveUser failed: {}", e.getMessage());
return null;
}
}
private static JobSource determineSource(HttpServletRequest request, Authentication auth) {
String automationHeader = request.getHeader(AUTOMATION_HEADER);
if (automationHeader != null && "true".equalsIgnoreCase(automationHeader.trim())) {
return JobSource.PIPELINE;
}
if (auth instanceof ApiKeyAuthenticationToken) {
return JobSource.API;
}
return JobSource.WEB;
}
}
@@ -0,0 +1,44 @@
package stirling.software.saas.payg.filter;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.annotation.Profile;
import org.springframework.stereotype.Component;
import lombok.Getter;
import lombok.Setter;
/**
* Configuration knobs for the PAYG filter + interceptor stack. See {@code PAYG_FILTER_DESIGN.md}
* §16 for the rationale on each default.
*
* <ul>
* <li>{@code payg.filter.enabled} — master kill switch. Restart-required; no
* {@code @RefreshScope}. When {@code false}, the wrapper filter passes through and the
* interceptor short-circuits in {@code preHandle}.
* <li>{@code payg.filter.response.in-memory-threshold-bytes} — the wrapper buffers below this in
* a {@link java.io.ByteArrayOutputStream}; above it spills to a {@code TempFile}.
* <li>{@code payg.filter.response.max-bytes} — optional ceiling. Responses exceeding this skip
* OUTPUT recording in {@code afterCompletion}. {@code null} = unbounded.
* </ul>
*/
@Component
@Profile("saas")
@ConfigurationProperties(prefix = "payg.filter")
@Getter
@Setter
public class PaygFilterProperties {
private boolean enabled = true;
private final Response response = new Response();
@Getter
@Setter
public static class Response {
/** 10 MiB. Tiny responses stay in RAM; large responses spill. */
private long inMemoryThresholdBytes = 10L * 1024L * 1024L;
/** Optional ceiling: when set, OUTPUT recording is skipped past this size. */
private Long maxBytes;
}
}
@@ -0,0 +1,168 @@
package stirling.software.saas.payg.filter;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardCopyOption;
import java.util.ArrayList;
import java.util.List;
import java.util.Locale;
import java.util.Objects;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;
import org.springframework.context.annotation.Profile;
import org.springframework.stereotype.Component;
import lombok.extern.slf4j.Slf4j;
import stirling.software.common.util.TempFile;
import stirling.software.common.util.TempFileManager;
/**
* Extracts the PDF artefacts from a controller's response body for lineage OUTPUT recording. Two
* content-type paths:
*
* <ul>
* <li>{@code application/pdf} — return the response body path verbatim (the wrapper already
* materialised it to disk).
* <li>{@code application/zip} — iterate entries; each that has a {@code .pdf} extension AND
* starts with the {@code %PDF-} magic-byte sequence becomes one extracted {@link Path}.
* </ul>
*
* <p>Anything else yields an empty list — non-PDF responses don't drive lineage matching. Encrypted
* ZIPs, malformed entries, and IO failures log at debug and yield an empty list (per design §15
* fail-open). Nested ZIP-of-ZIPs is intentionally out of scope: only outer-level PDFs are
* extracted.
*
* <p>Each extracted ZIP entry is written to its own {@link TempFile} returned in the result list.
* Callers are responsible for closing those temp files — typically the interceptor closes them in
* the same {@code afterCompletion} that calls this extractor. The whole-response body path is owned
* by the wrapper and is NOT in the returned list when extraction takes the direct-PDF path.
*/
@Slf4j
@Component
@Profile("saas")
public class PaygOutputExtractor {
/** {@code %PDF-} in ASCII. */
private static final byte[] PDF_MAGIC = {0x25, 0x50, 0x44, 0x46, 0x2D};
private static final String ZIP_CONTENT_TYPE = "application/zip";
private static final String PDF_CONTENT_TYPE = "application/pdf";
private final TempFileManager tempFileManager;
public PaygOutputExtractor(TempFileManager tempFileManager) {
this.tempFileManager = Objects.requireNonNull(tempFileManager, "tempFileManager");
}
/**
* Extract PDF paths from {@code bodyPath} according to {@code contentType}. The returned list
* carries {@link ExtractedPdf} records — each wraps a {@link Path} and an indicator of whether
* the caller owns the temp file lifecycle.
*
* @param contentType the response Content-Type header (may be null / parametrised — only the
* base media type is inspected)
* @param bodyPath the on-disk full response body (from the wrapper's {@code materialisedPath})
* @return zero-or-more PDFs extracted from the body. Empty when content type is not PDF/ZIP,
* when extraction failed, when no entries matched, or when {@code bodyPath} is null.
*/
public List<ExtractedPdf> extract(String contentType, Path bodyPath) {
if (bodyPath == null) {
return List.of();
}
String mediaType = stripParameters(contentType);
if (PDF_CONTENT_TYPE.equalsIgnoreCase(mediaType)) {
// Wrapper-owned path. Don't claim ownership.
return List.of(new ExtractedPdf(bodyPath, null));
}
if (ZIP_CONTENT_TYPE.equalsIgnoreCase(mediaType)) {
return extractZip(bodyPath);
}
return List.of();
}
private List<ExtractedPdf> extractZip(Path bodyPath) {
List<ExtractedPdf> results = new ArrayList<>();
try (InputStream rawIn = Files.newInputStream(bodyPath);
ZipInputStream zin = new ZipInputStream(rawIn)) {
ZipEntry entry;
while ((entry = zin.getNextEntry()) != null) {
try {
if (entry.isDirectory()) {
continue;
}
String name = entry.getName();
if (name == null || !name.toLowerCase(Locale.ROOT).endsWith(".pdf")) {
continue;
}
TempFile temp = tempFileManager.createManagedTempFile(".pdf");
Files.copy(zin, temp.getPath(), StandardCopyOption.REPLACE_EXISTING);
if (isPdfMagic(temp.getPath())) {
results.add(new ExtractedPdf(temp.getPath(), temp));
} else {
temp.close();
}
} finally {
zin.closeEntry();
}
}
} catch (IOException | IllegalArgumentException e) {
// ZipException extends IOException; IllegalArgumentException covers ZipEntry
// "MALFORMED" surfaces from zlib. Fail-open: caller still serves the response.
log.debug(
"ZIP unpack failed for response body {} ({}); skipping per-PDF OUTPUT recording",
bodyPath,
e.getClass().getSimpleName());
// Close anything we already opened.
for (ExtractedPdf p : results) {
p.close();
}
return List.of();
}
return results;
}
private boolean isPdfMagic(Path path) {
byte[] head = new byte[PDF_MAGIC.length];
try (InputStream in = Files.newInputStream(path)) {
int read = in.read(head);
if (read != PDF_MAGIC.length) {
return false;
}
for (int i = 0; i < PDF_MAGIC.length; i++) {
if (head[i] != PDF_MAGIC[i]) {
return false;
}
}
return true;
} catch (IOException e) {
log.debug("PDF magic-byte check failed for {}", path, e);
return false;
}
}
private static String stripParameters(String contentType) {
if (contentType == null) {
return null;
}
int semi = contentType.indexOf(';');
return (semi < 0 ? contentType : contentType.substring(0, semi)).trim();
}
/**
* One PDF extracted from the response body. {@link #ownedTempFile} is non-null only when this
* extractor created the temp file (ZIP entries); the {@link #path} for direct-PDF responses
* remains owned by the wrapper.
*/
public record ExtractedPdf(Path path, TempFile ownedTempFile) implements AutoCloseable {
@Override
public void close() {
if (ownedTempFile != null) {
ownedTempFile.close();
}
}
}
}
@@ -0,0 +1,289 @@
package stirling.software.saas.payg.filter;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Objects;
import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.WriteListener;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.servlet.http.HttpServletResponseWrapper;
import lombok.extern.slf4j.Slf4j;
import stirling.software.common.util.TempFile;
import stirling.software.common.util.TempFileManager;
/**
* Tees the controller's response body so the PAYG interceptor can hash it for OUTPUT lineage
* recording. Writes still flow through to the real client output stream unmodified — the wrapper
* just keeps a parallel copy.
*
* <p>Memory model: bytes accumulate in an in-memory {@link ByteArrayOutputStream} until the
* configured {@code in-memory-threshold-bytes} is crossed; after that the wrapper spills the
* existing buffer plus all subsequent writes to a {@link TempFile} owned by {@link
* TempFileManager}. Tiny responses (error JSONs, small PDFs) stay entirely in RAM with zero disk
* IO; large responses (split-to-ZIP, big compressed PDFs) spill cleanly. See design doc §8.
*
* <p>The interceptor calls {@link #materialisedPath()} from {@code afterCompletion} to get a {@code
* Path} suitable for the lineage detector + ZIP unpack. The wrapper guarantees a Path even if the
* response stayed in memory — it materialises the buffer to a {@link TempFile} on demand so callers
* have a uniform file-based interface.
*
* <p>{@link #close()} closes any {@link TempFile} the wrapper created. Callers MUST invoke close in
* a finally — typically the interceptor's {@code afterCompletion} after it's done hashing.
*
* <p><b>Thread safety:</b> all mutating methods (record paths, materialisedPath, close,
* resetBuffer) synchronize on the wrapper instance. The Servlet spec serialises controller writes
* onto a single dispatch thread, but the {@link jakarta.servlet.AsyncListener#onComplete} callback
* that closes the wrapper for async controllers runs on a container thread distinct from the
* dispatch thread that produced the body. The synchronization makes that handoff safe and also
* guards against future callers (e.g. tests) that might invoke {@link #materialisedPath} or {@link
* #close} from a non-dispatch thread.
*/
@Slf4j
public class PaygResponseBodyWrapper extends HttpServletResponseWrapper implements AutoCloseable {
private final TempFileManager tempFileManager;
private final long inMemoryThresholdBytes;
/** Lazily-created on the first getOutputStream() / getWriter() call. */
private TeeingServletOutputStream teeOut;
private PrintWriter writer;
/**
* In-memory accumulator until the threshold is crossed. Becomes null after spill (helps GC of
* potentially large buffers).
*/
private ByteArrayOutputStream memoryBuffer = new ByteArrayOutputStream();
/** Non-null once we've spilled. Owns the {@link TempFile} below. */
private OutputStream spillStream;
/** Non-null once we've spilled, OR once {@link #materialisedPath()} forced materialisation. */
private TempFile spillFile;
private long bytesWritten;
private boolean spilled;
public PaygResponseBodyWrapper(
HttpServletResponse response,
TempFileManager tempFileManager,
long inMemoryThresholdBytes) {
super(response);
this.tempFileManager = Objects.requireNonNull(tempFileManager, "tempFileManager");
if (inMemoryThresholdBytes < 0) {
throw new IllegalArgumentException(
"inMemoryThresholdBytes must be >= 0, got " + inMemoryThresholdBytes);
}
this.inMemoryThresholdBytes = inMemoryThresholdBytes;
}
@Override
public ServletOutputStream getOutputStream() throws IOException {
if (writer != null) {
// Servlet spec: getOutputStream() and getWriter() are mutually exclusive per request.
throw new IllegalStateException(
"getWriter() was already called on this response; cannot switch to getOutputStream()");
}
if (teeOut == null) {
teeOut = new TeeingServletOutputStream(super.getOutputStream());
}
return teeOut;
}
@Override
public PrintWriter getWriter() throws IOException {
if (teeOut != null) {
throw new IllegalStateException(
"getOutputStream() was already called on this response; cannot switch to getWriter()");
}
if (writer == null) {
String encoding = getCharacterEncoding() != null ? getCharacterEncoding() : "UTF-8";
// Wrap super.getOutputStream() directly with our TeeingServletOutputStream so writes
// through the Writer path also get tee'd. The Writer just adds character→byte encoding.
teeOut = new TeeingServletOutputStream(super.getOutputStream());
writer = new PrintWriter(new OutputStreamWriter(teeOut, encoding));
}
return writer;
}
@Override
public synchronized void resetBuffer() {
super.resetBuffer();
if (memoryBuffer != null) {
memoryBuffer.reset();
}
// If we'd already spilled, the only safe move is to abandon the spill file: the client
// hasn't seen the body yet (otherwise resetBuffer would be illegal), but our tee captured
// bytes we now want to forget.
if (spilled) {
closeSpillQuietly();
spillFile = null;
spillStream = null;
spilled = false;
memoryBuffer = new ByteArrayOutputStream();
}
bytesWritten = 0;
}
/**
* Returns a {@link Path} containing the full response body, or {@code null} if no bytes were
* written. The returned path is owned by this wrapper — do NOT delete or modify it. Use {@link
* #close()} to release.
*
* <p>If the response stayed under the threshold, this materialises the in-memory buffer to a
* {@link TempFile} on demand so the caller always gets a file-based handle (uniform with the
* spilled path).
*/
public synchronized Path materialisedPath() throws IOException {
if (bytesWritten == 0) {
return null;
}
// Flush the writer so any character data lands in the underlying byte stream first.
if (writer != null) {
writer.flush();
}
if (spilled) {
spillStream.flush();
return spillFile.getPath();
}
// Stayed in memory — materialise on demand for a uniform Path-based interface.
if (spillFile == null) {
spillFile = tempFileManager.createManagedTempFile(".body");
Files.write(spillFile.getPath(), memoryBuffer.toByteArray());
}
return spillFile.getPath();
}
public synchronized long bytesWritten() {
return bytesWritten;
}
@Override
public synchronized void close() {
closeSpillQuietly();
}
private void closeSpillQuietly() {
if (spillStream != null) {
try {
spillStream.close();
} catch (IOException e) {
log.debug("Ignoring close error on spill stream: {}", e.getMessage());
}
spillStream = null;
}
if (spillFile != null) {
spillFile.close(); // TempFile.close() deletes the file
spillFile = null;
}
}
/**
* Routes bytes both to the real client output stream AND to our buffer (in-memory or spilled).
* Writes are not re-batched — each call to the delegate corresponds exactly to one call here.
*/
private final class TeeingServletOutputStream extends ServletOutputStream {
private final ServletOutputStream delegate;
TeeingServletOutputStream(ServletOutputStream delegate) {
this.delegate = delegate;
}
@Override
public void write(int b) throws IOException {
delegate.write(b);
recordSingleByte((byte) b);
}
@Override
public void write(byte[] b, int off, int len) throws IOException {
delegate.write(b, off, len);
recordRange(b, off, len);
}
@Override
public void flush() throws IOException {
delegate.flush();
// Read spilled / spillStream under the outer monitor so we observe the publication
// written by spillToDisk() on another thread. The writers (recordSingleByte,
// recordRange, spillToDisk) already mutate these fields under
// synchronized(PaygResponseBodyWrapper.this); without matching locking here the
// JMM permits stale reads (false `spilled`, null `spillStream`) and Aikido AI
// flagged that gap.
synchronized (PaygResponseBodyWrapper.this) {
if (spilled && spillStream != null) {
spillStream.flush();
}
}
}
@Override
public void close() throws IOException {
delegate.close();
synchronized (PaygResponseBodyWrapper.this) {
if (spilled && spillStream != null) {
spillStream.flush();
}
}
}
@Override
public boolean isReady() {
return delegate.isReady();
}
@Override
public void setWriteListener(WriteListener writeListener) {
delegate.setWriteListener(writeListener);
}
}
private synchronized void recordSingleByte(byte b) throws IOException {
if (spilled) {
spillStream.write(b & 0xFF);
} else if (bytesWritten + 1 > inMemoryThresholdBytes) {
spillToDisk();
spillStream.write(b & 0xFF);
} else {
memoryBuffer.write(b & 0xFF);
}
bytesWritten++;
}
private synchronized void recordRange(byte[] b, int off, int len) throws IOException {
if (spilled) {
spillStream.write(b, off, len);
} else if (bytesWritten + len > inMemoryThresholdBytes) {
// This write crosses the threshold. Spill the existing in-memory buffer, then write
// this entire chunk to disk too — we don't bother splitting it for the sake of staying
// exactly at the threshold. Going over by at most one chunk is fine.
spillToDisk();
spillStream.write(b, off, len);
} else {
memoryBuffer.write(b, off, len);
}
bytesWritten += len;
}
private void spillToDisk() throws IOException {
spillFile = tempFileManager.createManagedTempFile(".body");
spillStream = Files.newOutputStream(spillFile.getPath());
memoryBuffer.writeTo(spillStream);
memoryBuffer = null; // help GC of potentially large buffer
spilled = true;
log.debug(
"PaygResponseBodyWrapper spilled to {} after {} bytes (threshold {})",
spillFile.getPath(),
bytesWritten,
inMemoryThresholdBytes);
}
}
@@ -0,0 +1,138 @@
package stirling.software.saas.payg.filter;
import java.io.IOException;
import org.springframework.context.annotation.Profile;
import org.springframework.stereotype.Component;
import org.springframework.web.filter.OncePerRequestFilter;
import jakarta.servlet.AsyncEvent;
import jakarta.servlet.AsyncListener;
import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import stirling.software.common.util.TempFileManager;
/**
* Wraps the {@link HttpServletResponse} with a {@link PaygResponseBodyWrapper} for every request,
* making the response body retrievable by the downstream {@code PaygChargeInterceptor} in {@code
* afterCompletion}. The wrapper is stashed as a request attribute under {@link #REQUEST_ATTRIBUTE}
* so the interceptor can find it.
*
* <p>Pure plumbing — no business logic. When {@code payg.filter.enabled=false} the filter passes
* through unchanged.
*
* <p>Lifecycle: the wrapper is closed in a {@code finally} after the chain returns for sync
* requests; for async controllers ({@code DeferredResult}, {@code CompletableFuture}), close is
* deferred to an {@link AsyncListener} so the wrapper survives the async window. Close is
* idempotent so a defensive call by the interceptor's {@code afterCompletion} is harmless.
*/
@Slf4j
@Component
@Profile("saas")
public class PaygResponseBodyWrapperFilter extends OncePerRequestFilter {
/** Request-attribute key under which the wrapper is exposed to the interceptor. */
public static final String REQUEST_ATTRIBUTE =
PaygResponseBodyWrapperFilter.class.getName() + ".WRAPPER";
private final TempFileManager tempFileManager;
private final PaygFilterProperties properties;
public PaygResponseBodyWrapperFilter(
TempFileManager tempFileManager, PaygFilterProperties properties) {
this.tempFileManager = tempFileManager;
this.properties = properties;
}
@Override
protected void doFilterInternal(
HttpServletRequest request, HttpServletResponse response, FilterChain chain)
throws ServletException, IOException {
if (!properties.isEnabled()) {
chain.doFilter(request, response);
return;
}
PaygResponseBodyWrapper wrapper;
try {
wrapper =
new PaygResponseBodyWrapper(
response,
tempFileManager,
properties.getResponse().getInMemoryThresholdBytes());
} catch (RuntimeException e) {
// Wrapper construction failure: fail-open. Pass through unwrapped — OUTPUT recording
// is lost but the customer's tool call still runs.
log.warn("PaygResponseBodyWrapper construction failed; passing through unwrapped", e);
chain.doFilter(request, response);
return;
}
request.setAttribute(REQUEST_ATTRIBUTE, wrapper);
boolean asyncStarted = false;
try {
chain.doFilter(request, wrapper);
asyncStarted = request.isAsyncStarted();
if (asyncStarted) {
// Async controller: defer attribute removal + close to async dispatch completion.
// The interceptor's afterCompletion fires on the async dispatch and needs the
// wrapper attribute still present at that point. close() is idempotent so a
// defensive call by the interceptor is harmless.
request.getAsyncContext().addListener(new ReleaseOnAsyncComplete(request, wrapper));
}
} finally {
if (!asyncStarted) {
// Sync path: interceptor.afterCompletion has already run inside chain.doFilter.
request.removeAttribute(REQUEST_ATTRIBUTE);
wrapper.close();
}
}
}
/**
* For async dispatches, removes the wrapper attribute and closes the wrapper after the async
* dispatch completes. The Servlet container fires exactly one of {@code onComplete} / {@code
* onError} / {@code onTimeout} per async context lifecycle.
*/
private static final class ReleaseOnAsyncComplete implements AsyncListener {
private final HttpServletRequest request;
private final PaygResponseBodyWrapper wrapper;
ReleaseOnAsyncComplete(HttpServletRequest request, PaygResponseBodyWrapper wrapper) {
this.request = request;
this.wrapper = wrapper;
}
private void release() {
request.removeAttribute(REQUEST_ATTRIBUTE);
wrapper.close();
}
@Override
public void onComplete(AsyncEvent event) {
release();
}
@Override
public void onTimeout(AsyncEvent event) {
release();
}
@Override
public void onError(AsyncEvent event) {
release();
}
@Override
public void onStartAsync(AsyncEvent event) {
// re-dispatch retains the listener — no-op
}
}
}
@@ -0,0 +1,60 @@
package stirling.software.saas.payg.filter;
import org.springframework.boot.web.servlet.FilterRegistrationBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Profile;
import org.springframework.web.servlet.config.annotation.InterceptorRegistry;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;
import lombok.RequiredArgsConstructor;
/**
* Wires the PAYG filter + interceptor into Spring MVC. Two registrations:
*
* <ul>
* <li>{@link PaygResponseBodyWrapperFilter} as a Servlet filter — registered with no explicit
* order so it sits at the end of the Spring filter chain (after all security filters). Pure
* response-wrapping plumbing.
* <li>{@link PaygChargeInterceptor} as a Spring MVC interceptor — registered AFTER {@code
* UnifiedCreditInterceptor} so legacy credit rejections short-circuit before we hash inputs.
* Both intercept {@code /api/**} with the same admin/info/health exclusions as the legacy
* config.
* </ul>
*/
@Configuration
@Profile("saas")
@RequiredArgsConstructor
public class PaygWebMvcConfig implements WebMvcConfigurer {
private final PaygChargeInterceptor paygChargeInterceptor;
@Bean
public FilterRegistrationBean<PaygResponseBodyWrapperFilter>
paygResponseBodyWrapperFilterRegistration(PaygResponseBodyWrapperFilter filter) {
FilterRegistrationBean<PaygResponseBodyWrapperFilter> reg =
new FilterRegistrationBean<>(filter);
reg.addUrlPatterns("/api/*");
return reg;
}
/**
* Interceptor ordering: the legacy {@code UnifiedCreditInterceptor} (registered with default
* order = 0 in {@code CreditInterceptorConfig}) must run BEFORE this one so credit rejections
* short-circuit before we hash inputs. Explicit positive order guarantees this regardless of
* {@code WebMvcConfigurer} bean discovery order.
*/
public static final int INTERCEPTOR_ORDER = 1000;
@Override
public void addInterceptors(InterceptorRegistry registry) {
registry.addInterceptor(paygChargeInterceptor)
.addPathPatterns("/api/**")
.excludePathPatterns(
"/api/v1/credits/**",
"/api/v1/config/**",
"/api/v1/info/**",
"/api/v1/admin/**")
.order(INTERCEPTOR_ORDER);
}
}
@@ -6,9 +6,12 @@ import java.time.Duration;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.UUID;
import org.springframework.beans.factory.annotation.Value;
@@ -20,6 +23,7 @@ import lombok.extern.slf4j.Slf4j;
import stirling.software.saas.payg.lineage.HashLineageDetector;
import stirling.software.saas.payg.lineage.LineageMatch;
import stirling.software.saas.payg.lineage.LineageSignature;
import stirling.software.saas.payg.model.ArtifactKind;
import stirling.software.saas.payg.model.JobStatus;
import stirling.software.saas.payg.model.JobStepStatus;
@@ -86,7 +90,15 @@ public class JobService {
throw new IllegalArgumentException("inputs must not be empty");
}
Optional<LineageMatch> bestMatch = findBestMatch(ctx.ownerUserId(), inputs);
// Extract signatures ONCE per input, then reuse for both the lineage lookup and the
// post-decision record() call. Avoids hashing every input twice on the hot path.
Map<Path, Set<LineageSignature>> signaturesByInput = new HashMap<>(inputs.size());
for (Path input : inputs) {
signaturesByInput.put(input, detector.extractSignatures(input));
}
Optional<LineageMatch> bestMatch =
findBestMatch(ctx.ownerUserId(), inputs, signaturesByInput);
if (bestMatch.isPresent()) {
ProcessingJob existing =
@@ -100,7 +112,7 @@ public class JobService {
+ " but no such ProcessingJob row"
+ " exists (stale signature?)"));
if (existing.getStepCount() < ctx.stepLimit()) {
return joinExisting(existing, inputs);
return joinExisting(existing, signaturesByInput);
}
// Step-limit hit: spawn a fresh job. The new job will share input signatures with
// the existing chain so future tool calls still lineage-match into the workflow,
@@ -111,7 +123,7 @@ public class JobService {
ctx.stepLimit());
}
return openFresh(ctx, inputs);
return openFresh(ctx, signaturesByInput);
}
/**
@@ -199,25 +211,26 @@ public class JobService {
return stale.size();
}
private Optional<LineageMatch> findBestMatch(Long userId, List<Path> inputs)
throws IOException {
private Optional<LineageMatch> findBestMatch(
Long userId, List<Path> inputs, Map<Path, Set<LineageSignature>> signaturesByInput) {
List<LineageMatch> matches = new ArrayList<>(inputs.size());
for (Path input : inputs) {
detector.detect(userId, input).ifPresent(matches::add);
detector.detect(userId, signaturesByInput.get(input)).ifPresent(matches::add);
}
return matches.stream().max(Comparator.comparing(LineageMatch::jobLastStepAt));
}
private JoinOrOpenResult joinExisting(ProcessingJob existing, List<Path> inputs)
throws IOException {
private JoinOrOpenResult joinExisting(
ProcessingJob existing, Map<Path, Set<LineageSignature>> signaturesByInput) {
existing.setStepCount(existing.getStepCount() + 1);
existing.setLastStepAt(LocalDateTime.now());
ProcessingJob saved = jobRepository.save(existing);
recordAllInputs(saved.getId(), inputs);
recordAllInputs(saved.getId(), signaturesByInput);
return new JoinOrOpenResult(saved, JoinOrOpenResult.Disposition.JOINED);
}
private JoinOrOpenResult openFresh(JobContext ctx, List<Path> inputs) throws IOException {
private JoinOrOpenResult openFresh(
JobContext ctx, Map<Path, Set<LineageSignature>> signaturesByInput) {
ProcessingJob fresh = new ProcessingJob();
fresh.setId(UUID.randomUUID());
fresh.setOwnerUserId(ctx.ownerUserId());
@@ -231,13 +244,13 @@ public class JobService {
fresh.setLastStepAt(now);
fresh.setStatus(JobStatus.OPEN);
ProcessingJob saved = jobRepository.save(fresh);
recordAllInputs(saved.getId(), inputs);
recordAllInputs(saved.getId(), signaturesByInput);
return new JoinOrOpenResult(saved, JoinOrOpenResult.Disposition.OPENED);
}
private void recordAllInputs(UUID jobId, List<Path> inputs) throws IOException {
for (Path input : inputs) {
detector.record(jobId, input, ArtifactKind.INPUT);
private void recordAllInputs(UUID jobId, Map<Path, Set<LineageSignature>> signaturesByInput) {
for (Set<LineageSignature> signatures : signaturesByInput.values()) {
detector.record(jobId, signatures, ArtifactKind.INPUT);
}
}
}
@@ -58,37 +58,45 @@ public class DefaultHashLineageDetector implements HashLineageDetector {
@Override
public Optional<LineageMatch> detect(Long userId, Path inputFile) throws IOException {
Objects.requireNonNull(userId, "userId");
Objects.requireNonNull(inputFile, "inputFile");
Set<LineageSignature> signatures = extractAll(inputFile);
if (signatures.isEmpty()) {
// No extractor produced anything for this content. Treat as no-match.
log.debug("No signatures extracted from {}; lineage check returns empty.", inputFile);
return Optional.empty();
return detect(userId, extractSignatures(inputFile));
}
@Override
public Optional<LineageMatch> detect(Long userId, Set<LineageSignature> signatures) {
Objects.requireNonNull(userId, "userId");
Objects.requireNonNull(signatures, "signatures");
if (signatures.isEmpty()) {
// No extractor produced anything for this content. Treat as no-match.
return Optional.empty();
}
return store.findOpenJobForSignatures(userId, signatures, workflowWindow);
}
@Override
public void record(UUID jobId, Path file, ArtifactKind kind) throws IOException {
Objects.requireNonNull(jobId, "jobId");
Objects.requireNonNull(file, "file");
Objects.requireNonNull(kind, "kind");
record(jobId, extractSignatures(file), kind);
}
Set<LineageSignature> signatures = extractAll(file);
@Override
public void record(UUID jobId, Set<LineageSignature> signatures, ArtifactKind kind) {
Objects.requireNonNull(jobId, "jobId");
Objects.requireNonNull(signatures, "signatures");
Objects.requireNonNull(kind, "kind");
if (signatures.isEmpty()) {
log.debug(
"No signatures extracted from {} for job {} ({}); nothing recorded.",
file,
jobId,
kind);
log.debug("No signatures to record for job {} ({}); skipping.", jobId, kind);
return;
}
store.record(jobId, signatures, kind);
}
@Override
public Set<LineageSignature> extractSignatures(Path file) {
Objects.requireNonNull(file, "file");
return extractAll(file);
}
private Set<LineageSignature> extractAll(Path file) {
Set<LineageSignature> union = new HashSet<>();
for (LineageSignatureExtractor extractor : extractors) {
@@ -3,6 +3,7 @@ package stirling.software.saas.payg.lineage;
import java.io.IOException;
import java.nio.file.Path;
import java.util.Optional;
import java.util.Set;
import java.util.UUID;
import stirling.software.saas.payg.model.ArtifactKind;
@@ -33,4 +34,25 @@ public interface HashLineageDetector {
/** Records the file's signatures against the given job as either INPUT or OUTPUT. */
void record(UUID jobId, Path file, ArtifactKind kind) throws IOException;
/**
* Pre-computes the signature set for {@code file} so callers can avoid hashing the same bytes
* twice when they need both {@link #detect} and {@link #record} for the same file. Returned set
* may be empty (no extractor recognised the content) — treat as "no signatures to match or
* record" by both consumers.
*/
Set<LineageSignature> extractSignatures(Path file);
/**
* Same as {@link #detect(Long, Path)} but operating on pre-computed signatures. Useful when the
* caller has already extracted them (e.g. via {@link #extractSignatures}) and wants to avoid
* re-hashing. Empty {@code signatures} short-circuits to {@link Optional#empty}.
*/
Optional<LineageMatch> detect(Long userId, Set<LineageSignature> signatures);
/**
* Same as {@link #record(UUID, Path, ArtifactKind)} but operating on pre-computed signatures.
* Empty {@code signatures} is a no-op.
*/
void record(UUID jobId, Set<LineageSignature> signatures, ArtifactKind kind);
}
@@ -0,0 +1,18 @@
package stirling.software.saas.payg.model;
/**
* Lifecycle state for a {@code payg_shadow_charge} row. Mimics the Stripe meter_event_adjustment
* (type=cancel) mechanism that real-mode will invoke when a freshly-opened process fails on its
* first step. The reconciliation report's "true net" query is {@code SUM(payg_units) WHERE status =
* 'CHARGED'}.
*/
public enum ShadowChargeStatus {
/** Default at process open; the would-be charge stands. */
CHARGED,
/**
* Set when a 5xx first-step failure was observed in the same request's {@code afterCompletion}.
* Real-mode equivalent is a successful Stripe {@code meter_event_adjustment} posted in the same
* transaction-after-commit hook.
*/
REFUNDED
}
@@ -79,19 +79,28 @@ public class PricingPolicyService {
* {@code is_default = TRUE}. Throws {@link IllegalStateException} if no default exists — the
* seed migration is expected to put one there.
*
* <p>A {@code null} {@code teamId} (admin-created user, deleted team, pre-team-migration
* account) is valid and returns the default policy directly. Throwing here would silently
* exclude the team-less cohort from PAYG via the filter's fail-open path, biasing the
* shadow-vs-legacy reconciliation.
*
* <p>{@link Transactional}({@code readOnly = true}) so the eager-loaded {@code stepLimits} and
* {@code stripePriceIds} collections initialize inside the same session.
*/
@Transactional(readOnly = true)
public PricingPolicy getEffectivePolicy(Long teamId) {
Objects.requireNonNull(teamId, "teamId");
if (teamId == null) {
return loadDefaultPolicy();
}
return byTeamCache.get(teamId, this::loadEffectivePolicy);
}
/** Bypasses the cache. Useful for admin endpoints that want a fresh read after a mutation. */
@Transactional(readOnly = true)
public PricingPolicy getEffectivePolicyUncached(Long teamId) {
Objects.requireNonNull(teamId, "teamId");
if (teamId == null) {
return loadDefaultPolicy();
}
return loadEffectivePolicy(teamId);
}
@@ -243,6 +252,10 @@ public class PricingPolicyService {
teamId,
id);
}
return loadDefaultPolicy();
}
private PricingPolicy loadDefaultPolicy() {
return policyRepository
.findFirstByIsDefaultTrue()
.orElseThrow(
@@ -2,6 +2,8 @@ package stirling.software.saas.payg.repository;
import java.time.LocalDateTime;
import java.util.List;
import java.util.Optional;
import java.util.UUID;
import org.springframework.data.jpa.repository.JpaRepository;
import org.springframework.data.jpa.repository.Query;
@@ -19,4 +21,12 @@ public interface PaygShadowChargeRepository extends JpaRepository<PaygShadowChar
+ " ORDER BY s.occurredAt DESC")
List<PaygShadowCharge> findInWindow(
@Param("from") LocalDateTime from, @Param("to") LocalDateTime to);
/**
* The shadow row written when the given process was opened. At most one row per {@code jobId}
* exists by construction — {@code openProcess} writes exactly one row on OPENED and zero on
* JOINED — so callers can treat the result as a single optional. Returns the first row by id
* defensively if a duplicate ever appears.
*/
Optional<PaygShadowCharge> findFirstByJobIdOrderByIdAsc(UUID jobId);
}
@@ -8,6 +8,8 @@ import org.hibernate.annotations.CreationTimestamp;
import jakarta.persistence.Column;
import jakarta.persistence.Entity;
import jakarta.persistence.EnumType;
import jakarta.persistence.Enumerated;
import jakarta.persistence.GeneratedValue;
import jakarta.persistence.GenerationType;
import jakarta.persistence.Id;
@@ -17,10 +19,16 @@ import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.Setter;
import stirling.software.saas.payg.model.ShadowChargeStatus;
/**
* Per-job comparison row written while a team is in {@code PAYG_SHADOW} mode: what the legacy
* engine actually charged vs. what the PAYG engine would have charged. Aggregated daily by the
* shadow-reconciliation report; deletable after promotion.
*
* <p>The {@link #status} column mirrors the eventual Stripe meter_event_adjustment(cancel) flow:
* rows are written {@code CHARGED}; a freshly-opened process whose first step 5xx's is flipped to
* {@code REFUNDED} in the same request's {@code afterCompletion}.
*/
@Entity
@Table(name = "payg_shadow_charge")
@@ -55,6 +63,18 @@ public class PaygShadowCharge implements Serializable {
@Column(name = "diff_pct", nullable = false)
private Integer diffPct;
@Enumerated(EnumType.STRING)
@Column(name = "status", nullable = false, length = 16)
private ShadowChargeStatus status = ShadowChargeStatus.CHARGED;
/** Set when {@link #status} flips to {@link ShadowChargeStatus#REFUNDED}. */
@Column(name = "refunded_at")
private LocalDateTime refundedAt;
/** Free-form reason, e.g. {@code "first-step-5xx:503"}. */
@Column(name = "refund_reason", length = 128)
private String refundReason;
@CreationTimestamp
@Column(name = "occurred_at", nullable = false, updatable = false)
private LocalDateTime occurredAt;
@@ -0,0 +1,20 @@
-- Refund tracking on shadow rows. Shadow models the eventual Stripe meter_event_adjustment(cancel)
-- by flipping status from CHARGED to REFUNDED in the same request's afterCompletion when a
-- freshly-opened process fails with 5xx on its first step.
--
-- Reconciliation report selects SUM(payg_units) WHERE status = 'CHARGED' to get the true net
-- Stripe would bill.
ALTER TABLE payg_shadow_charge
ADD COLUMN IF NOT EXISTS status VARCHAR(16) NOT NULL DEFAULT 'CHARGED',
ADD COLUMN IF NOT EXISTS refunded_at TIMESTAMP,
ADD COLUMN IF NOT EXISTS refund_reason VARCHAR(128);
CREATE INDEX IF NOT EXISTS idx_payg_shadow_status_time
ON payg_shadow_charge (status, occurred_at);
-- Hot-path index for findFirstByJobIdOrderByIdAsc: hit on every 5xx-first-step refund to flip
-- the row to REFUNDED. UNIQUE because at most one shadow row exists per processing_job by
-- construction (openProcess writes exactly one on OPENED, zero on JOINED).
CREATE UNIQUE INDEX IF NOT EXISTS uq_payg_shadow_job_id
ON payg_shadow_charge (job_id);
@@ -36,9 +36,11 @@ import stirling.software.saas.payg.job.ProcessingJob;
import stirling.software.saas.payg.model.JobSource;
import stirling.software.saas.payg.model.JobStatus;
import stirling.software.saas.payg.model.ProcessType;
import stirling.software.saas.payg.model.ShadowChargeStatus;
import stirling.software.saas.payg.policy.PricingPolicy;
import stirling.software.saas.payg.policy.PricingPolicyService;
import stirling.software.saas.payg.repository.PaygShadowChargeRepository;
import stirling.software.saas.payg.repository.ProcessingJobRepository;
import stirling.software.saas.payg.shadow.PaygShadowCharge;
/**
@@ -51,6 +53,7 @@ class JobChargeServiceTest {
private PricingPolicyService policyService;
private DocumentClassifier classifier;
private PaygShadowChargeRepository shadowRepo;
private ProcessingJobRepository jobRepo;
private JobChargeService service;
@BeforeEach
@@ -59,7 +62,8 @@ class JobChargeServiceTest {
policyService = Mockito.mock(PricingPolicyService.class);
classifier = Mockito.mock(DocumentClassifier.class);
shadowRepo = Mockito.mock(PaygShadowChargeRepository.class);
service = new JobChargeService(jobService, policyService, classifier, shadowRepo);
jobRepo = Mockito.mock(ProcessingJobRepository.class);
service = new JobChargeService(jobService, policyService, classifier, shadowRepo, jobRepo);
}
@Test
@@ -97,7 +101,7 @@ class JobChargeServiceTest {
.thenReturn(new JoinOrOpenResult(newJob, JoinOrOpenResult.Disposition.OPENED));
JobInput in = jobInput(tmp, "in.pdf", "application/pdf");
when(classifier.classify(any(MultipartFile.class), eq(policy)))
when(classifier.classify(any(MultipartFile.class), any(Path.class), eq(policy)))
.thenReturn(new DocumentMetrics(50, 1024L, "application/pdf", 4));
ChargeOutcome out =
@@ -107,9 +111,10 @@ class JobChargeServiceTest {
assertThat(out.disposition()).isEqualTo(ChargeOutcome.Disposition.OPENED);
assertThat(out.units()).isEqualTo(4);
// Single-file path called single-file classifier overload, not the list one.
verify(classifier, times(1)).classify(any(MultipartFile.class), eq(policy));
verify(classifier, never()).classify(anyList(), eq(policy));
// Single-file path called single-file classifier overload (with Path), not the list one.
verify(classifier, times(1))
.classify(any(MultipartFile.class), any(Path.class), eq(policy));
verify(classifier, never()).classify(anyList(), anyList(), eq(policy));
ArgumentCaptor<PaygShadowCharge> captor = ArgumentCaptor.forClass(PaygShadowCharge.class);
verify(shadowRepo).save(captor.capture());
@@ -136,7 +141,7 @@ class JobChargeServiceTest {
JobInput a = jobInput(tmp, "a.pdf", "application/pdf");
JobInput b = jobInput(tmp, "b.pdf", "application/pdf");
when(classifier.classify(anyList(), eq(policy)))
when(classifier.classify(anyList(), anyList(), eq(policy)))
.thenReturn(new DocumentMetrics(100, 2048L, "application/pdf", 7));
ChargeOutcome out =
@@ -145,8 +150,8 @@ class JobChargeServiceTest {
List.of(a, b));
assertThat(out.units()).isEqualTo(7);
verify(classifier, never()).classify(any(MultipartFile.class), any());
verify(classifier, times(1)).classify(anyList(), eq(policy));
verify(classifier, never()).classify(any(MultipartFile.class), any(Path.class), any());
verify(classifier, times(1)).classify(anyList(), anyList(), eq(policy));
}
@Test
@@ -157,7 +162,7 @@ class JobChargeServiceTest {
ProcessingJob newJob = openJob(UUID.randomUUID());
when(jobService.joinOrOpen(any(JobContext.class), anyList()))
.thenReturn(new JoinOrOpenResult(newJob, JoinOrOpenResult.Disposition.OPENED));
when(classifier.classify(any(MultipartFile.class), eq(policy)))
when(classifier.classify(any(MultipartFile.class), any(Path.class), eq(policy)))
.thenReturn(new DocumentMetrics(10, 1024L, "application/pdf", 2));
ChargeOutcome out =
@@ -178,7 +183,7 @@ class JobChargeServiceTest {
ProcessingJob newJob = openJob(UUID.randomUUID());
when(jobService.joinOrOpen(any(JobContext.class), anyList()))
.thenReturn(new JoinOrOpenResult(newJob, JoinOrOpenResult.Disposition.OPENED));
when(classifier.classify(any(MultipartFile.class), eq(policy)))
when(classifier.classify(any(MultipartFile.class), any(Path.class), eq(policy)))
.thenReturn(new DocumentMetrics(1, 100L, "application/pdf", 1));
service.openProcess(
@@ -202,7 +207,7 @@ class JobChargeServiceTest {
ProcessingJob newJob = openJob(UUID.randomUUID());
when(jobService.joinOrOpen(any(JobContext.class), anyList()))
.thenReturn(new JoinOrOpenResult(newJob, JoinOrOpenResult.Disposition.OPENED));
when(classifier.classify(any(MultipartFile.class), eq(policy)))
when(classifier.classify(any(MultipartFile.class), any(Path.class), eq(policy)))
.thenReturn(new DocumentMetrics(1, 100L, "application/pdf", 1));
service.openProcess(
@@ -226,6 +231,105 @@ class JobChargeServiceTest {
.hasMessageContaining("inputs must not be empty");
}
@Test
void markFirstStepFailed_flipsShadowRowAndClosesProcess() {
UUID jobId = UUID.randomUUID();
PaygShadowCharge row = new PaygShadowCharge();
row.setJobId(jobId);
row.setStatus(ShadowChargeStatus.CHARGED);
when(shadowRepo.findFirstByJobIdOrderByIdAsc(jobId)).thenReturn(java.util.Optional.of(row));
ProcessingJob job = openJob(jobId);
when(jobRepo.findById(jobId)).thenReturn(java.util.Optional.of(job));
service.markFirstStepFailed(jobId, "first-step-5xx:503");
assertThat(row.getStatus()).isEqualTo(ShadowChargeStatus.REFUNDED);
assertThat(row.getRefundedAt()).isNotNull();
assertThat(row.getRefundReason()).isEqualTo("first-step-5xx:503");
assertThat(job.getStatus()).isEqualTo(JobStatus.CLOSED);
assertThat(job.getClosedAt()).isNotNull();
verify(shadowRepo).save(row);
verify(jobRepo).save(job);
}
@Test
void markFirstStepFailed_alreadyRefunded_isNoOp() {
UUID jobId = UUID.randomUUID();
PaygShadowCharge row = new PaygShadowCharge();
row.setJobId(jobId);
row.setStatus(ShadowChargeStatus.REFUNDED);
when(shadowRepo.findFirstByJobIdOrderByIdAsc(jobId)).thenReturn(java.util.Optional.of(row));
ProcessingJob job = openJob(jobId);
job.setStatus(JobStatus.CLOSED);
when(jobRepo.findById(jobId)).thenReturn(java.util.Optional.of(job));
service.markFirstStepFailed(jobId, "first-step-5xx:500");
verify(shadowRepo, never()).save(any());
verify(jobRepo, never()).save(any());
}
@Test
void markFirstStepFailed_noShadowRow_stillClosesProcess() {
UUID jobId = UUID.randomUUID();
when(shadowRepo.findFirstByJobIdOrderByIdAsc(jobId)).thenReturn(java.util.Optional.empty());
ProcessingJob job = openJob(jobId);
when(jobRepo.findById(jobId)).thenReturn(java.util.Optional.of(job));
service.markFirstStepFailed(jobId, "first-step-5xx:503");
assertThat(job.getStatus()).isEqualTo(JobStatus.CLOSED);
verify(jobRepo).save(job);
}
@Test
void markFirstStepFailed_trimsLongRefundReason() {
UUID jobId = UUID.randomUUID();
PaygShadowCharge row = new PaygShadowCharge();
row.setStatus(ShadowChargeStatus.CHARGED);
when(shadowRepo.findFirstByJobIdOrderByIdAsc(jobId)).thenReturn(java.util.Optional.of(row));
when(jobRepo.findById(jobId)).thenReturn(java.util.Optional.empty());
String oversized = "x".repeat(200);
service.markFirstStepFailed(jobId, oversized);
assertThat(row.getRefundReason()).hasSize(128);
}
@Test
void decrementStepCount_decrementsByOne() {
UUID jobId = UUID.randomUUID();
ProcessingJob job = openJob(jobId);
job.setStepCount(3);
when(jobRepo.findById(jobId)).thenReturn(java.util.Optional.of(job));
service.decrementStepCount(jobId);
assertThat(job.getStepCount()).isEqualTo(2);
verify(jobRepo).save(job);
}
@Test
void decrementStepCount_floorAtOne_neverDrivesNegative() {
UUID jobId = UUID.randomUUID();
ProcessingJob job = openJob(jobId);
job.setStepCount(1);
when(jobRepo.findById(jobId)).thenReturn(java.util.Optional.of(job));
service.decrementStepCount(jobId);
assertThat(job.getStepCount()).isEqualTo(1);
verify(jobRepo, never()).save(any());
}
@Test
void decrementStepCount_missingJob_isNoOp() {
UUID jobId = UUID.randomUUID();
when(jobRepo.findById(jobId)).thenReturn(java.util.Optional.empty());
service.decrementStepCount(jobId); // must not throw
verify(jobRepo, never()).save(any());
}
// --- helpers --------------------------------------------------------------------------------
private static PricingPolicy stubPolicy(int minCharge, Map<JobSource, Integer> stepLimits) {
@@ -0,0 +1,420 @@
package stirling.software.saas.payg.filter;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyList;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.when;
import java.io.IOException;
import java.lang.reflect.Method;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Optional;
import java.util.UUID;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;
import org.springframework.mock.web.MockHttpServletRequest;
import org.springframework.mock.web.MockHttpServletResponse;
import org.springframework.mock.web.MockMultipartFile;
import org.springframework.mock.web.MockMultipartHttpServletRequest;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.authority.SimpleGrantedAuthority;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.web.method.HandlerMethod;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
import stirling.software.common.annotations.AutoJobPostMapping;
import stirling.software.common.model.ApplicationProperties;
import stirling.software.common.util.TempFileManager;
import stirling.software.common.util.TempFileRegistry;
import stirling.software.proprietary.security.database.repository.UserRepository;
import stirling.software.proprietary.security.model.User;
import stirling.software.saas.payg.charge.ChargeOutcome;
import stirling.software.saas.payg.charge.JobChargeService;
import stirling.software.saas.payg.job.JobService;
import stirling.software.saas.payg.model.JobStepStatus;
/**
* Pure-Mockito tests for {@link PaygChargeInterceptor}. Real {@link TempFileManager} but mocked
* downstream charge/job services so the test runs without a Spring context.
*/
class PaygChargeInterceptorTest {
private JobChargeService chargeService;
private JobService jobService;
private UserRepository userRepository;
private PaygOutputExtractor outputExtractor;
private PaygFilterProperties properties;
private MeterRegistry meterRegistry;
private TempFileManager tempFileManager;
private PaygChargeInterceptor interceptor;
@BeforeEach
void setUp() {
chargeService = org.mockito.Mockito.mock(JobChargeService.class);
jobService = org.mockito.Mockito.mock(JobService.class);
userRepository = org.mockito.Mockito.mock(UserRepository.class);
outputExtractor = org.mockito.Mockito.mock(PaygOutputExtractor.class);
properties = new PaygFilterProperties();
meterRegistry = new SimpleMeterRegistry();
tempFileManager = new TempFileManager(new TempFileRegistry(), new ApplicationProperties());
interceptor =
new PaygChargeInterceptor(
chargeService,
jobService,
userRepository,
tempFileManager,
outputExtractor,
properties,
meterRegistry);
SecurityContextHolder.clearContext();
}
@Test
void preHandle_filterDisabled_isShortCircuitNoop() throws Exception {
properties.setEnabled(false);
MockHttpServletRequest req = new MockHttpServletRequest();
MockHttpServletResponse res = new MockHttpServletResponse();
boolean cont = interceptor.preHandle(req, res, handlerMethodForFakeController());
assertThat(cont).isTrue();
verifyNoInteractions(chargeService);
}
@Test
void preHandle_handlerNotAnnotated_isShortCircuitNoop() throws Exception {
MockHttpServletRequest req = new MockHttpServletRequest();
MockHttpServletResponse res = new MockHttpServletResponse();
boolean cont = interceptor.preHandle(req, res, handlerMethodForPlain());
assertThat(cont).isTrue();
verifyNoInteractions(chargeService);
}
@Test
void preHandle_noAuth_isShortCircuitNoop() throws Exception {
MockMultipartHttpServletRequest req = newMultipart();
req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "hi".getBytes()));
boolean cont =
interceptor.preHandle(
req, new MockHttpServletResponse(), handlerMethodForFakeController());
assertThat(cont).isTrue();
verifyNoInteractions(chargeService);
}
@Test
void preHandle_noMultipartParts_isShortCircuitNoop() throws Exception {
// Authenticated but no file parts.
authenticateWithUser(makeUser(1L, null));
MockHttpServletRequest req = new MockHttpServletRequest();
boolean cont =
interceptor.preHandle(
req, new MockHttpServletResponse(), handlerMethodForFakeController());
assertThat(cont).isTrue();
verifyNoInteractions(chargeService);
}
@Test
void preHandle_openedDisposition_stashesJobIdAndDisposition() throws Exception {
authenticateWithUser(makeUser(7L, 42L));
UUID jobId = UUID.randomUUID();
when(chargeService.openProcess(any(), anyList()))
.thenReturn(new ChargeOutcome(jobId, 4, ChargeOutcome.Disposition.OPENED));
MockMultipartHttpServletRequest req = newMultipart();
req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "abc".getBytes()));
boolean cont =
interceptor.preHandle(
req, new MockHttpServletResponse(), handlerMethodForFakeController());
assertThat(cont).isTrue();
assertThat(req.getAttribute(PaygChargeInterceptor.ATTR_JOB_ID)).isEqualTo(jobId);
assertThat(req.getAttribute(PaygChargeInterceptor.ATTR_DISPOSITION))
.isEqualTo(ChargeOutcome.Disposition.OPENED);
assertThat(req.getAttribute(PaygChargeInterceptor.ATTR_INPUT_TEMP_FILES)).isNotNull();
assertThat(meterRegistry.counter("payg.filter.calls", "disposition", "OPENED").count())
.isEqualTo(1.0);
}
@Test
void preHandle_chargeServiceThrows_failsOpenAndIncrementsErrorCounter() throws Exception {
authenticateWithUser(makeUser(7L, 42L));
when(chargeService.openProcess(any(), anyList())).thenThrow(new RuntimeException("boom"));
MockMultipartHttpServletRequest req = newMultipart();
req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "abc".getBytes()));
boolean cont =
interceptor.preHandle(
req, new MockHttpServletResponse(), handlerMethodForFakeController());
assertThat(cont).isTrue(); // fail-open
assertThat(req.getAttribute(PaygChargeInterceptor.ATTR_FAILED)).isEqualTo(Boolean.TRUE);
assertThat(meterRegistry.counter("payg.filter.errors").count()).isEqualTo(1.0);
}
@Test
void afterCompletion_2xx_appendsStepAndRecordsOutputs() throws Exception {
authenticateWithUser(makeUser(7L, 42L));
UUID jobId = UUID.randomUUID();
when(chargeService.openProcess(any(), anyList()))
.thenReturn(new ChargeOutcome(jobId, 1, ChargeOutcome.Disposition.OPENED));
MockMultipartHttpServletRequest req = newMultipart();
req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "abc".getBytes()));
MockHttpServletResponse res = new MockHttpServletResponse();
res.setStatus(200);
res.setContentType("application/pdf");
// Install a wrapper as the filter would. Pre-populate it with some bytes so
// materialisedPath returns non-null.
PaygResponseBodyWrapper wrapper = new PaygResponseBodyWrapper(res, tempFileManager, 1024);
wrapper.getOutputStream().write("body".getBytes(StandardCharsets.UTF_8));
req.setAttribute(PaygResponseBodyWrapperFilter.REQUEST_ATTRIBUTE, wrapper);
interceptor.preHandle(req, res, handlerMethodForFakeController());
when(outputExtractor.extract(eq("application/pdf"), any()))
.thenReturn(
List.of(
new PaygOutputExtractor.ExtractedPdf(
wrapper.materialisedPath(), null)));
interceptor.afterCompletion(req, res, handlerMethodForFakeController(), null);
ArgumentCaptor<JobStepStatus> status = ArgumentCaptor.forClass(JobStepStatus.class);
verify(jobService).appendStep(eq(jobId), any(), status.capture(), any(), any(), any());
assertThat(status.getValue()).isEqualTo(JobStepStatus.OK);
verify(jobService).recordOutput(eq(jobId), any());
verify(chargeService, never()).markFirstStepFailed(any(), any());
verify(chargeService, never()).decrementStepCount(any());
}
@Test
void afterCompletion_5xx_opened_callsMarkFirstStepFailed() throws Exception {
authenticateWithUser(makeUser(7L, 42L));
UUID jobId = UUID.randomUUID();
when(chargeService.openProcess(any(), anyList()))
.thenReturn(new ChargeOutcome(jobId, 1, ChargeOutcome.Disposition.OPENED));
MockMultipartHttpServletRequest req = newMultipart();
req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "abc".getBytes()));
MockHttpServletResponse res = new MockHttpServletResponse();
res.setStatus(503);
interceptor.preHandle(req, res, handlerMethodForFakeController());
interceptor.afterCompletion(req, res, handlerMethodForFakeController(), null);
verify(chargeService).markFirstStepFailed(eq(jobId), eq("first-step-5xx:503"));
verify(chargeService, never()).decrementStepCount(any());
verify(jobService, never()).recordOutput(any(), any());
verify(jobService)
.appendStep(eq(jobId), any(), eq(JobStepStatus.FAILED), any(), any(), eq("503"));
assertThat(meterRegistry.counter("payg.filter.refunds").count()).isEqualTo(1.0);
}
@Test
void afterCompletion_5xx_joined_callsDecrementStepCount() throws Exception {
authenticateWithUser(makeUser(7L, 42L));
UUID jobId = UUID.randomUUID();
when(chargeService.openProcess(any(), anyList()))
.thenReturn(new ChargeOutcome(jobId, 0, ChargeOutcome.Disposition.JOINED));
MockMultipartHttpServletRequest req = newMultipart();
req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "abc".getBytes()));
MockHttpServletResponse res = new MockHttpServletResponse();
res.setStatus(500);
interceptor.preHandle(req, res, handlerMethodForFakeController());
interceptor.afterCompletion(req, res, handlerMethodForFakeController(), null);
verify(chargeService).decrementStepCount(jobId);
verify(chargeService, never()).markFirstStepFailed(any(), any());
}
@Test
void afterCompletion_4xx_appendsFailedStepNoRefundNoOutputs() throws Exception {
authenticateWithUser(makeUser(7L, 42L));
UUID jobId = UUID.randomUUID();
when(chargeService.openProcess(any(), anyList()))
.thenReturn(new ChargeOutcome(jobId, 1, ChargeOutcome.Disposition.OPENED));
MockMultipartHttpServletRequest req = newMultipart();
req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "abc".getBytes()));
MockHttpServletResponse res = new MockHttpServletResponse();
res.setStatus(422);
interceptor.preHandle(req, res, handlerMethodForFakeController());
interceptor.afterCompletion(req, res, handlerMethodForFakeController(), null);
verify(chargeService, never()).markFirstStepFailed(any(), any());
verify(chargeService, never()).decrementStepCount(any());
verify(jobService, never()).recordOutput(any(), any());
verify(jobService)
.appendStep(eq(jobId), any(), eq(JobStepStatus.FAILED), any(), any(), eq("422"));
}
@Test
void afterCompletion_preHandleFailedFlag_cleansUpWithoutCharging() throws Exception {
MockMultipartHttpServletRequest req = newMultipart();
req.setAttribute(PaygChargeInterceptor.ATTR_FAILED, Boolean.TRUE);
MockHttpServletResponse res = new MockHttpServletResponse();
interceptor.afterCompletion(req, res, handlerMethodForFakeController(), null);
verifyNoInteractions(chargeService);
verifyNoInteractions(jobService);
}
@Test
void afterCompletion_noJobId_isNoop() throws Exception {
MockMultipartHttpServletRequest req = newMultipart();
MockHttpServletResponse res = new MockHttpServletResponse();
interceptor.afterCompletion(req, res, handlerMethodForFakeController(), null);
verifyNoInteractions(chargeService);
verifyNoInteractions(jobService);
}
@Test
void afterCompletion_maxBytesExceeded_skipsOutputRecording() throws Exception {
properties.getResponse().setMaxBytes(2L);
authenticateWithUser(makeUser(7L, 42L));
UUID jobId = UUID.randomUUID();
when(chargeService.openProcess(any(), anyList()))
.thenReturn(new ChargeOutcome(jobId, 1, ChargeOutcome.Disposition.OPENED));
MockMultipartHttpServletRequest req = newMultipart();
req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "abc".getBytes()));
MockHttpServletResponse res = new MockHttpServletResponse();
res.setStatus(200);
res.setContentType("application/pdf");
PaygResponseBodyWrapper wrapper = new PaygResponseBodyWrapper(res, tempFileManager, 1024);
wrapper.getOutputStream().write("1234567890".getBytes(StandardCharsets.UTF_8));
req.setAttribute(PaygResponseBodyWrapperFilter.REQUEST_ATTRIBUTE, wrapper);
interceptor.preHandle(req, res, handlerMethodForFakeController());
interceptor.afterCompletion(req, res, handlerMethodForFakeController(), null);
verify(outputExtractor, never()).extract(any(), any());
verify(jobService, never()).recordOutput(any(), any());
}
@Test
void preHandle_pipelineHeader_setsJobSourcePipeline() throws Exception {
authenticateWithUser(makeUser(7L, 42L));
UUID jobId = UUID.randomUUID();
when(chargeService.openProcess(any(), anyList()))
.thenReturn(new ChargeOutcome(jobId, 1, ChargeOutcome.Disposition.OPENED));
org.mockito.ArgumentCaptor<stirling.software.saas.payg.charge.ChargeContext> ctxCaptor =
org.mockito.ArgumentCaptor.forClass(
stirling.software.saas.payg.charge.ChargeContext.class);
MockMultipartHttpServletRequest req = newMultipart();
req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "abc".getBytes()));
req.addHeader("X-Stirling-Automation", "true");
interceptor.preHandle(req, new MockHttpServletResponse(), handlerMethodForFakeController());
verify(chargeService).openProcess(ctxCaptor.capture(), anyList());
assertThat(ctxCaptor.getValue().source())
.isEqualTo(stirling.software.saas.payg.model.JobSource.PIPELINE);
}
// --- helpers --------------------------------------------------------------------------------
private MockMultipartHttpServletRequest newMultipart() {
MockMultipartHttpServletRequest r = new MockMultipartHttpServletRequest();
r.setRequestURI("/api/v1/security/test-tool");
return r;
}
private void authenticateWithUser(User user) {
UsernamePasswordAuthenticationToken token =
new UsernamePasswordAuthenticationToken(
"supabase-id-here", null, List.of(new SimpleGrantedAuthority("ROLE_USER")));
SecurityContextHolder.getContext().setAuthentication(token);
// resolveUser does UUID.fromString on the name, so use a real UUID string.
String supabaseId = UUID.randomUUID().toString();
UsernamePasswordAuthenticationToken realToken =
new UsernamePasswordAuthenticationToken(
supabaseId, null, List.of(new SimpleGrantedAuthority("ROLE_USER")));
SecurityContextHolder.getContext().setAuthentication(realToken);
when(userRepository.findBySupabaseId(UUID.fromString(supabaseId)))
.thenReturn(Optional.of(user));
}
private static User makeUser(Long id, Long teamId) {
User user = new User();
try {
// User entity doesn't expose Lombok setters for all fields. Set via reflection.
java.lang.reflect.Field idField = User.class.getDeclaredField("id");
idField.setAccessible(true);
idField.set(user, id);
} catch (ReflectiveOperationException e) {
throw new RuntimeException(e);
}
if (teamId != null) {
stirling.software.proprietary.model.Team team =
new stirling.software.proprietary.model.Team();
team.setId(teamId);
user.setTeam(team);
}
return user;
}
private static HandlerMethod handlerMethodForFakeController() {
try {
Method m = FakeController.class.getDeclaredMethod("handleAuto");
return new HandlerMethod(new FakeController(), m);
} catch (NoSuchMethodException e) {
throw new RuntimeException(e);
}
}
private static HandlerMethod handlerMethodForPlain() {
try {
Method m = FakeController.class.getDeclaredMethod("handlePlain");
return new HandlerMethod(new FakeController(), m);
} catch (NoSuchMethodException e) {
throw new RuntimeException(e);
}
}
static class FakeController {
@AutoJobPostMapping(value = "/x", resourceWeight = 1)
public void handleAuto() {}
public void handlePlain() {}
}
/** Placeholder so AutoCloseable resources flow in some helper methods. */
@SuppressWarnings("unused")
private static void closeQuietly(AutoCloseable c) {
try {
c.close();
} catch (Exception ignored) {
// ignored
}
}
@SuppressWarnings("unused")
private static IOException unused() {
return null;
}
}
@@ -0,0 +1,120 @@
package stirling.software.saas.payg.filter;
import static org.assertj.core.api.Assertions.assertThat;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.List;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
import stirling.software.common.model.ApplicationProperties;
import stirling.software.common.util.TempFileManager;
import stirling.software.common.util.TempFileRegistry;
class PaygOutputExtractorTest {
private final TempFileManager tempFileManager =
new TempFileManager(new TempFileRegistry(), new ApplicationProperties());
private final PaygOutputExtractor extractor = new PaygOutputExtractor(tempFileManager);
@Test
void pdfContentType_returnsBodyPathVerbatim(@TempDir Path tmp) throws IOException {
Path body = tmp.resolve("body.pdf");
Files.write(body, pdfBytes("hello"));
List<PaygOutputExtractor.ExtractedPdf> out = extractor.extract("application/pdf", body);
assertThat(out).hasSize(1);
assertThat(out.get(0).path()).isEqualTo(body);
assertThat(out.get(0).ownedTempFile()).isNull();
}
@Test
void pdfContentType_withParameters_stillMatches(@TempDir Path tmp) throws IOException {
Path body = tmp.resolve("body.pdf");
Files.write(body, pdfBytes("x"));
List<PaygOutputExtractor.ExtractedPdf> out =
extractor.extract("application/pdf; charset=binary", body);
assertThat(out).hasSize(1);
}
@Test
void zipContentType_extractsOnlyPdfEntriesWithValidMagicBytes(@TempDir Path tmp)
throws IOException {
Path zip = tmp.resolve("out.zip");
try (ZipOutputStream zos = new ZipOutputStream(Files.newOutputStream(zip))) {
writeEntry(zos, "doc1.pdf", pdfBytes("one"));
writeEntry(zos, "doc2.pdf", pdfBytes("two"));
writeEntry(zos, "fake.pdf", "not a pdf at all".getBytes(StandardCharsets.UTF_8));
writeEntry(zos, "notes.txt", "ignored".getBytes(StandardCharsets.UTF_8));
}
List<PaygOutputExtractor.ExtractedPdf> out = extractor.extract("application/zip", zip);
try {
assertThat(out).hasSize(2);
for (PaygOutputExtractor.ExtractedPdf p : out) {
byte[] head = new byte[5];
Files.newInputStream(p.path()).read(head);
assertThat(head).startsWith("%PDF-".getBytes(StandardCharsets.UTF_8));
assertThat(p.ownedTempFile()).isNotNull();
}
} finally {
for (PaygOutputExtractor.ExtractedPdf p : out) {
p.close();
}
}
}
@Test
void nonPdfNonZipContentType_returnsEmpty(@TempDir Path tmp) throws IOException {
Path body = tmp.resolve("body.json");
Files.writeString(body, "{\"error\":\"bad request\"}");
assertThat(extractor.extract("application/json", body)).isEmpty();
assertThat(extractor.extract("text/plain", body)).isEmpty();
assertThat(extractor.extract(null, body)).isEmpty();
}
@Test
void nullBodyPath_returnsEmpty() {
assertThat(extractor.extract("application/pdf", null)).isEmpty();
}
@Test
void corruptZip_failsClosedAndReturnsEmpty(@TempDir Path tmp) throws IOException {
Path corrupt = tmp.resolve("garbage.zip");
Files.write(corrupt, "not a zip file at all".getBytes(StandardCharsets.UTF_8));
// Should NOT throw — fail-open: empty list, response still serves normally.
List<PaygOutputExtractor.ExtractedPdf> out = extractor.extract("application/zip", corrupt);
assertThat(out).isEmpty();
}
@Test
void zipContentType_emptyZip_returnsEmpty(@TempDir Path tmp) throws IOException {
Path zip = tmp.resolve("empty.zip");
try (ZipOutputStream zos = new ZipOutputStream(Files.newOutputStream(zip))) {
// no entries
}
assertThat(extractor.extract("application/zip", zip)).isEmpty();
}
private static void writeEntry(ZipOutputStream zos, String name, byte[] data)
throws IOException {
zos.putNextEntry(new ZipEntry(name));
zos.write(data);
zos.closeEntry();
}
private static byte[] pdfBytes(String payload) throws IOException {
ByteArrayOutputStream out = new ByteArrayOutputStream();
// Minimal "looks like a PDF" — just the magic-byte prefix + filler. The extractor only
// checks magic bytes, not full PDF validity.
out.write("%PDF-1.4\n".getBytes(StandardCharsets.UTF_8));
out.write(payload.getBytes(StandardCharsets.UTF_8));
return out.toByteArray();
}
}
@@ -0,0 +1,210 @@
package stirling.software.saas.payg.filter;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import java.io.IOException;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import org.junit.jupiter.api.Test;
import org.springframework.mock.web.MockHttpServletResponse;
import stirling.software.common.model.ApplicationProperties;
import stirling.software.common.util.TempFileManager;
import stirling.software.common.util.TempFileRegistry;
class PaygResponseBodyWrapperTest {
private final TempFileManager tempFileManager =
new TempFileManager(new TempFileRegistry(), new ApplicationProperties());
@Test
void inMemory_smallWrites_clientReceivesAndPathContainsSameBytes() throws Exception {
MockHttpServletResponse downstream = new MockHttpServletResponse();
try (PaygResponseBodyWrapper wrapper =
new PaygResponseBodyWrapper(downstream, tempFileManager, 1024)) {
wrapper.getOutputStream().write("hello world".getBytes(StandardCharsets.UTF_8));
wrapper.getOutputStream().flush();
assertThat(downstream.getContentAsString()).isEqualTo("hello world");
assertThat(wrapper.bytesWritten()).isEqualTo(11);
Path materialised = wrapper.materialisedPath();
assertThat(materialised).isNotNull();
assertThat(Files.readString(materialised, StandardCharsets.UTF_8))
.isEqualTo("hello world");
}
}
@Test
void spill_crossThresholdMidChunk_capturesEverything() throws Exception {
MockHttpServletResponse downstream = new MockHttpServletResponse();
try (PaygResponseBodyWrapper wrapper =
new PaygResponseBodyWrapper(downstream, tempFileManager, 8)) {
// First write: 5 bytes — fits in memory.
wrapper.getOutputStream().write("HELLO".getBytes(StandardCharsets.UTF_8));
// Second write: 6 bytes — crosses the 8-byte threshold mid-chunk; whole chunk
// ends up on disk per the design (we don't split the crossing chunk).
wrapper.getOutputStream().write(" WORLD".getBytes(StandardCharsets.UTF_8));
wrapper.getOutputStream().flush();
assertThat(downstream.getContentAsString()).isEqualTo("HELLO WORLD");
assertThat(wrapper.bytesWritten()).isEqualTo(11);
Path path = wrapper.materialisedPath();
assertThat(Files.readString(path, StandardCharsets.UTF_8)).isEqualTo("HELLO WORLD");
}
}
@Test
void spill_largeSingleWrite_capturedToDisk() throws Exception {
MockHttpServletResponse downstream = new MockHttpServletResponse();
byte[] payload = new byte[64 * 1024]; // 64 KiB
for (int i = 0; i < payload.length; i++) {
payload[i] = (byte) (i % 251);
}
try (PaygResponseBodyWrapper wrapper =
new PaygResponseBodyWrapper(downstream, tempFileManager, 1024)) {
wrapper.getOutputStream().write(payload);
wrapper.getOutputStream().flush();
assertThat(downstream.getContentAsByteArray()).isEqualTo(payload);
assertThat(wrapper.bytesWritten()).isEqualTo(payload.length);
byte[] disk = Files.readAllBytes(wrapper.materialisedPath());
assertThat(disk).isEqualTo(payload);
}
}
@Test
void writer_pathTeesThroughToBuffer() throws Exception {
MockHttpServletResponse downstream = new MockHttpServletResponse();
downstream.setCharacterEncoding("UTF-8");
try (PaygResponseBodyWrapper wrapper =
new PaygResponseBodyWrapper(downstream, tempFileManager, 1024)) {
PrintWriter writer = wrapper.getWriter();
writer.write("a-string-from-writer");
writer.flush();
assertThat(downstream.getContentAsString()).isEqualTo("a-string-from-writer");
Path path = wrapper.materialisedPath();
assertThat(Files.readString(path, StandardCharsets.UTF_8))
.isEqualTo("a-string-from-writer");
}
}
@Test
void mixingOutputStreamAndWriter_throwsPerServletSpec() throws Exception {
MockHttpServletResponse downstream = new MockHttpServletResponse();
try (PaygResponseBodyWrapper wrapper =
new PaygResponseBodyWrapper(downstream, tempFileManager, 1024)) {
wrapper.getOutputStream();
assertThatThrownBy(wrapper::getWriter).isInstanceOf(IllegalStateException.class);
}
MockHttpServletResponse downstream2 = new MockHttpServletResponse();
try (PaygResponseBodyWrapper wrapper2 =
new PaygResponseBodyWrapper(downstream2, tempFileManager, 1024)) {
wrapper2.getWriter();
assertThatThrownBy(wrapper2::getOutputStream).isInstanceOf(IllegalStateException.class);
}
}
@Test
void noBytesWritten_materialisedPathIsNull() throws Exception {
MockHttpServletResponse downstream = new MockHttpServletResponse();
try (PaygResponseBodyWrapper wrapper =
new PaygResponseBodyWrapper(downstream, tempFileManager, 1024)) {
// Don't touch the output stream.
assertThat(wrapper.materialisedPath()).isNull();
assertThat(wrapper.bytesWritten()).isZero();
}
}
@Test
void resetBuffer_inMemory_clearsAccumulatedBytes() throws Exception {
MockHttpServletResponse downstream = new MockHttpServletResponse();
try (PaygResponseBodyWrapper wrapper =
new PaygResponseBodyWrapper(downstream, tempFileManager, 1024)) {
wrapper.getOutputStream().write("draft".getBytes(StandardCharsets.UTF_8));
wrapper.resetBuffer();
assertThat(wrapper.bytesWritten()).isZero();
assertThat(wrapper.materialisedPath()).isNull();
wrapper.getOutputStream().write("final".getBytes(StandardCharsets.UTF_8));
wrapper.getOutputStream().flush();
assertThat(Files.readString(wrapper.materialisedPath(), StandardCharsets.UTF_8))
.isEqualTo("final");
}
}
@Test
void resetBuffer_afterSpill_dropsSpillFileAndStartsFresh() throws Exception {
MockHttpServletResponse downstream = new MockHttpServletResponse();
try (PaygResponseBodyWrapper wrapper =
new PaygResponseBodyWrapper(downstream, tempFileManager, 4)) {
wrapper.getOutputStream().write("draftover".getBytes(StandardCharsets.UTF_8)); // spill
assertThat(wrapper.bytesWritten()).isEqualTo(9);
wrapper.resetBuffer();
assertThat(wrapper.bytesWritten()).isZero();
assertThat(wrapper.materialisedPath()).isNull();
// Write again — should land in fresh memory, not the dropped spill.
wrapper.getOutputStream().write("ok".getBytes(StandardCharsets.UTF_8));
wrapper.getOutputStream().flush();
assertThat(Files.readString(wrapper.materialisedPath(), StandardCharsets.UTF_8))
.isEqualTo("ok");
}
}
@Test
void close_isIdempotent() throws IOException {
MockHttpServletResponse downstream = new MockHttpServletResponse();
PaygResponseBodyWrapper wrapper =
new PaygResponseBodyWrapper(downstream, tempFileManager, 8);
wrapper.getOutputStream().write("more-than-eight-bytes".getBytes(StandardCharsets.UTF_8));
wrapper.materialisedPath(); // forces flush
wrapper.close();
wrapper.close(); // second call must not throw
}
@Test
void materialisedPath_isStableAcrossCalls_whenInMemory() throws Exception {
MockHttpServletResponse downstream = new MockHttpServletResponse();
try (PaygResponseBodyWrapper wrapper =
new PaygResponseBodyWrapper(downstream, tempFileManager, 1024)) {
wrapper.getOutputStream().write("abc".getBytes(StandardCharsets.UTF_8));
Path first = wrapper.materialisedPath();
Path second = wrapper.materialisedPath();
assertThat(first).isEqualTo(second);
}
}
@Test
void singleByteWrites_areCorrectlyAccountedAcrossThreshold() throws Exception {
MockHttpServletResponse downstream = new MockHttpServletResponse();
try (PaygResponseBodyWrapper wrapper =
new PaygResponseBodyWrapper(downstream, tempFileManager, 3)) {
for (int b : "ABCDE".getBytes(StandardCharsets.UTF_8)) {
wrapper.getOutputStream().write(b);
}
wrapper.getOutputStream().flush();
assertThat(downstream.getContentAsString()).isEqualTo("ABCDE");
assertThat(wrapper.bytesWritten()).isEqualTo(5);
assertThat(Files.readString(wrapper.materialisedPath(), StandardCharsets.UTF_8))
.isEqualTo("ABCDE");
}
}
@Test
void negativeThreshold_rejected() {
assertThatThrownBy(
() ->
new PaygResponseBodyWrapper(
new MockHttpServletResponse(), tempFileManager, -1))
.isInstanceOf(IllegalArgumentException.class);
}
}
@@ -322,10 +322,17 @@ class JobServiceTest {
return p;
}
/** Test double: programmable detector + records observations. */
/**
* Test double: programmable detector + records observations. Synthesises one signature per
* path; every path seen by extractSignatures is reverse-mapped so the post-dedupe flow
* (extractSignatures → detect/record by signature) keeps the path-based test assertions working
* unchanged.
*/
private static class FakeDetector implements HashLineageDetector {
private final Map<Path, LineageMatch> matches = new HashMap<>();
private final Set<String> observations = new java.util.HashSet<>();
private final Map<stirling.software.saas.payg.lineage.LineageSignature, Path>
pathBySignature = new HashMap<>();
void willMatch(Path input, LineageMatch match) {
matches.put(input, match);
@@ -335,14 +342,53 @@ class JobServiceTest {
return observations.contains(jobId + "|" + file + "|" + kind);
}
private stirling.software.saas.payg.lineage.LineageSignature sigFor(Path file) {
stirling.software.saas.payg.lineage.LineageSignature sig =
new stirling.software.saas.payg.lineage.LineageSignature(
"test", Integer.toHexString(file.toString().hashCode()));
pathBySignature.put(sig, file);
return sig;
}
@Override
public Optional<LineageMatch> detect(Long userId, Path inputFile) {
return Optional.ofNullable(matches.get(inputFile));
}
@Override
public Optional<LineageMatch> detect(
Long userId, Set<stirling.software.saas.payg.lineage.LineageSignature> signatures) {
for (stirling.software.saas.payg.lineage.LineageSignature sig : signatures) {
Path p = pathBySignature.get(sig);
if (p != null && matches.containsKey(p)) {
return Optional.of(matches.get(p));
}
}
return Optional.empty();
}
@Override
public void record(UUID jobId, Path file, ArtifactKind kind) {
observations.add(jobId + "|" + file + "|" + kind);
}
@Override
public void record(
UUID jobId,
Set<stirling.software.saas.payg.lineage.LineageSignature> signatures,
ArtifactKind kind) {
for (stirling.software.saas.payg.lineage.LineageSignature sig : signatures) {
Path p = pathBySignature.get(sig);
if (p != null) {
observations.add(jobId + "|" + p + "|" + kind);
}
}
}
@Override
public Set<stirling.software.saas.payg.lineage.LineageSignature> extractSignatures(
Path file) {
return Set.of(sigFor(file));
}
}
}
@@ -98,6 +98,33 @@ class PricingPolicyServiceTest {
.hasMessageContaining("No default pricing_policy row");
}
@Test
void nullTeamId_returnsDefaultPolicyDirectly_noCacheLookup() {
PricingPolicy result = service.getEffectivePolicy(null);
assertThat(result).isEqualTo(defaultPolicy);
verify(extensionsRepo, never()).findById(any());
// Team-less users go straight to the default — no cache pollution either.
assertThat(service.cacheSize()).isZero();
}
@Test
void nullTeamIdUncached_returnsDefaultPolicyDirectly() {
PricingPolicy result = service.getEffectivePolicyUncached(null);
assertThat(result).isEqualTo(defaultPolicy);
verify(extensionsRepo, never()).findById(any());
}
@Test
void nullTeamId_noDefaultExists_throws() {
when(policyRepo.findFirstByIsDefaultTrue()).thenReturn(Optional.empty());
assertThatThrownBy(() -> service.getEffectivePolicy(null))
.isInstanceOf(IllegalStateException.class)
.hasMessageContaining("No default pricing_policy row");
}
@Test
void secondCallHitsCache_noRepoLookup() {
when(extensionsRepo.findById(42L)).thenReturn(Optional.empty());