mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-16 19:33:11 +02:00
fix(payg): attribute policy runs to the owner so usage is charged (#6620)
## Problem A policy ran over a file but the owner's **free usage was never consumed**. A policy run executes on a **background virtual thread** (`PolicyEngine.submit` → `asyncExecutor`), and Spring's `SecurityContextHolder` is thread-local — so the worker thread has no identity. When `PolicyExecutor` → `InternalApiClient.post` resolves the tool-call API key via `UserService.getCurrentUsername()`, it finds nothing and falls back to the **`INTERNAL_API_USER`** key. The loopback tool calls then authenticate as that system account, so `PaygChargeInterceptor` attributes the charge to *its* team (or none) — the real owner's free grant is untouched. Folder-watch / scheduled triggers are even further removed (fired from a background watch loop with no request context at all). The charging *mechanism* was fine (AUTOMATION, multipart, `openProcess`); only the **attribution** was wrong. ## Fix Propagate the acting identity onto the worker thread using the **audit-principal MDC key** that `UserService.getCurrentUsername()` already reads as its documented async fallback (the same mechanism used for other async jobs). No new plumbing through the executor. - **`runPolicy`** (stored policies — covers triggers *and* manual `runWith`) → bill the **policy owner**. `Policy.owner` is the username stamped at creation, so `getApiKeyForUser(owner)` resolves it. - **`submit`** (ad-hoc Automate/AI one-offs) → bill the **submitting user**, captured on the request thread (it doesn't survive the hop to the worker otherwise). With the principal set, `InternalApiClient` dispatches each tool call as that user → the interceptor resolves the right team → free grant draws / Stripe meters correctly. ## Tests `PolicyEngineTest`: - `runPolicyDispatchesToolCallsAsTheOwner` — asserts MDC `auditPrincipal` == the policy owner at the moment `InternalApiClient.post` is invoked. - `adHocRunDispatchesToolCallsAsTheSubmittingUser` — asserts it's the submitting user for an ad-hoc run. `:proprietary:test` + `:saas:test` + spotless green; coverage gates met. ## Heads-up (not in this PR) Once attributed, **automatic folder-watch / scheduled runs consume free grant (or bill) per file** — set up once, runs forever. That's automation-is-billable working as intended, but a set-and-forget policy can drain an allowance fast, so it may warrant a per-policy cap or a heads-up in the UI. Flagging for a product decision.
This commit is contained in:
+80
-7
@@ -8,8 +8,11 @@ import java.util.UUID;
|
|||||||
import java.util.concurrent.CompletableFuture;
|
import java.util.concurrent.CompletableFuture;
|
||||||
import java.util.concurrent.ExecutorService;
|
import java.util.concurrent.ExecutorService;
|
||||||
|
|
||||||
|
import org.slf4j.MDC;
|
||||||
import org.springframework.core.io.Resource;
|
import org.springframework.core.io.Resource;
|
||||||
import org.springframework.http.ResponseEntity;
|
import org.springframework.http.ResponseEntity;
|
||||||
|
import org.springframework.security.core.Authentication;
|
||||||
|
import org.springframework.security.core.context.SecurityContextHolder;
|
||||||
import org.springframework.stereotype.Service;
|
import org.springframework.stereotype.Service;
|
||||||
|
|
||||||
import lombok.RequiredArgsConstructor;
|
import lombok.RequiredArgsConstructor;
|
||||||
@@ -71,6 +74,25 @@ public class PolicyEngine {
|
|||||||
*/
|
*/
|
||||||
public PolicyRunHandle submit(
|
public PolicyRunHandle submit(
|
||||||
PipelineDefinition definition, PolicyInputs inputs, PolicyProgressListener listener) {
|
PipelineDefinition definition, PolicyInputs inputs, PolicyProgressListener listener) {
|
||||||
|
// Ad-hoc run (no stored policy): bill whoever kicked it off. Capture the principal on this
|
||||||
|
// (request) thread — it does not survive the hop onto the async worker.
|
||||||
|
return submitForPrincipal(currentActingPrincipal(), definition, inputs, listener);
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Run a stored policy on demand. {@code enabled} gates triggers, not explicit runs. */
|
||||||
|
public PolicyRunHandle runPolicy(
|
||||||
|
Policy policy, PolicyInputs inputs, PolicyProgressListener listener) {
|
||||||
|
// Bill the policy owner. Trigger-fired runs have no security context at all, and even an
|
||||||
|
// on-demand run executes on a background worker that doesn't inherit the caller's context —
|
||||||
|
// so the owner (a username stamped at policy creation) is the reliable billing identity.
|
||||||
|
return submitForPrincipal(policy.owner(), policy.toDefinition(), inputs, listener);
|
||||||
|
}
|
||||||
|
|
||||||
|
private PolicyRunHandle submitForPrincipal(
|
||||||
|
String actingPrincipal,
|
||||||
|
PipelineDefinition definition,
|
||||||
|
PolicyInputs inputs,
|
||||||
|
PolicyProgressListener listener) {
|
||||||
// Scope the run id to the current user (this request thread) so the file-download
|
// Scope the run id to the current user (this request thread) so the file-download
|
||||||
// ownership check passes. No-op when security is off.
|
// ownership check passes. No-op when security is off.
|
||||||
String runId = jobOwnershipService.createScopedJobKey(UUID.randomUUID().toString());
|
String runId = jobOwnershipService.createScopedJobKey(UUID.randomUUID().toString());
|
||||||
@@ -79,7 +101,16 @@ public class PolicyEngine {
|
|||||||
registry.register(run);
|
registry.register(run);
|
||||||
CompletableFuture<PolicyRun> completion = new CompletableFuture<>();
|
CompletableFuture<PolicyRun> completion = new CompletableFuture<>();
|
||||||
PolicyProgressListener tracking = trackingListener(runId, run, listener);
|
PolicyProgressListener tracking = trackingListener(runId, run, listener);
|
||||||
Runnable task = () -> runToCompletion(run, inputs, tracking, completion);
|
// Re-establish the acting principal as the audit principal on the worker thread. Each tool
|
||||||
|
// step dispatches via InternalApiClient, which resolves the caller from
|
||||||
|
// UserService.getCurrentUsername() — that has an MDC `auditPrincipal` fallback for async
|
||||||
|
// threads. Without this the worker has no identity, tool calls fall back to the
|
||||||
|
// INTERNAL_API_USER, and PAYG charges that system account instead of the owner's team.
|
||||||
|
Runnable task =
|
||||||
|
() ->
|
||||||
|
runAsPrincipal(
|
||||||
|
actingPrincipal,
|
||||||
|
() -> runToCompletion(run, inputs, tracking, completion));
|
||||||
|
|
||||||
// One admission unit per run; steps run synchronously within it, so this gates heavy work
|
// One admission unit per run; steps run synchronously within it, so this gates heavy work
|
||||||
// without the pool-within-pool risk of queueing each tool call.
|
// without the pool-within-pool risk of queueing each tool call.
|
||||||
@@ -100,12 +131,6 @@ public class PolicyEngine {
|
|||||||
return new PolicyRunHandle(runId, completion);
|
return new PolicyRunHandle(runId, completion);
|
||||||
}
|
}
|
||||||
|
|
||||||
/** Run a stored policy on demand. {@code enabled} gates triggers, not explicit runs. */
|
|
||||||
public PolicyRunHandle runPolicy(
|
|
||||||
Policy policy, PolicyInputs inputs, PolicyProgressListener listener) {
|
|
||||||
return submit(policy.toDefinition(), inputs, listener);
|
|
||||||
}
|
|
||||||
|
|
||||||
public PolicyRun getRun(String runId) {
|
public PolicyRun getRun(String runId) {
|
||||||
return registry.get(runId);
|
return registry.get(runId);
|
||||||
}
|
}
|
||||||
@@ -241,4 +266,52 @@ public class PolicyEngine {
|
|||||||
"The %s tool did not respond within %d seconds and was aborted.",
|
"The %s tool did not respond within %d seconds and was aborted.",
|
||||||
e.getEndpointPath(), e.getReadTimeout().toSeconds());
|
e.getEndpointPath(), e.getReadTimeout().toSeconds());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* MDC key {@code UserService.getCurrentUsername()} reads as its async fallback (stamped by the
|
||||||
|
* controller audit aspect on request threads). We reuse it to carry the billing identity onto
|
||||||
|
* the policy worker thread.
|
||||||
|
*/
|
||||||
|
private static final String AUDIT_PRINCIPAL_MDC_KEY = "auditPrincipal";
|
||||||
|
|
||||||
|
/**
|
||||||
|
* The username to bill an ad-hoc run to, captured on the submitting (request) thread. Prefers
|
||||||
|
* the audit principal the controller aspect already stamped; falls back to the security context
|
||||||
|
* name. {@code anonymousUser} (and no identity) resolve to null so we don't try to bill it.
|
||||||
|
*/
|
||||||
|
private static String currentActingPrincipal() {
|
||||||
|
String mdc = MDC.get(AUDIT_PRINCIPAL_MDC_KEY);
|
||||||
|
if (mdc != null && !mdc.isBlank()) {
|
||||||
|
return mdc;
|
||||||
|
}
|
||||||
|
Authentication auth = SecurityContextHolder.getContext().getAuthentication();
|
||||||
|
if (auth == null) {
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
String name = auth.getName();
|
||||||
|
return "anonymousUser".equals(name) ? null : name;
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Run {@code body} with {@code principal} set as the audit principal in MDC, so async tool
|
||||||
|
* dispatch attributes (and charges) usage to that user. A null/blank principal runs as-is.
|
||||||
|
* Restores the previous MDC value afterward (defensive — worker threads aren't pooled).
|
||||||
|
*/
|
||||||
|
private static void runAsPrincipal(String principal, Runnable body) {
|
||||||
|
if (principal == null || principal.isBlank()) {
|
||||||
|
body.run();
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
String previous = MDC.get(AUDIT_PRINCIPAL_MDC_KEY);
|
||||||
|
MDC.put(AUDIT_PRINCIPAL_MDC_KEY, principal);
|
||||||
|
try {
|
||||||
|
body.run();
|
||||||
|
} finally {
|
||||||
|
if (previous != null) {
|
||||||
|
MDC.put(AUDIT_PRINCIPAL_MDC_KEY, previous);
|
||||||
|
} else {
|
||||||
|
MDC.remove(AUDIT_PRINCIPAL_MDC_KEY);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+81
@@ -28,6 +28,7 @@ import org.junit.jupiter.api.extension.ExtendWith;
|
|||||||
import org.junit.jupiter.api.io.TempDir;
|
import org.junit.jupiter.api.io.TempDir;
|
||||||
import org.mockito.Mock;
|
import org.mockito.Mock;
|
||||||
import org.mockito.junit.jupiter.MockitoExtension;
|
import org.mockito.junit.jupiter.MockitoExtension;
|
||||||
|
import org.slf4j.MDC;
|
||||||
import org.springframework.core.io.ByteArrayResource;
|
import org.springframework.core.io.ByteArrayResource;
|
||||||
import org.springframework.core.io.Resource;
|
import org.springframework.core.io.Resource;
|
||||||
import org.springframework.http.ResponseEntity;
|
import org.springframework.http.ResponseEntity;
|
||||||
@@ -204,6 +205,86 @@ class PolicyEngineTest {
|
|||||||
verify(internalApiClient).post(eq(ROTATE), any());
|
verify(internalApiClient).post(eq(ROTATE), any());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void runPolicyDispatchesToolCallsAsTheOwner() throws Exception {
|
||||||
|
// Billing-attribution regression: the pipeline runs on a background worker thread, but the
|
||||||
|
// policy owner must be propagated as the audit principal so InternalApiClient (and thus
|
||||||
|
// PAYG) attributes each tool call to the owner — not the INTERNAL_API_USER fallback.
|
||||||
|
when(toolMetadataService.isMultiInput(anyString())).thenReturn(false);
|
||||||
|
when(toolMetadataService.shouldUnpackZipResponse(anyString())).thenReturn(false);
|
||||||
|
int[] counter = {0};
|
||||||
|
when(fileStorage.storeInputStream(any(InputStream.class), anyString()))
|
||||||
|
.thenAnswer(
|
||||||
|
inv ->
|
||||||
|
new StoredFile(
|
||||||
|
"file-" + ++counter[0],
|
||||||
|
((InputStream) inv.getArgument(0)).readAllBytes().length));
|
||||||
|
|
||||||
|
String[] principalAtDispatch = {"<none>"};
|
||||||
|
when(internalApiClient.post(eq(ROTATE), any()))
|
||||||
|
.thenAnswer(
|
||||||
|
inv -> {
|
||||||
|
principalAtDispatch[0] = MDC.get("auditPrincipal");
|
||||||
|
return ResponseEntity.ok(pdf("rotated", "rotated.pdf"));
|
||||||
|
});
|
||||||
|
|
||||||
|
Policy policy =
|
||||||
|
new Policy(
|
||||||
|
"p1",
|
||||||
|
"rotate",
|
||||||
|
"alice",
|
||||||
|
true,
|
||||||
|
null,
|
||||||
|
List.of(new PipelineStep(ROTATE, Map.of())),
|
||||||
|
OutputSpec.inline());
|
||||||
|
|
||||||
|
engine.runPolicy(
|
||||||
|
policy,
|
||||||
|
PolicyInputs.of(List.of(pdf("input", "input.pdf"))),
|
||||||
|
PolicyProgressListener.NOOP)
|
||||||
|
.completion()
|
||||||
|
.get(10, TimeUnit.SECONDS);
|
||||||
|
|
||||||
|
assertEquals("alice", principalAtDispatch[0]);
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void adHocRunDispatchesToolCallsAsTheSubmittingUser() throws Exception {
|
||||||
|
// Ad-hoc runs (no stored policy) bill whoever kicked them off; the principal is captured on
|
||||||
|
// the request thread (here simulated via MDC) and re-established on the worker thread.
|
||||||
|
when(toolMetadataService.isMultiInput(anyString())).thenReturn(false);
|
||||||
|
when(toolMetadataService.shouldUnpackZipResponse(anyString())).thenReturn(false);
|
||||||
|
int[] counter = {0};
|
||||||
|
when(fileStorage.storeInputStream(any(InputStream.class), anyString()))
|
||||||
|
.thenAnswer(
|
||||||
|
inv ->
|
||||||
|
new StoredFile(
|
||||||
|
"file-" + ++counter[0],
|
||||||
|
((InputStream) inv.getArgument(0)).readAllBytes().length));
|
||||||
|
|
||||||
|
String[] principalAtDispatch = {"<none>"};
|
||||||
|
when(internalApiClient.post(eq(ROTATE), any()))
|
||||||
|
.thenAnswer(
|
||||||
|
inv -> {
|
||||||
|
principalAtDispatch[0] = MDC.get("auditPrincipal");
|
||||||
|
return ResponseEntity.ok(pdf("rotated", "rotated.pdf"));
|
||||||
|
});
|
||||||
|
|
||||||
|
MDC.put("auditPrincipal", "bob"); // the request thread's audit principal
|
||||||
|
try {
|
||||||
|
engine.submit(
|
||||||
|
definition(new PipelineStep(ROTATE, Map.of())),
|
||||||
|
PolicyInputs.of(List.of(pdf("input", "input.pdf"))),
|
||||||
|
PolicyProgressListener.NOOP)
|
||||||
|
.completion()
|
||||||
|
.get(10, TimeUnit.SECONDS);
|
||||||
|
} finally {
|
||||||
|
MDC.remove("auditPrincipal");
|
||||||
|
}
|
||||||
|
|
||||||
|
assertEquals("bob", principalAtDispatch[0]);
|
||||||
|
}
|
||||||
|
|
||||||
@Test
|
@Test
|
||||||
void runIsQueuedUnderResourcePressure() {
|
void runIsQueuedUnderResourcePressure() {
|
||||||
when(resourceMonitor.shouldQueueJob(anyInt())).thenReturn(true);
|
when(resourceMonitor.shouldQueueJob(anyInt())).thenReturn(true);
|
||||||
|
|||||||
Reference in New Issue
Block a user