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:
ConnorYoh
2026-06-11 18:28:58 +01:00
committed by GitHub
parent 9ee0bc4b32
commit 9e5fe2f4ca
2 changed files with 161 additions and 7 deletions
@@ -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<PolicyRun> 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);
}
}
}
}
@@ -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 = {"<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
void runIsQueuedUnderResourcePressure() {
when(resourceMonitor.shouldQueueJob(anyInt())).thenReturn(true);