fe(payg): show the usage-limit modal when the limit is hit (direct + policy) (#6626)

app/proprietary/src/main/java/stirling/software/proprietary/policy/engine/PolicyEngine.java

@@ -14,6 +14,7 @@ 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.web.client.RestClientResponseException;
 
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
@@ -199,6 +200,26 @@ public class PolicyEngine {
                     e.getMessage());
             run.fail(message);
             taskManager.setError(runId, message);
+        } catch (RestClientResponseException e) {
+            // A downstream tool call returned an error status. When it's a structured entitlement
+            // response (401/402 with a JSON `error` sentinel), surface that code onto the run so
+            // the
+            // client can react — e.g. pop the usage-limit modal — instead of only seeing a generic
+            // 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);
+            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));
+                taskManager.setError(runId, message);
+            } else {
+                String message = "Policy run failed: " + e.getMessage();
+                log.error("Policy run {} failed (downstream HTTP error)", runId, e);
+                run.fail(message);
+                taskManager.setError(runId, message);
+            }
         } catch (Exception e) {
             String message = "Policy run failed: " + e.getMessage();
             log.error("Policy run {} failed", runId, e);
@@ -281,6 +302,49 @@ 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/policy/model/PolicyRun.java

@@ -27,6 +27,21 @@ public class PolicyRun {
 
     private volatile WaitState waitState;
     private volatile String error;
+
+    /**
+     * Stable, machine-readable failure code the client can branch on — e.g. an entitlement-limit
+     * sentinel ({@code PAYG_LIMIT_REACHED} / {@code FEATURE_DEGRADED}) propagated from a downstream
+     * tool call's 402 — alongside the human-readable {@link #error}. Null unless set on failure.
+     */
+    private volatile String errorCode;
+
+    /**
+     * For an entitlement-limit failure, whether the team was subscribed (over its spending cap) vs
+     * un-subscribed (free allowance spent) — taken from the blocking 402 body. Drives which
+     * usage-limit modal the client shows. Null unless {@link #errorCode} is an entitlement code.
+     */
+    private volatile Boolean errorSubscribed;
+
     private volatile List<ResultFile> outputs = List.of();
     private volatile Instant updatedAt = Instant.now();
 
@@ -61,6 +76,17 @@ public class PolicyRun {
         touch();
     }
 
+    /**
+     * Fail with a stable {@code errorCode} the client can branch on (e.g. an entitlement-limit
+     * sentinel from a downstream 402), plus the optional {@code subscribed} flag from that
+     * response, in addition to the human-readable message.
+     */
+    public synchronized void failWithCode(String message, String errorCode, Boolean subscribed) {
+        this.errorCode = errorCode;
+        this.errorSubscribed = subscribed;
+        fail(message);
+    }
+
     public synchronized void waitForInput(WaitState wait) {
         this.waitState = wait;
         this.status = PolicyRunStatus.WAITING_FOR_INPUT;

app/proprietary/src/main/java/stirling/software/proprietary/policy/model/PolicyRunView.java

@@ -14,6 +14,8 @@ public record PolicyRunView(
         int currentStep,
         int stepCount,
         String error,
+        String errorCode,
+        Boolean errorSubscribed,
         List<ResultFile> outputs) {
 
     public static PolicyRunView of(PolicyRun run) {
@@ -23,6 +25,8 @@ public record PolicyRunView(
                 run.getCurrentStep(),
                 run.stepCount(),
                 run.getError(),
+                run.getErrorCode(),
+                run.getErrorSubscribed(),
                 run.getOutputs());
     }
 }

app/proprietary/src/test/java/stirling/software/proprietary/policy/engine/PolicyEngineTest.java

@@ -31,7 +31,10 @@ 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.HttpHeaders;
+import org.springframework.http.HttpStatus;
 import org.springframework.http.ResponseEntity;
+import org.springframework.web.client.HttpClientErrorException;
 
 import stirling.software.common.model.ApplicationProperties;
 import stirling.software.common.service.FileStorage;
@@ -171,6 +174,35 @@ class PolicyEngineTest {
         verify(taskManager, never()).setComplete(runId);
     }
 
+    @Test
+    void runBlockedByUsageLimit_surfacesErrorCodeAndSubscribed() throws Exception {
+        // A downstream tool call gets a 402 entitlement block. The run fails, but its errorCode +
+        // subscribed are taken from the 402 body so the client can pop the right usage-limit modal
+        // (the policy 402 happens server-side, out of reach of the apiClient interceptor).
+        when(toolMetadataService.isMultiInput(ROTATE)).thenReturn(false);
+        String body = "{\"error\":\"PAYG_LIMIT_REACHED\",\"subscribed\":true}";
+        when(internalApiClient.post(eq(ROTATE), any()))
+                .thenThrow(
+                        HttpClientErrorException.create(
+                                HttpStatus.PAYMENT_REQUIRED,
+                                "Payment Required",
+                                HttpHeaders.EMPTY,
+                                body.getBytes(java.nio.charset.StandardCharsets.UTF_8),
+                                java.nio.charset.StandardCharsets.UTF_8));
+
+        PolicyRun run =
+                engine.submit(
+                                definition(new PipelineStep(ROTATE, Map.of())),
+                                PolicyInputs.of(List.of(pdf("input", "input.pdf"))),
+                                PolicyProgressListener.NOOP)
+                        .completion()
+                        .get(10, TimeUnit.SECONDS);
+
+        assertEquals(PolicyRunStatus.FAILED, run.getStatus());
+        assertEquals("PAYG_LIMIT_REACHED", run.getErrorCode());
+        assertEquals(Boolean.TRUE, run.getErrorSubscribed());
+    }
+
     @Test
     void runPolicyExecutesThePolicysPipeline() throws Exception {
         when(toolMetadataService.isMultiInput(anyString())).thenReturn(false);

app/saas/src/main/java/stirling/software/saas/payg/entitlement/EntitlementGuard.java

@@ -299,6 +299,10 @@ public class EntitlementGuard implements HandlerInterceptor {
         deniedDegradedCounter.increment();
         Map<String, Object> body = new LinkedHashMap<>();
         body.put("error", "FEATURE_DEGRADED");
+        // subscribed tells the client which usage-limit modal to show: a subscribed team is over
+        // its spending cap; an un-subscribed one has spent its free allowance. (PAYG_LIMIT_REACHED
+        // already carries this; mirror it here so the JWT/web path can pick the right modal too.)
+        body.put("subscribed", snapshot.subscribed());
         body.put("missingGates", missingGates(required, snapshot.enabledGates()));
         body.put("state", snapshot.state().name());
         body.put(

app/saas/src/test/java/stirling/software/saas/payg/entitlement/EntitlementGuardTest.java

@@ -308,6 +308,8 @@ class EntitlementGuardTest {
         JsonNode body = json.readTree(res.getContentAsByteArray());
         assertThat(body.get("error").asText()).isEqualTo("FEATURE_DEGRADED");
         assertThat(body.get("state").asText()).isEqualTo("DEGRADED");
+        // subscribed drives the client's modal choice (free-limit vs spend-cap).
+        assertThat(body.get("subscribed").asBoolean()).isFalse();
         assertThat(body.get("capUnits").asLong()).isEqualTo(500L);
         assertThat(body.get("spendUnits").asLong()).isEqualTo(500L);
         assertThat(body.get("missingGates").isArray()).isTrue();