History

Add pre-execution balance gate for org execution limits

Check the org's Autumn execution balance before execute/executeWithPause (never resume, so paused executions can always complete). Blocked orgs get a descriptive ExecuteResult.error instead of running. Fails open on any billing error or a 2s timeout, and caches per-org outcomes for 60s.

Rhys Sullivan committed 4 months ago

Commit 45f9afa1 parent 018cb591 file · +205 −0Browse files

Unified · Split

apps/cloud/src/engine/execution-gate.ts+205−0
@@ -0,0 +1,205 @@
1+// ---------------------------------------------------------------------------
2+// Pre-execution balance gate — blocks new executions once an organization has
3+// used up its plan's included executions.
4+//
5+// Usage is tracked to Autumn after every execution (execution-usage.ts), but
6+// nothing ever CHECKED the balance before running, so a free-tier org could run
7+// unbounded executions past its quota. This gate consults
8+// `AutumnService.checkExecutionBalance` before `execute` / `executeWithPause`.
9+// `resume` is never gated: a paused execution already consumed its quota slot
10+// when it started, and blocking resume would strand approved work forever.
11+//
12+// FAIL OPEN is a hard requirement: any Autumn error, timeout, or missing
13+// customer/feature allows the execution (logged + Sentry, mirroring
14+// `trackExecution`'s reporting). Autumn being down must never block
15+// executions or add meaningful latency, so the check is bounded by a short
16+// timeout and its outcome is cached per organization.
17+//
18+// Error surfacing: the gate fails internally with the typed
19+// `ExecutionLimitReachedError`, and the engine seam folds it into the
20+// descriptive `ExecuteResult.error` channel. That is the codebase's mechanism
21+// for getting a domain failure's message to the MCP client verbatim — engine
22+// error-channel failures are deliberately rendered opaque by the MCP host
23+// ("Internal tool error [correlation id]"), exactly like the runtimes fold
24+// `CodeCompilationError` / `SandboxRuntimeError` into `ExecuteResult.error`.
25+// ---------------------------------------------------------------------------
26+
27+import * as Sentry from "@sentry/cloudflare";
28+import { Data, Effect } from "effect";
29+import type * as Cause from "effect/Cause";
30+
31+import type { ExecutionEngine, ExecutionResult } from "@executor-js/execution";
32+
33+// The engine's completed-result payload (`ExecuteResult` in codemode-core),
34+// derived from the public execution types so this app package doesn't need a
35+// direct dependency on the kernel package that declares it.
36+type EngineExecuteResult = Extract<ExecutionResult, { status: "completed" }>["result"];
37+
38+// One check per org per minute is fresh enough for a monthly quota; both
39+// allowed AND blocked outcomes are cached (errors never are).
40+const BALANCE_CACHE_TTL_MS = 60_000;
41+// Autumn slower than this => fail open rather than stall a user-facing
42+// execution behind the billing provider.
43+const BALANCE_CHECK_TIMEOUT_MS = 2_000;
44+// Sweep guard so a long-lived worker isolate serving many orgs can't grow the
45+// cache map unbounded.
46+const BALANCE_CACHE_MAX_ENTRIES = 10_000;
47+
48+export const EXECUTION_LIMIT_BLOCKED_MESSAGE =
49+ "Execution limit reached: your plan's included executions for this billing period are used up. Upgrade your plan or wait for the reset to continue.";
50+
51+export class ExecutionLimitReachedError extends Data.TaggedError("ExecutionLimitReachedError")<{
52+ readonly organizationId: string;
53+ readonly message: string;
54+}> {}
55+
56+/** Internal sentinel for a balance check that exceeded its time budget. */
57+class GateCheckTimeoutError extends Data.TaggedError("GateCheckTimeoutError")<{
58+ readonly timeoutMs: number;
59+}> {}
60+
61+// ---------------------------------------------------------------------------
62+// Shared gate seam — used by this gate and the rate-limit backstop.
63+// ---------------------------------------------------------------------------
64+
65+export type GateDecision =
66+ | { readonly blocked: false }
67+ | { readonly blocked: true; readonly error: { readonly message: string } };
68+
69+/**
70+ * Wrap an engine so `decide` runs before `execute` / `executeWithPause`. A
71+ * blocked decision short-circuits to a descriptive `ExecuteResult.error`
72+ * (which `formatExecuteResult` renders as a clean `isError` MCP tool result)
73+ * WITHOUT invoking the inner engine — so a blocked execution is neither run
74+ * nor usage-tracked. `resume` and all read-only members pass through
75+ * untouched: paused executions must always be able to complete.
76+ */
77+export const withPreExecutionGate = <E extends Cause.YieldableError>(
78+ engine: ExecutionEngine<E>,
79+ decide: Effect.Effect<GateDecision>,
80+): ExecutionEngine<E> => ({
81+ execute: (code, options) =>
82+ Effect.flatMap(
83+ decide,
84+ (decision): Effect.Effect<EngineExecuteResult, E> =>
85+ decision.blocked
86+ ? Effect.succeed({ result: null, error: decision.error.message })
87+ : engine.execute(code, options),
88+ ),
89+ executeWithPause: (code, options) =>
90+ Effect.flatMap(
91+ decide,
92+ (decision): Effect.Effect<ExecutionResult, E> =>
93+ decision.blocked
94+ ? Effect.succeed({
95+ status: "completed",
96+ result: { result: null, error: decision.error.message },
97+ })
98+ : engine.executeWithPause(code, options),
99+ ),
100+ // resume is never gated: paused executions must be able to complete.
101+ resume: (executionId, response) => engine.resume(executionId, response),
102+ getPausedExecution: (executionId) => engine.getPausedExecution(executionId),
103+ pausedExecutionCount: () => engine.pausedExecutionCount(),
104+ hasPausedExecutions: () => engine.hasPausedExecutions(),
105+ getDescription: engine.getDescription,
106+});
107+
108+// ---------------------------------------------------------------------------
109+// Balance gate factory
110+// ---------------------------------------------------------------------------
111+
112+export type ExecutionBalanceCheck = (
113+ organizationId: string,
114+) => Effect.Effect<{ readonly allowed: boolean }, unknown>;
115+
116+/**
117+ * Build a balance gate around `checkBalance` (in production:
118+ * `AutumnService.checkExecutionBalance`). One gate instance holds one
119+ * per-organization outcome cache; the metered decorator layer creates a single
120+ * instance so all engines it decorates share it.
121+ *
122+ * `options.timeoutMs` exists for tests; production uses the default budget.
123+ */
124+export const makeExecutionLimitGate = (
125+ checkBalance: ExecutionBalanceCheck,
126+ options?: { readonly timeoutMs?: number },
127+) => {
128+ const timeoutMs = options?.timeoutMs ?? BALANCE_CHECK_TIMEOUT_MS;
129+ const cache = new Map<string, { readonly allowed: boolean; readonly expiresAtMs: number }>();
130+
131+ const writeCache = (organizationId: string, allowed: boolean, nowMs: number): void => {
132+ if (cache.size >= BALANCE_CACHE_MAX_ENTRIES) {
133+ for (const [key, entry] of cache) {
134+ if (entry.expiresAtMs <= nowMs) cache.delete(key);
135+ }
136+ // Still saturated after dropping expired entries: reset rather than grow.
137+ if (cache.size >= BALANCE_CACHE_MAX_ENTRIES) cache.clear();
138+ }
139+ cache.set(organizationId, { allowed, expiresAtMs: nowMs + BALANCE_CACHE_TTL_MS });
140+ };
141+
142+ const toDecision = (organizationId: string, allowed: boolean): GateDecision =>
143+ allowed
144+ ? { blocked: false }
145+ : {
146+ blocked: true,
147+ error: new ExecutionLimitReachedError({
148+ organizationId,
149+ message: EXECUTION_LIMIT_BLOCKED_MESSAGE,
150+ }),
151+ };
152+
153+ const decide = (organizationId: string): Effect.Effect<GateDecision> =>
154+ Effect.suspend(() => {
155+ const nowMs = Date.now();
156+ const cached = cache.get(organizationId);
157+ if (cached && cached.expiresAtMs > nowMs) {
158+ return Effect.succeed(toDecision(organizationId, cached.allowed));
159+ }
160+ return checkBalance(organizationId).pipe(
161+ Effect.timeoutOrElse({
162+ duration: `${timeoutMs} millis`,
163+ orElse: () => Effect.fail(new GateCheckTimeoutError({ timeoutMs })),
164+ }),
165+ Effect.map(({ allowed }) => {
166+ writeCache(organizationId, allowed, nowMs);
167+ return toDecision(organizationId, allowed);
168+ }),
169+ // FAIL OPEN: Autumn errors, timeouts, and missing customers/features
170+ // must never block executions. Reported like `trackExecution` so a
171+ // billing outage still pages; the error outcome is never cached.
172+ Effect.catch((error: unknown) =>
173+ Effect.sync((): GateDecision => {
174+ console.warn("[billing] execution balance check failed open:", error);
175+ Sentry.captureException(error);
176+ return { blocked: false };
177+ }),
178+ ),
179+ );
180+ });
181+
182+ return {
183+ /**
184+ * Test/diagnostic surface: succeeds when the execution is allowed (or the
185+ * check failed open), fails with the typed `ExecutionLimitReachedError`
186+ * when the organization is out of executions.
187+ */
188+ check: (organizationId: string): Effect.Effect<void, ExecutionLimitReachedError> =>
189+ Effect.flatMap(decide(organizationId), (decision) =>
190+ decision.blocked
191+ ? Effect.fail(
192+ new ExecutionLimitReachedError({
193+ organizationId,
194+ message: EXECUTION_LIMIT_BLOCKED_MESSAGE,
195+ }),
196+ )
197+ : Effect.void,
198+ ),
199+ /** Wrap an engine so the balance gate runs before each new execution. */
200+ decorate: <E extends Cause.YieldableError>(
201+ organizationId: string,
202+ engine: ExecutionEngine<E>,
203+ ): ExecutionEngine<E> => withPreExecutionGate(engine, decide(organizationId)),
204+ };
205+};