arsvendg/Stirling-PDF
1
0
Fork 0
mirror of https://github.com/arsvendg/Stirling-PDF.git synced 2026-07-14 10:34:06 +02:00
Code Issues Actions 4 Packages Projects Releases 1 Wiki Activity

fix(payg): fire the usage-limit modal when an AI agent run hits the limit (#6638)

Browse Source 
## Problem

We're getting 402s when an AI **agent** (chat) run hits the free
allowance / spending cap, but the frontend handles them poorly and never
pops the usage-limit modal.

The agent runs its tool calls **server-side** (loopback HTTP via
`PolicyExecutor`), so the 402 never reaches the `apiClient` interceptor
that pops the modal for direct calls. It was caught by the generic
tool-failure handler and flattened into a `CANNOT_CONTINUE` reason
string (`"The /api/v1/… tool failed: 402…"`), streamed as a `result`
event, and rendered as a scary chat bubble. This is the same gap the
policy auto-run path bridges (#6626) — one layer up.

## Fix

**Backend** (`proprietary`)
- `AiWorkflowResponse` gains `errorCode` + `errorSubscribed`.
- `AiWorkflowService` detects a downstream 401/402 entitlement sentinel
in its three tool-exec catch sites (`onToolCall`, `runPlan`,
`onConvertMarkdown`) and surfaces the structured code (+ `subscribed`)
on the terminal response instead of the raw failure text.
- Factored the 401/402 body extraction `PolicyEngine` already had into a
shared `DownstreamEntitlementError` util so the two server-side paths
can't drift.

**Frontend**
- New `usageLimitBridge` (`PAYG_LIMIT_REACHED_EVENT` +
`dispatchPaygLimitReached`) generalises the previously policy-only
bridge. Proprietary can't import the saas modal API (layering), so
server-side limit hits broadcast a window event the saas
`UsageLimitModalHost` opens the modal from. Migrated the policy path
onto it.
- `ChatContext` fires the matching modal (free → subscribe, subscribed →
raise cap) on the limit result **and** on a direct 402, replacing the
raw reason with a brief friendly line
(`chat.responses.usage_limit_reached`).

No Python engine changes — the charge/402 happens on the Java tool
endpoint that Java itself calls.

## Test plan

- [x] `:proprietary:compileJava` + `spotlessCheck` clean
- [x] `AiWorkflowServiceTest` + `PolicyEngineTest` green
- [x] eslint, proprietary + saas typechecks clean
- [ ] Manual: drive an agent run over the limit → brief line in chat +
the right modal (free vs cap)

> Note: proprietary test compilation is currently blocked on the
pre-existing `InitialSecuritySetupTest` 6-arg ctor break (unrelated,
tracked separately); verified locally by temporarily patching it.
This commit is contained in:
ConnorYoh
2026-06-12 11:38:07 +01:00
committed by GitHub
parent eb2527fc7f
commit 87723d3ce2
10 changed files with 251 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowResponse.java
+15
View File
@@ -105,4 +105,19 @@ public class AiWorkflowResponse {
+ " body or via the X-Stirling-Tool-Report header. May be null for tools"
+ " that produce only a file.")
private JsonNode report;
@Schema(
description =
"Structured error code when a downstream tool call was blocked (e.g."
+ " PAYG_LIMIT_REACHED). Lets the client react — such as opening the"
+ " usage-limit modal — instead of only seeing a generic failure. Null"
+ " for ordinary outcomes.")
private String errorCode;
@Schema(
description =
"Whether the team is subscribed, carried from a downstream usage-limit response."
+ " Selects which limit modal the client shows (free → subscribe,"
+ " subscribed → raise cap). Null when the downstream body omitted it.")
private Boolean errorSubscribed;
}
app/proprietary/src/main/java/stirling/software/proprietary/policy/engine/PolicyEngine.java
+3 -45
View File
@@ -36,6 +36,7 @@ import stirling.software.proprietary.policy.model.PolicyRun;
import stirling.software.proprietary.policy.model.WaitState;
import stirling.software.proprietary.policy.output.PolicyOutputSink;
import stirling.software.proprietary.policy.progress.PolicyProgressListener;
import stirling.software.proprietary.service.DownstreamEntitlementError;
/**
* Runs pipelines asynchronously as tracked jobs. {@link #submit} returns a run id immediately; the
@@ -208,11 +209,11 @@ public class PolicyEngine {
// failure. We don't interpret the code here (that would couple this module to the saas
// billing layer); we just pass it through for the client to map. Other statuses fall
// through to the generic failure below.
String code = extractDownstreamErrorCode(e);
String code = DownstreamEntitlementError.extractCode(e);
if (code != null) {
log.info("Policy run {} blocked by downstream entitlement gate ({})", runId, code);
String message = "Usage limit reached";
run.failWithCode(message, code, extractDownstreamSubscribed(e));
run.failWithCode(message, code, DownstreamEntitlementError.extractSubscribed(e));
taskManager.setError(runId, message);
} else {
String message = "Policy run failed: " + e.getMessage();
@@ -302,49 +303,6 @@ public class PolicyEngine {
e.getEndpointPath(), e.getReadTimeout().toSeconds());
}
/** Matches the {@code "error":"CODE"} field of a small JSON error body. */
private static final java.util.regex.Pattern ERROR_CODE_FIELD =
java.util.regex.Pattern.compile("\"error\"\\s*:\\s*\"([^\"]+)\"");
/**
* Pull the {@code error} sentinel out of a downstream 401/402 JSON body — e.g. the saas
* EntitlementGuard's {@code {"error":"PAYG_LIMIT_REACHED",...}}. Regex (not a JSON parse) on
* purpose: the body is a small, server-controlled shape and this keeps the proprietary module
* free of any billing-layer coupling. Returns null for other statuses or an unmatched body, in
* which case the caller treats it as a generic failure.
*/
private static String extractDownstreamErrorCode(RestClientResponseException e) {
int status = e.getStatusCode().value();
if (status != 401 && status != 402) {
return null;
}
String body = e.getResponseBodyAsString();
if (body == null || body.isBlank()) {
return null;
}
java.util.regex.Matcher m = ERROR_CODE_FIELD.matcher(body);
return m.find() ? m.group(1) : null;
}
/** Matches the {@code "subscribed":true|false} field of a small JSON error body. */
private static final java.util.regex.Pattern SUBSCRIBED_FIELD =
java.util.regex.Pattern.compile("\"subscribed\"\\s*:\\s*(true|false)");
/**
* Pull the {@code subscribed} flag out of a downstream 401/402 JSON body (present on the saas
* {@code PAYG_LIMIT_REACHED} response). Null when absent — the client then defaults to the
* free-limit modal. Regex for the same dependency-free reason as {@link
* #extractDownstreamErrorCode}.
*/
private static Boolean extractDownstreamSubscribed(RestClientResponseException e) {
String body = e.getResponseBodyAsString();
if (body == null || body.isBlank()) {
return null;
}
java.util.regex.Matcher m = SUBSCRIBED_FIELD.matcher(body);
return m.find() ? Boolean.valueOf(m.group(1)) : null;
}
/**
* 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
app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java
+50
View File
@@ -18,6 +18,7 @@ import org.springframework.http.MediaType;
import org.springframework.http.MediaTypeFactory;
import org.springframework.stereotype.Service;
import org.springframework.web.client.HttpServerErrorException;
import org.springframework.web.client.RestClientResponseException;
import org.springframework.web.multipart.MultipartFile;
import io.github.pixee.security.Filenames;
@@ -385,6 +386,13 @@ public class AiWorkflowService {
return new WorkflowState.Terminal(
cannotContinue(toolTimeoutMessage(PDF_TO_MARKDOWN_ENDPOINT, e)));
} catch (Exception e) {
AiWorkflowResponse limit = paygLimitResponseOrNull(e);
if (limit != null) {
log.info(
"AI markdown conversion blocked by downstream entitlement gate ({})",
limit.getErrorCode());
return new WorkflowState.Terminal(limit);
}
log.error("Failed to convert PDF to Markdown: {}", e.getMessage(), e);
return new WorkflowState.Terminal(
cannotContinue(toolFailureMessage(PDF_TO_MARKDOWN_ENDPOINT, e)));
@@ -470,6 +478,14 @@ public class AiWorkflowService {
log.error("Tool {} timed out: {}", endpointPath, e.getMessage());
return new WorkflowState.Terminal(cannotContinue(toolTimeoutMessage(endpointPath, e)));
} catch (Exception e) {
AiWorkflowResponse limit = paygLimitResponseOrNull(e);
if (limit != null) {
log.info(
"AI workflow tool {} blocked by downstream entitlement gate ({})",
endpointPath,
limit.getErrorCode());
return new WorkflowState.Terminal(limit);
}
log.error("Failed to execute tool {}: {}", endpointPath, e.getMessage(), e);
return new WorkflowState.Terminal(cannotContinue(toolFailureMessage(endpointPath, e)));
}
@@ -585,6 +601,13 @@ public class AiWorkflowService {
log.error("Plan step failed (HTTP {}): {}", e.getStatusCode(), reason);
return new WorkflowState.Terminal(cannotContinue(reason));
} catch (Exception e) {
AiWorkflowResponse limit = paygLimitResponseOrNull(e);
if (limit != null) {
log.info(
"AI workflow plan blocked by downstream entitlement gate ({})",
limit.getErrorCode());
return new WorkflowState.Terminal(limit);
}
log.error("Failed to execute plan: {}", e.getMessage(), e);
return new WorkflowState.Terminal(
cannotContinue("Plan execution failed: " + e.getMessage()));
@@ -722,6 +745,33 @@ public class AiWorkflowService {
return response;
}
/**
* If {@code e} is a downstream usage-limit block — a 401/402 from a tool call carrying the saas
* EntitlementGuard's {@code error} sentinel — build a terminal response that carries the
* structured code (+ {@code subscribed}) through to the client, so it can pop the matching
* usage-limit modal instead of surfacing the raw "tool failed: 402…" text. Returns null for any
* other failure, so the caller falls back to its normal tool-failure handling.
*
* <p>The agent's tool calls run server-side (loopback HTTP via {@link PolicyExecutor}), so this
* 402 never reaches the frontend's API-client interceptor that pops the modal for direct calls
* — same gap the policy auto-run path bridges in {@code PolicyEngine}.
*/
private AiWorkflowResponse paygLimitResponseOrNull(Throwable e) {
if (!(e instanceof RestClientResponseException rce)) {
return null;
}
String code = DownstreamEntitlementError.extractCode(rce);
if (code == null) {
return null;
}
AiWorkflowResponse response = new AiWorkflowResponse();
response.setOutcome(AiWorkflowOutcome.CANNOT_CONTINUE);
response.setReason("You've reached your current usage limit.");
response.setErrorCode(code);
response.setErrorSubscribed(DownstreamEntitlementError.extractSubscribed(rce));
return response;
}
/**
* Drive the engine's streaming orchestrator endpoint. Progress events are forwarded to {@code
* listener} as they arrive (each one keeps the SSE connection to the frontend alive too). The
app/proprietary/src/main/java/stirling/software/proprietary/service/DownstreamEntitlementError.java
+64
View File
@@ -0,0 +1,64 @@
package stirling.software.proprietary.service;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import org.springframework.web.client.RestClientResponseException;
/**
* Reads the {@code error} sentinel and {@code subscribed} flag out of a downstream 401/402 JSON
* body — e.g. the saas EntitlementGuard's {@code
* {"error":"PAYG_LIMIT_REACHED","subscribed":false}}.
*
* <p>Server-side run paths (policy auto-run, AI agent workflows) execute tool calls via loopback
* HTTP, so a usage-limit 402 surfaces as a {@link RestClientResponseException} rather than reaching
* the frontend's API-client interceptor. These helpers let those paths pass the structured code
* through to the client, which maps it to the right usage-limit modal instead of showing a generic
* failure.
*
* <p>Regex (not a JSON parse) on purpose: the body is a small, server-controlled shape and this
* keeps the proprietary module free of any billing-layer (saas) coupling.
*/
public final class DownstreamEntitlementError {
private DownstreamEntitlementError() {}
/** Matches the {@code "error":"CODE"} field of a small JSON error body. */
private static final Pattern ERROR_CODE_FIELD =
Pattern.compile("\"error\"\\s*:\\s*\"([^\"]+)\"");
/** Matches the {@code "subscribed":true|false} field of a small JSON error body. */
private static final Pattern SUBSCRIBED_FIELD =
Pattern.compile("\"subscribed\"\\s*:\\s*(true|false)");
/**
* Pull the {@code error} sentinel out of a downstream 401/402 JSON body. Returns null for other
* statuses or an unmatched body, in which case the caller treats it as a generic failure.
*/
public static String extractCode(RestClientResponseException e) {
int status = e.getStatusCode().value();
if (status != 401 && status != 402) {
return null;
}
String body = e.getResponseBodyAsString();
if (body == null || body.isBlank()) {
return null;
}
Matcher m = ERROR_CODE_FIELD.matcher(body);
return m.find() ? m.group(1) : null;
}
/**
* Pull the {@code subscribed} flag out of the body (present on the saas {@code
* PAYG_LIMIT_REACHED}/{@code FEATURE_DEGRADED} responses). Null when absent — the client then
* defaults to the free-limit modal.
*/
public static Boolean extractSubscribed(RestClientResponseException e) {
String body = e.getResponseBodyAsString();
if (body == null || body.isBlank()) {
return null;
}
Matcher m = SUBSCRIBED_FIELD.matcher(body);
return m.find() ? Boolean.valueOf(m.group(1)) : null;
}
}