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.
Unified · Split
| @@ -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 | }; | ||