feat(payg): pop the usage-limit modal when a policy run is blocked

#6626 wired the limit modals for direct browser→API calls, but a policy's tool calls run server-side, so their 402 entitlement block never reaches the apiClient interceptor — the run just showed up "failed" with no modal. Bridge it: BE (proprietary, billing-layer-agnostic): - PolicyEngine.runToCompletion catches a downstream RestClientResponseException; on a 401/402 it regex-extracts the body's `error` sentinel + `subscribed` flag and surfaces them on the run (PolicyRun.errorCode/errorSubscribed → PolicyRunView), instead of only a generic failure string. The code is passed through, not interpreted, so proprietary stays free of saas coupling. FE: - usePolicyAutoRun: when a polled run finishes with an entitlement errorCode, broadcast POLICY_LIMIT_REACHED_EVENT (with `subscribed`), deduped per run so a folder-watch burst opens the modal once. The proprietary hook can't import the saas modal API, so it bridges via a window event. - UsageLimitModalHost (saas): listens for that event and opens the spend-cap (subscribed) or free-limit (free) modal — the same modals direct calls use. Tests: PolicyEngineTest asserts a downstream 402 surfaces errorCode + errorSubscribed on the run. proprietary compileJava, PolicyEngineTest, eslint, tsc (proprietary + saas), and policies/interceptor vitest all green.
2026-06-11 21:21:29 +01:00
parent 124f4af6dd
commit da9b5bd48c
8 changed files with 225 additions and 7 deletions
@@ -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;
@@ -185,6 +186,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);
@@ -267,6 +288,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
@@ -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;
@@ -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());
    }
 }
@@ -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);
@@ -31,6 +31,8 @@ export interface PolicyRunRecord {
   *  which marks the policy's OUTPUT — not the input it ran on. */
  outputFileIds?: string[];
  error: string | null;
+  /** Stable backend failure code (e.g. an entitlement sentinel) when FAILED; null otherwise. */
+  errorCode?: string | null;
  /** Epoch ms when the run was dispatched. */
  startedAt: number;
 }
@@ -11,7 +11,7 @@
 * in the run store), so re-renders and remounts don't re-fire.
 */

-import { useEffect, useRef } from "react";
+import { useCallback, useEffect, useRef } from "react";
 import {
  useAllFiles,
  useFileManagement,
@@ -24,7 +24,12 @@ import {
  getPolicyRun,
  downloadPolicyOutput,
 } from "@app/services/policyApi";
-import type { PolicyRunStatus } from "@app/services/policyPipeline";
+import type {
+  PolicyRunStatus,
+  PolicyRunView,
+  PolicyLimitReachedDetail,
+} from "@app/services/policyPipeline";
+import { POLICY_LIMIT_REACHED_EVENT } from "@app/services/policyPipeline";
 import type { FileId } from "@app/types/file";
 import { createStirlingFilesAndStubs } from "@app/services/fileStubHelpers";
 import type { StirlingFile, StirlingFileStub } from "@app/types/fileContext";
@@ -69,6 +74,30 @@ export function usePolicyAutoRun(): void {
  const importing = useRef<Set<string>>(new Set());
  const dispatching = useRef<Set<string>>(new Set());

+  // A policy's tool calls run server-side, so a usage-limit 402 never reaches the apiClient
+  // interceptor (and thus never pops the modal that direct calls get). The backend surfaces the
+  // limit sentinel on the run's errorCode; when a run we polled finishes blocked, broadcast a
+  // window event. A saas-layer listener (which can read the wallet + open the modal — this
+  // proprietary hook can't import the saas modal API) decides free-limit vs spend-cap. Dedupe per
+  // run so a folder-watch burst opens the modal once, not once per file.
+  const firedLimitModal = useRef<Set<string>>(new Set());
+
+  const onRunFinished = useCallback((view: PolicyRunView) => {
+    const code = view.errorCode;
+    if (code !== "PAYG_LIMIT_REACHED" && code !== "FEATURE_DEGRADED") return;
+    if (firedLimitModal.current.has(view.runId)) return;
+    firedLimitModal.current.add(view.runId);
+    try {
+      window.dispatchEvent(
+        new CustomEvent<PolicyLimitReachedDetail>(POLICY_LIMIT_REACHED_EVENT, {
+          detail: { subscribed: view.errorSubscribed ?? null },
+        }),
+      );
+    } catch {
+      // non-browser env (tests / SSR) — no-op.
+    }
+  }, []);
+
  // Dispatch: for each active policy × each session file not yet run, fire a run.
  useEffect(() => {
    if (!POLICIES_ENABLED) return;
@@ -111,9 +140,11 @@ export function usePolicyAutoRun(): void {
    for (const run of runs) {
      if (isTerminal(run.status) || polling.current.has(run.runId)) continue;
      polling.current.add(run.runId);
-      void poll(run.runId).finally(() => polling.current.delete(run.runId));
+      void poll(run.runId, onRunFinished).finally(() =>
+        polling.current.delete(run.runId),
+      );
    }
-  }, [runs]);
+  }, [runs, onRunFinished]);

  // Import each completed run's outputs into the workspace (each output once),
  // so the enforced file appears in the app rather than only on the backend.
@@ -300,8 +331,16 @@ export async function runPolicyOnFile(
  }
 }

-/** Poll a run's status until it reaches a terminal state (or the cap). */
-async function poll(runId: string): Promise<void> {
+/**
+ * Poll a run's status until it reaches a terminal state (or the cap). Calls {@code onTerminal} once
+ * with the final view when it terminates — the caller uses that to pop the usage-limit modal when a
+ * run was blocked. Only runs polled this session fire it (terminal runs aren't re-polled), so a
+ * persisted failed run never re-triggers a modal on reload.
+ */
+async function poll(
+  runId: string,
+  onTerminal?: (view: PolicyRunView) => void,
+): Promise<void> {
  for (let i = 0; i < MAX_POLLS; i++) {
    await delay(POLL_MS);
    let view;
@@ -314,7 +353,11 @@ async function poll(runId: string): Promise<void> {
      status: view.status,
      outputs: view.outputs,
      error: view.error,
+      errorCode: view.errorCode ?? null,
    });
-    if (isTerminal(view.status)) return;
+    if (isTerminal(view.status)) {
+      onTerminal?.(view);
+      return;
+    }
  }
 }
@@ -56,6 +56,24 @@ export interface BackendPolicy {
  output: BackendOutputSpec;
 }

+/**
+ * Window event fired when a policy run is blocked by a usage limit (its tool call got a 402
+ * entitlement sentinel). Dispatched from the proprietary auto-run poll; a saas-layer listener
+ * reads the wallet and opens the matching usage-limit modal. Kept here (proprietary) so both
+ * layers share one name — proprietary can't import the saas modal API directly.
+ */
+export const POLICY_LIMIT_REACHED_EVENT = "payg:policyLimitReached";
+
+/** Detail carried on {@link POLICY_LIMIT_REACHED_EVENT}. */
+export interface PolicyLimitReachedDetail {
+  /**
+   * Whether the blocked team was subscribed (over its spending cap) vs un-subscribed (free
+   * allowance spent), from the blocking 402. The listener uses it to choose the spend-cap vs
+   * free-limit modal. Null when unknown → treat as free-limit.
+   */
+  subscribed: boolean | null;
+}
+
 /** Lifecycle states of a backend run (mirrors PolicyRunStatus). */
 export type PolicyRunStatus =
  | "PENDING"
@@ -77,6 +95,17 @@ export interface PolicyRunView {
  currentStep: number;
  stepCount: number;
  error: string | null;
+  /**
+   * Stable failure code from the backend (e.g. an entitlement sentinel
+   * {@code PAYG_LIMIT_REACHED} / {@code FEATURE_DEGRADED} propagated from a
+   * downstream tool's 402). Null/absent for ordinary failures.
+   */
+  errorCode?: string | null;
+  /**
+   * For an entitlement-limit failure, the {@code subscribed} flag from the blocking 402 — picks the
+   * spend-cap (true) vs free-limit (false) modal. Null/absent otherwise.
+   */
+  errorSubscribed?: boolean | null;
  outputs: BackendResultFile[];
 }

@@ -5,6 +5,10 @@ import {
  FREE_LIMIT_MODAL_EVENT,
  SPEND_CAP_MODAL_EVENT,
 } from "@app/components/usageLimitModals";
+import {
+  POLICY_LIMIT_REACHED_EVENT,
+  type PolicyLimitReachedDetail,
+} from "@app/services/policyPipeline";

 /**
 * Always-mounted host for the usage-limit warning modals. Mount once (in
@@ -12,6 +16,12 @@ import {
 * (see usageLimitModals.ts) fire their bridge events. Each modal is mounted
 * only while open, so it reads the wallet (and animates in) on open rather
 * than on app load.
+ *
+ * <p>Also bridges the policy auto-run path: a policy's tool calls run server-side,
+ * so their usage-limit 402 never reaches the apiClient interceptor that pops these
+ * modals for direct calls. The proprietary auto-run hook broadcasts {@link
+ * POLICY_LIMIT_REACHED_EVENT} (with the blocking 402's {@code subscribed} flag)
+ * instead; we open the matching modal here.
 */
 export default function UsageLimitModalHost() {
  const [freeOpen, setFreeOpen] = useState(false);
@@ -20,11 +30,19 @@ export default function UsageLimitModalHost() {
  useEffect(() => {
    const onFree = () => setFreeOpen(true);
    const onSpend = () => setSpendOpen(true);
+    const onPolicyLimit = (e: Event) => {
+      const subscribed = (e as CustomEvent<PolicyLimitReachedDetail>).detail
+        ?.subscribed;
+      if (subscribed) setSpendOpen(true);
+      else setFreeOpen(true);
+    };
    window.addEventListener(FREE_LIMIT_MODAL_EVENT, onFree);
    window.addEventListener(SPEND_CAP_MODAL_EVENT, onSpend);
+    window.addEventListener(POLICY_LIMIT_REACHED_EVENT, onPolicyLimit);
    return () => {
      window.removeEventListener(FREE_LIMIT_MODAL_EVENT, onFree);
      window.removeEventListener(SPEND_CAP_MODAL_EVENT, onSpend);
+      window.removeEventListener(POLICY_LIMIT_REACHED_EVENT, onPolicyLimit);
    };
  }, []);