From da9b5bd48ce439724de3038083151dcdc0553a02 Mon Sep 17 00:00:00 2001 From: Connor Yoh Date: Thu, 11 Jun 2026 21:21:29 +0100 Subject: [PATCH] feat(payg): pop the usage-limit modal when a policy run is blocked MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #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. --- .../policy/engine/PolicyEngine.java | 64 +++++++++++++++++++ .../proprietary/policy/model/PolicyRun.java | 26 ++++++++ .../policy/model/PolicyRunView.java | 4 ++ .../policy/engine/PolicyEngineTest.java | 32 ++++++++++ .../components/policies/policyRunStore.ts | 2 + .../components/policies/usePolicyAutoRun.ts | 57 +++++++++++++++-- .../proprietary/services/policyPipeline.ts | 29 +++++++++ .../saas/components/UsageLimitModalHost.tsx | 18 ++++++ 8 files changed, 225 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 1118d0058d..94b4a86a33 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 @@ -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 diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/policy/model/PolicyRun.java b/app/proprietary/src/main/java/stirling/software/proprietary/policy/model/PolicyRun.java index bdecaacde5..a24114eda9 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/policy/model/PolicyRun.java +++ b/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 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; diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/policy/model/PolicyRunView.java b/app/proprietary/src/main/java/stirling/software/proprietary/policy/model/PolicyRunView.java index bba12a3c5c..61a6c61f66 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/policy/model/PolicyRunView.java +++ b/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 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()); } } 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 af9d7a427a..ddb4eced36 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 @@ -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); diff --git a/frontend/editor/src/proprietary/components/policies/policyRunStore.ts b/frontend/editor/src/proprietary/components/policies/policyRunStore.ts index 434b84dfac..31175dbe20 100644 --- a/frontend/editor/src/proprietary/components/policies/policyRunStore.ts +++ b/frontend/editor/src/proprietary/components/policies/policyRunStore.ts @@ -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; } diff --git a/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.ts b/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.ts index 5dd143e3f1..52ba7b2abf 100644 --- a/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.ts +++ b/frontend/editor/src/proprietary/components/policies/usePolicyAutoRun.ts @@ -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>(new Set()); const dispatching = useRef>(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>(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(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 { +/** + * 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 { for (let i = 0; i < MAX_POLLS; i++) { await delay(POLL_MS); let view; @@ -314,7 +353,11 @@ async function poll(runId: string): Promise { 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; + } } } diff --git a/frontend/editor/src/proprietary/services/policyPipeline.ts b/frontend/editor/src/proprietary/services/policyPipeline.ts index 29f38695a0..3049ef28cd 100644 --- a/frontend/editor/src/proprietary/services/policyPipeline.ts +++ b/frontend/editor/src/proprietary/services/policyPipeline.ts @@ -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[]; } diff --git a/frontend/editor/src/saas/components/UsageLimitModalHost.tsx b/frontend/editor/src/saas/components/UsageLimitModalHost.tsx index c6e2a28876..4c7802cac6 100644 --- a/frontend/editor/src/saas/components/UsageLimitModalHost.tsx +++ b/frontend/editor/src/saas/components/UsageLimitModalHost.tsx @@ -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. + * + *

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).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); }; }, []);