From 9e5fe2f4ca6c1bc1c7766a0826f494a792828088 Mon Sep 17 00:00:00 2001 From: ConnorYoh <40631091+ConnorYoh@users.noreply.github.com> Date: Thu, 11 Jun 2026 18:28:58 +0100 Subject: [PATCH] fix(payg): attribute policy runs to the owner so usage is charged (#6620) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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. --- .../policy/engine/PolicyEngine.java | 87 +++++++++++++++++-- .../policy/engine/PolicyEngineTest.java | 81 +++++++++++++++++ 2 files changed, 161 insertions(+), 7 deletions(-) diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/policy/engine/PolicyEngine.java b/app/proprietary/src/main/java/stirling/software/proprietary/policy/engine/PolicyEngine.java index 26a291589..1118d0058 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/policy/engine/PolicyEngine.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/policy/engine/PolicyEngine.java @@ -8,8 +8,11 @@ import java.util.UUID; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; +import org.slf4j.MDC; import org.springframework.core.io.Resource; import org.springframework.http.ResponseEntity; +import org.springframework.security.core.Authentication; +import org.springframework.security.core.context.SecurityContextHolder; import org.springframework.stereotype.Service; import lombok.RequiredArgsConstructor; @@ -71,6 +74,25 @@ public class PolicyEngine { */ public PolicyRunHandle submit( 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 // ownership check passes. No-op when security is off. String runId = jobOwnershipService.createScopedJobKey(UUID.randomUUID().toString()); @@ -79,7 +101,16 @@ public class PolicyEngine { registry.register(run); CompletableFuture completion = new CompletableFuture<>(); 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 // without the pool-within-pool risk of queueing each tool call. @@ -100,12 +131,6 @@ public class PolicyEngine { 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) { return registry.get(runId); } @@ -241,4 +266,52 @@ public class PolicyEngine { "The %s tool did not respond within %d seconds and was aborted.", 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); + } + } + } } diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/policy/engine/PolicyEngineTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/policy/engine/PolicyEngineTest.java index e47916dbe..af9d7a427 100644 --- a/app/proprietary/src/test/java/stirling/software/proprietary/policy/engine/PolicyEngineTest.java +++ b/app/proprietary/src/test/java/stirling/software/proprietary/policy/engine/PolicyEngineTest.java @@ -28,6 +28,7 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.junit.jupiter.api.io.TempDir; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; +import org.slf4j.MDC; import org.springframework.core.io.ByteArrayResource; import org.springframework.core.io.Resource; import org.springframework.http.ResponseEntity; @@ -204,6 +205,86 @@ class PolicyEngineTest { 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 = {""}; + 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 = {""}; + 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 void runIsQueuedUnderResourcePressure() { when(resourceMonitor.shouldQueueJob(anyInt())).thenReturn(true);