* Add execution balance check to billing service * 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. * Add per-org execution rate-limit backstop 1000 execute calls per org per hour, fixed window, counted in a minimal per-org Durable Object (EXECUTION_RATE_LIMITER, one instance per org via idFromName). Independent of billing so runaway automation is caught even when Autumn is down; fails open if the counter DO errors or is slow. The DO stores a single {windowId, count} record and purges itself by alarm after two idle windows. * Wire execution guards into the metered engine decorator Order: rate-limit backstop first (cheap counter), then the balance gate, then usage tracking. Guards wrap outside the tracker so a blocked execution is neither run nor tracked. Covers both planes (HTTP executor plane and MCP session DO) since both build engines through this layer. * Test the execution balance gate and rate-limit backstop Covers allow/block, the typed errors, fail-open on billing errors and timeouts (asserting the check was attempted AND the execution ran), the 60s per-org outcome cache (allowed and blocked cached, errors never), window rollover, per-org isolation, and that resume is never gated. * Use typed test errors and instanceof assertions in gate tests * Replace execution-limit unit tests with e2e coverage Delete the balance-gate and rate-limit unit tests: the guards are now exercised end to end against the real workerd + Autumn emulator, which is a truer test of behaviour that mattered here (blocked executions not reaching the meter, failing open on billing errors). Drop the test-only affordances the unit tests were the only callers of: the gate's `timeoutMs` option knob and both limiters' `check` diagnostic methods. Move the two blocked-execution messages into a small Cloudflare-free module so the e2e suite can import them as assertion ground truth without pulling the guards' worker-runtime dependencies into its typecheck. * Make the per-org execution rate limit env-configurable Add an optional EXECUTION_RATE_LIMIT_PER_HOUR override, parsed as a positive integer and falling back to the 1000/hour default when unset or unparseable. Production leaves it unset; it exists so e2e can drive the abuse backstop with a handful of real executions instead of a thousand. The e2e cloud dev server sets it to 3, reaching the worker the same way ALLOW_LOCAL_NETWORK does. * Add Autumn e2e helpers for fault injection and balance exhaustion Bump the emulator to 0.13.2 for its balances.check support and one-shot fault injection, and extend the Autumn surface to use them: armFault/clearFaults drive the emulator's fault control plane, exhaustExecutions burns an org's whole included allotment in one track (the amount read from the plan seed, never hardcoded), and ledgerFor reads the request ledger so a scenario can prove a faulted balances.check was actually attempted. * Add e2e scenarios for the pre-execution limit guards Four cloud scenarios drive a real MCP client against the workerd + session-DO topology and read Autumn's ledger as ground truth: an org at its cap is blocked before the code runs and not metered; a balance-check 500 and a balance-check timeout each fail open and still run the execution (with the faulted check visible in the ledger as proof the gate consulted Autumn); and the rate-limit backstop blocks the execution past the hourly cap without metering it. Faults are one-shot and cleared in finalizers, and each scenario uses a fresh org. * Raise the e2e execution rate cap above other scenarios' usage The e2e worker's per-org hourly cap of 3 was low enough that ordinary scenarios (toolkits-mcp, no-auth-connection, connect-handoff) tripped the backstop mid-test once the suite ran them against the shared worker. Lift it to 20 (still exhaustible by the backstop scenario) and move the value into one shared constant consumed by both the boot env and the test.
Showing the changes to apps/cloud/src/engine (4 of 15 files). Show all
Unified · Split
| @@ -0,0 +1,185 @@ | |||
| 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 | import { EXECUTION_LIMIT_BLOCKED_MESSAGE } from "./execution-limit-messages"; | ||
| 34 | |||
| 35 | // The engine's completed-result payload (`ExecuteResult` in codemode-core), | ||
| 36 | // derived from the public execution types so this app package doesn't need a | ||
| 37 | // direct dependency on the kernel package that declares it. | ||
| 38 | type EngineExecuteResult = Extract<ExecutionResult, { status: "completed" }>["result"]; | ||
| 39 | |||
| 40 | // One check per org per minute is fresh enough for a monthly quota; both | ||
| 41 | // allowed AND blocked outcomes are cached (errors never are). | ||
| 42 | const BALANCE_CACHE_TTL_MS = 60_000; | ||
| 43 | // Autumn slower than this => fail open rather than stall a user-facing | ||
| 44 | // execution behind the billing provider. | ||
| 45 | const BALANCE_CHECK_TIMEOUT_MS = 2_000; | ||
| 46 | // Sweep guard so a long-lived worker isolate serving many orgs can't grow the | ||
| 47 | // cache map unbounded. | ||
| 48 | const BALANCE_CACHE_MAX_ENTRIES = 10_000; | ||
| 49 | |||
| 50 | export { EXECUTION_LIMIT_BLOCKED_MESSAGE }; | ||
| 51 | |||
| 52 | export class ExecutionLimitReachedError extends Data.TaggedError("ExecutionLimitReachedError")<{ | ||
| 53 | readonly organizationId: string; | ||
| 54 | readonly message: string; | ||
| 55 | }> {} | ||
| 56 | |||
| 57 | /** Internal sentinel for a balance check that exceeded its time budget. */ | ||
| 58 | class GateCheckTimeoutError extends Data.TaggedError("GateCheckTimeoutError")<{ | ||
| 59 | readonly timeoutMs: number; | ||
| 60 | }> {} | ||
| 61 | |||
| 62 | // --------------------------------------------------------------------------- | ||
| 63 | // Shared gate seam — used by this gate and the rate-limit backstop. | ||
| 64 | // --------------------------------------------------------------------------- | ||
| 65 | |||
| 66 | export type GateDecision = | ||
| 67 | | { readonly blocked: false } | ||
| 68 | | { readonly blocked: true; readonly error: { readonly message: string } }; | ||
| 69 | |||
| 70 | /** | ||
| 71 | * Wrap an engine so `decide` runs before `execute` / `executeWithPause`. A | ||
| 72 | * blocked decision short-circuits to a descriptive `ExecuteResult.error` | ||
| 73 | * (which `formatExecuteResult` renders as a clean `isError` MCP tool result) | ||
| 74 | * WITHOUT invoking the inner engine — so a blocked execution is neither run | ||
| 75 | * nor usage-tracked. `resume` and all read-only members pass through | ||
| 76 | * untouched: paused executions must always be able to complete. | ||
| 77 | */ | ||
| 78 | export const withPreExecutionGate = <E extends Cause.YieldableError>( | ||
| 79 | engine: ExecutionEngine<E>, | ||
| 80 | decide: Effect.Effect<GateDecision>, | ||
| 81 | ): ExecutionEngine<E> => ({ | ||
| 82 | execute: (code, options) => | ||
| 83 | Effect.flatMap( | ||
| 84 | decide, | ||
| 85 | (decision): Effect.Effect<EngineExecuteResult, E> => | ||
| 86 | decision.blocked | ||
| 87 | ? Effect.succeed({ result: null, error: decision.error.message }) | ||
| 88 | : engine.execute(code, options), | ||
| 89 | ), | ||
| 90 | executeWithPause: (code, options) => | ||
| 91 | Effect.flatMap( | ||
| 92 | decide, | ||
| 93 | (decision): Effect.Effect<ExecutionResult, E> => | ||
| 94 | decision.blocked | ||
| 95 | ? Effect.succeed({ | ||
| 96 | status: "completed", | ||
| 97 | result: { result: null, error: decision.error.message }, | ||
| 98 | }) | ||
| 99 | : engine.executeWithPause(code, options), | ||
| 100 | ), | ||
| 101 | // resume is never gated: paused executions must be able to complete. | ||
| 102 | resume: (executionId, response) => engine.resume(executionId, response), | ||
| 103 | getPausedExecution: (executionId) => engine.getPausedExecution(executionId), | ||
| 104 | pausedExecutionCount: () => engine.pausedExecutionCount(), | ||
| 105 | hasPausedExecutions: () => engine.hasPausedExecutions(), | ||
| 106 | getDescription: engine.getDescription, | ||
| 107 | }); | ||
| 108 | |||
| 109 | // --------------------------------------------------------------------------- | ||
| 110 | // Balance gate factory | ||
| 111 | // --------------------------------------------------------------------------- | ||
| 112 | |||
| 113 | export type ExecutionBalanceCheck = ( | ||
| 114 | organizationId: string, | ||
| 115 | ) => Effect.Effect<{ readonly allowed: boolean }, unknown>; | ||
| 116 | |||
| 117 | /** | ||
| 118 | * Build a balance gate around `checkBalance` (in production: | ||
| 119 | * `AutumnService.checkExecutionBalance`). One gate instance holds one | ||
| 120 | * per-organization outcome cache; the metered decorator layer creates a single | ||
| 121 | * instance so all engines it decorates share it. | ||
| 122 | */ | ||
| 123 | export const makeExecutionLimitGate = (checkBalance: ExecutionBalanceCheck) => { | ||
| 124 | const timeoutMs = BALANCE_CHECK_TIMEOUT_MS; | ||
| 125 | const cache = new Map<string, { readonly allowed: boolean; readonly expiresAtMs: number }>(); | ||
| 126 | |||
| 127 | const writeCache = (organizationId: string, allowed: boolean, nowMs: number): void => { | ||
| 128 | if (cache.size >= BALANCE_CACHE_MAX_ENTRIES) { | ||
| 129 | for (const [key, entry] of cache) { | ||
| 130 | if (entry.expiresAtMs <= nowMs) cache.delete(key); | ||
| 131 | } | ||
| 132 | // Still saturated after dropping expired entries: reset rather than grow. | ||
| 133 | if (cache.size >= BALANCE_CACHE_MAX_ENTRIES) cache.clear(); | ||
| 134 | } | ||
| 135 | cache.set(organizationId, { allowed, expiresAtMs: nowMs + BALANCE_CACHE_TTL_MS }); | ||
| 136 | }; | ||
| 137 | |||
| 138 | const toDecision = (organizationId: string, allowed: boolean): GateDecision => | ||
| 139 | allowed | ||
| 140 | ? { blocked: false } | ||
| 141 | : { | ||
| 142 | blocked: true, | ||
| 143 | error: new ExecutionLimitReachedError({ | ||
| 144 | organizationId, | ||
| 145 | message: EXECUTION_LIMIT_BLOCKED_MESSAGE, | ||
| 146 | }), | ||
| 147 | }; | ||
| 148 | |||
| 149 | const decide = (organizationId: string): Effect.Effect<GateDecision> => | ||
| 150 | Effect.suspend(() => { | ||
| 151 | const nowMs = Date.now(); | ||
| 152 | const cached = cache.get(organizationId); | ||
| 153 | if (cached && cached.expiresAtMs > nowMs) { | ||
| 154 | return Effect.succeed(toDecision(organizationId, cached.allowed)); | ||
| 155 | } | ||
| 156 | return checkBalance(organizationId).pipe( | ||
| 157 | Effect.timeoutOrElse({ | ||
| 158 | duration: `${timeoutMs} millis`, | ||
| 159 | orElse: () => Effect.fail(new GateCheckTimeoutError({ timeoutMs })), | ||
| 160 | }), | ||
| 161 | Effect.map(({ allowed }) => { | ||
| 162 | writeCache(organizationId, allowed, nowMs); | ||
| 163 | return toDecision(organizationId, allowed); | ||
| 164 | }), | ||
| 165 | // FAIL OPEN: Autumn errors, timeouts, and missing customers/features | ||
| 166 | // must never block executions. Reported like `trackExecution` so a | ||
| 167 | // billing outage still pages; the error outcome is never cached. | ||
| 168 | Effect.catch((error: unknown) => | ||
| 169 | Effect.sync((): GateDecision => { | ||
| 170 | console.warn("[billing] execution balance check failed open:", error); | ||
| 171 | Sentry.captureException(error); | ||
| 172 | return { blocked: false }; | ||
| 173 | }), | ||
| 174 | ), | ||
| 175 | ); | ||
| 176 | }); | ||
| 177 | |||
| 178 | return { | ||
| 179 | /** Wrap an engine so the balance gate runs before each new execution. */ | ||
| 180 | decorate: <E extends Cause.YieldableError>( | ||
| 181 | organizationId: string, | ||
| 182 | engine: ExecutionEngine<E>, | ||
| 183 | ): ExecutionEngine<E> => withPreExecutionGate(engine, decide(organizationId)), | ||
| 184 | }; | ||
| 185 | }; | ||
| @@ -0,0 +1,11 @@ | |||
| 1 | // The user-facing messages a blocked execution surfaces to the MCP client, kept | ||
| 2 | // in their own Cloudflare-free module so the e2e suite can import them as | ||
| 3 | // ground truth without pulling the guards' `cloudflare:workers` dependencies | ||
| 4 | // into its (worker-types-free) typecheck. The guard modules re-export these, so | ||
| 5 | // production call sites are unchanged. | ||
| 6 | |||
| 7 | export const EXECUTION_LIMIT_BLOCKED_MESSAGE = | ||
| 8 | "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."; | ||
| 9 | |||
| 10 | export const RATE_LIMIT_BLOCKED_MESSAGE = | ||
| 11 | "Rate limit exceeded: too many executions this hour. This is an abuse backstop — contact support if you hit this legitimately."; | ||
| @@ -0,0 +1,220 @@ | |||
| 1 | // --------------------------------------------------------------------------- | ||
| 2 | // Per-org execution rate limit — an abuse backstop independent of billing. | ||
| 3 | // | ||
| 4 | // The balance gate (execution-gate.ts) depends on Autumn and fails open, so a | ||
| 5 | // billing outage plus runaway automation could still run unbounded executions. | ||
| 6 | // This limiter counts `execute` calls per organization in a fixed hourly | ||
| 7 | // window, backed by a minimal counter Durable Object (cross-session state: | ||
| 8 | // each MCP session lives in its own DO instance, so an in-memory counter | ||
| 9 | // would be per-session and trivially bypassed by opening more sessions). | ||
| 10 | // | ||
| 11 | // Like the balance gate it FAILS OPEN: an unreachable counter DO, a missing | ||
| 12 | // binding, or a slow call allows the execution (warn + Sentry). The backstop | ||
| 13 | // must never take executions down with it. | ||
| 14 | // --------------------------------------------------------------------------- | ||
| 15 | |||
| 16 | import { DurableObject, env } from "cloudflare:workers"; | ||
| 17 | import * as Sentry from "@sentry/cloudflare"; | ||
| 18 | import { Data, Effect } from "effect"; | ||
| 19 | import type * as Cause from "effect/Cause"; | ||
| 20 | |||
| 21 | import type { ExecutionEngine } from "@executor-js/execution"; | ||
| 22 | |||
| 23 | import { withPreExecutionGate, type GateDecision } from "./execution-gate"; | ||
| 24 | import { RATE_LIMIT_BLOCKED_MESSAGE } from "./execution-limit-messages"; | ||
| 25 | |||
| 26 | // Fixed window: all executions in the same clock hour share one counter. | ||
| 27 | export const RATE_LIMIT_WINDOW_MS = 3_600_000; | ||
| 28 | // Calibration: the heaviest legitimate org runs ~1.1k executions per MONTH, | ||
| 29 | // so 1000 per HOUR is far above any human-driven usage and only trips on | ||
| 30 | // runaway automation (the incident this backstops: ~18k in 30 days would | ||
| 31 | // still pass, which is fine — that class of overrun is the balance gate's | ||
| 32 | // job; this catches tight loops). | ||
| 33 | export const EXECUTIONS_PER_ORG_PER_HOUR = 1000; | ||
| 34 | // Counter DO slower than this => fail open rather than stall executions. | ||
| 35 | const RATE_LIMIT_CHECK_TIMEOUT_MS = 2_000; | ||
| 36 | // The DO purges its storage this long after the last increment, so idle orgs | ||
| 37 | // cost nothing. Two windows: long enough that an active window never purges. | ||
| 38 | const COUNTER_PURGE_AFTER_MS = 2 * RATE_LIMIT_WINDOW_MS; | ||
| 39 | |||
| 40 | export { RATE_LIMIT_BLOCKED_MESSAGE }; | ||
| 41 | |||
| 42 | export class ExecutionRateLimitExceededError extends Data.TaggedError( | ||
| 43 | "ExecutionRateLimitExceededError", | ||
| 44 | )<{ | ||
| 45 | readonly organizationId: string; | ||
| 46 | readonly message: string; | ||
| 47 | }> {} | ||
| 48 | |||
| 49 | /** Internal sentinel for a counter call that exceeded its time budget. */ | ||
| 50 | class RateLimitCheckTimeoutError extends Data.TaggedError("RateLimitCheckTimeoutError")<{ | ||
| 51 | readonly timeoutMs: number; | ||
| 52 | }> {} | ||
| 53 | |||
| 54 | // --------------------------------------------------------------------------- | ||
| 55 | // Counter Durable Object — one instance per organization (idFromName(orgId)). | ||
| 56 | // Stores a single { windowId, count } record: an increment in a new window | ||
| 57 | // resets the count, so old windows never accumulate. An alarm purges storage | ||
| 58 | // after inactivity. | ||
| 59 | // --------------------------------------------------------------------------- | ||
| 60 | |||
| 61 | const WINDOW_RECORD_KEY = "window"; | ||
| 62 | |||
| 63 | type WindowRecord = { | ||
| 64 | readonly windowId: number; | ||
| 65 | readonly count: number; | ||
| 66 | }; | ||
| 67 | |||
| 68 | export class ExecutionRateLimiterDO extends DurableObject { | ||
| 69 | private readonly counterStorage: DurableObjectState["storage"]; | ||
| 70 | |||
| 71 | constructor(ctx: DurableObjectState, doEnv: Env) { | ||
| 72 | super(ctx, doEnv); | ||
| 73 | // Kept on an own field (not just inherited `this.ctx`) so tests can run | ||
| 74 | // the class against a fake storage under the `cloudflare:workers` stub. | ||
| 75 | this.counterStorage = ctx.storage; | ||
| 76 | } | ||
| 77 | |||
| 78 | /** Add one execution to `windowId`'s counter and return the new count. */ | ||
| 79 | async increment(windowId: number): Promise<number> { | ||
| 80 | const stored = await this.counterStorage.get<WindowRecord>(WINDOW_RECORD_KEY); | ||
| 81 | const count = stored && stored.windowId === windowId ? stored.count + 1 : 1; | ||
| 82 | await this.counterStorage.put(WINDOW_RECORD_KEY, { windowId, count }); | ||
| 83 | await this.counterStorage.setAlarm(Date.now() + COUNTER_PURGE_AFTER_MS); | ||
| 84 | return count; | ||
| 85 | } | ||
| 86 | |||
| 87 | async alarm(): Promise<void> { | ||
| 88 | await this.counterStorage.deleteAll(); | ||
| 89 | } | ||
| 90 | } | ||
| 91 | |||
| 92 | // --------------------------------------------------------------------------- | ||
| 93 | // Client | ||
| 94 | // --------------------------------------------------------------------------- | ||
| 95 | |||
| 96 | /** Count one execution for (organizationId, windowId); returns the new count. */ | ||
| 97 | export type RateLimitIncrement = ( | ||
| 98 | organizationId: string, | ||
| 99 | windowId: number, | ||
| 100 | ) => Effect.Effect<number, unknown>; | ||
| 101 | |||
| 102 | export type ExecutionRateLimiter = { | ||
| 103 | readonly decorate: <E extends Cause.YieldableError>( | ||
| 104 | organizationId: string, | ||
| 105 | engine: ExecutionEngine<E>, | ||
| 106 | ) => ExecutionEngine<E>; | ||
| 107 | }; | ||
| 108 | |||
| 109 | /** | ||
| 110 | * Build a rate limiter around an increment function (in production: the | ||
| 111 | * counter DO). `options.limit` is the per-org hourly cap (production sets it | ||
| 112 | * from the env override in `makeCloudExecutionRateLimiter`); the rest tune the | ||
| 113 | * window and time budget. | ||
| 114 | */ | ||
| 115 | export const makeExecutionRateLimiter = ( | ||
| 116 | increment: RateLimitIncrement, | ||
| 117 | options?: { | ||
| 118 | readonly limit?: number; | ||
| 119 | readonly windowMs?: number; | ||
| 120 | readonly timeoutMs?: number; | ||
| 121 | readonly now?: () => number; | ||
| 122 | }, | ||
| 123 | ): ExecutionRateLimiter => { | ||
| 124 | const limit = options?.limit ?? EXECUTIONS_PER_ORG_PER_HOUR; | ||
| 125 | const windowMs = options?.windowMs ?? RATE_LIMIT_WINDOW_MS; | ||
| 126 | const timeoutMs = options?.timeoutMs ?? RATE_LIMIT_CHECK_TIMEOUT_MS; | ||
| 127 | const now = options?.now ?? Date.now; | ||
| 128 | |||
| 129 | const decide = (organizationId: string): Effect.Effect<GateDecision> => | ||
| 130 | Effect.suspend(() => { | ||
| 131 | const windowId = Math.floor(now() / windowMs); | ||
| 132 | return increment(organizationId, windowId).pipe( | ||
| 133 | Effect.timeoutOrElse({ | ||
| 134 | duration: `${timeoutMs} millis`, | ||
| 135 | orElse: () => Effect.fail(new RateLimitCheckTimeoutError({ timeoutMs })), | ||
| 136 | }), | ||
| 137 | Effect.map( | ||
| 138 | (count): GateDecision => | ||
| 139 | count > limit | ||
| 140 | ? { | ||
| 141 | blocked: true, | ||
| 142 | error: new ExecutionRateLimitExceededError({ | ||
| 143 | organizationId, | ||
| 144 | message: RATE_LIMIT_BLOCKED_MESSAGE, | ||
| 145 | }), | ||
| 146 | } | ||
| 147 | : { blocked: false }, | ||
| 148 | ), | ||
| 149 | // FAIL OPEN: the backstop must never block executions because its | ||
| 150 | // counter is unreachable or slow. | ||
| 151 | Effect.catch((error: unknown) => | ||
| 152 | Effect.sync((): GateDecision => { | ||
| 153 | console.warn("[rate-limit] execution rate limit check failed open:", error); | ||
| 154 | Sentry.captureException(error); | ||
| 155 | return { blocked: false }; | ||
| 156 | }), | ||
| 157 | ), | ||
| 158 | ); | ||
| 159 | }); | ||
| 160 | |||
| 161 | return { | ||
| 162 | decorate: (organizationId, engine) => withPreExecutionGate(engine, decide(organizationId)), | ||
| 163 | }; | ||
| 164 | }; | ||
| 165 | |||
| 166 | // --------------------------------------------------------------------------- | ||
| 167 | // Cloud wiring — reads the EXECUTION_RATE_LIMITER binding from the worker env. | ||
| 168 | // --------------------------------------------------------------------------- | ||
| 169 | |||
| 170 | // The DO stub's RPC surface. The binding is declared untyped in | ||
| 171 | // env-augment.d.ts (matching the BLOBS precedent), so the call site narrows it | ||
| 172 | // to the one method the class exposes. | ||
| 173 | type ExecutionRateLimiterStub = { | ||
| 174 | readonly increment: (windowId: number) => Promise<number>; | ||
| 175 | }; | ||
| 176 | |||
| 177 | type RateLimiterNamespace = { | ||
| 178 | readonly idFromName: (name: string) => DurableObjectId; | ||
| 179 | readonly get: (id: DurableObjectId) => unknown; | ||
| 180 | }; | ||
| 181 | |||
| 182 | /** | ||
| 183 | * Production rate limiter backed by the `EXECUTION_RATE_LIMITER` counter DO. | ||
| 184 | * When the binding is absent (unit-test workers, older local setups) the | ||
| 185 | * limiter is disabled: every check passes, logged once at construction. | ||
| 186 | */ | ||
| 187 | export const makeCloudExecutionRateLimiter = (): ExecutionRateLimiter => { | ||
| 188 | const limit = resolveRateLimit(); | ||
| 189 | const namespace = (env as { EXECUTION_RATE_LIMITER?: RateLimiterNamespace }) | ||
| 190 | .EXECUTION_RATE_LIMITER; | ||
| 191 | if (!namespace) { | ||
| 192 | console.warn( | ||
| 193 | "[rate-limit] EXECUTION_RATE_LIMITER binding missing; execution rate limiting disabled", | ||
| 194 | ); | ||
| 195 | return makeExecutionRateLimiter(() => Effect.succeed(0)); | ||
| 196 | } | ||
| 197 | return makeExecutionRateLimiter( | ||
| 198 | (organizationId, windowId) => | ||
| 199 | Effect.tryPromise(() => { | ||
| 200 | const stub = namespace.get( | ||
| 201 | namespace.idFromName(organizationId), | ||
| 202 | ) as ExecutionRateLimiterStub; | ||
| 203 | return stub.increment(windowId); | ||
| 204 | }), | ||
| 205 | { limit }, | ||
| 206 | ); | ||
| 207 | }; | ||
| 208 | |||
| 209 | /** | ||
| 210 | * The per-org hourly cap: the `EXECUTION_RATE_LIMIT_PER_HOUR` env override | ||
| 211 | * (parsed as a positive integer) or `EXECUTIONS_PER_ORG_PER_HOUR` when it's | ||
| 212 | * unset or unparseable. The override exists so e2e can drive the backstop with | ||
| 213 | * a small number of real executions; production leaves the var unset. | ||
| 214 | */ | ||
| 215 | const resolveRateLimit = (): number => { | ||
| 216 | const raw = (env as { EXECUTION_RATE_LIMIT_PER_HOUR?: string }).EXECUTION_RATE_LIMIT_PER_HOUR; | ||
| 217 | if (raw === undefined) return EXECUTIONS_PER_ORG_PER_HOUR; | ||
| 218 | const parsed = Number.parseInt(raw, 10); | ||
| 219 | return Number.isInteger(parsed) && parsed > 0 ? parsed : EXECUTIONS_PER_ORG_PER_HOUR; | ||
| 220 | }; | ||
| @@ -28,19 +28,40 @@ import { · 27 unmodified lines | |||
| 28 | import { AutumnService } from "../extensions/billing/service"; | 28 | import { AutumnService } from "../extensions/billing/service"; |
| 29 | import type { DbService } from "../db/db"; | 29 | import type { DbService } from "../db/db"; |
| 30 | import { CloudExecutionSeamsLayer } from "../engine/execution-stack"; | 30 | import { CloudExecutionSeamsLayer } from "../engine/execution-stack"; |
| 31 | import { makeExecutionLimitGate } from "./execution-gate"; | ||
| 32 | import { makeCloudExecutionRateLimiter } from "./execution-rate-limit"; | ||
| 31 | import { withExecutionUsageTracking } from "./execution-usage"; | 33 | import { withExecutionUsageTracking } from "./execution-usage"; |
| 32 | 34 | ||
| 33 | // Usage-metering decorator bound to the billing service. `trackExecution` is | 35 | // Usage-metering decorator bound to the billing service, plus the two |
| 34 | // fire-and-forget (`Effect.runFork`) so the billing call can't stall a | 36 | // pre-execution guards this layer owns, ordered cheapest first: |
| 35 | // user-facing execution. | 37 | // 1. rate-limit backstop (counter DO, independent of billing) |
| 38 | // 2. execution balance gate (Autumn check, cached 60s, fails open) | ||
| 39 | // 3. usage tracking — fire-and-forget (`Effect.runFork`) so the billing | ||
| 40 | // call can't stall a user-facing execution. | ||
| 41 | // The guards wrap OUTSIDE the tracker so a blocked execution is neither run | ||
| 42 | // nor tracked. One gate/limiter instance per layer build: the balance cache is | ||
| 43 | // shared by every engine this decorator produces (in the MCP session DO that's | ||
| 44 | // the session's engines; on the HTTP plane the boot layer builds it once). | ||
| 36 | export const CloudMeteringEngineDecorator: Layer.Layer<EngineDecorator, never, AutumnService> = | 45 | export const CloudMeteringEngineDecorator: Layer.Layer<EngineDecorator, never, AutumnService> = |
| 37 | Layer.effect(EngineDecorator)( | 46 | Layer.effect(EngineDecorator)( |
| 38 | Effect.map(AutumnService.asEffect(), (autumn): EngineDecorator["Service"] => ({ | 47 | Effect.map(AutumnService.asEffect(), (autumn): EngineDecorator["Service"] => { |
| 39 | decorate: (engine, identity: EngineStackIdentity) => | 48 | const balanceGate = makeExecutionLimitGate((organizationId) => |
| 40 | withExecutionUsageTracking(identity.organizationId, engine, (organizationId) => | 49 | autumn.checkExecutionBalance(organizationId), |
| 41 | Effect.runFork(autumn.trackExecution(organizationId)), | 50 | ); |
| 42 | ), | 51 | const rateLimiter = makeCloudExecutionRateLimiter(); |
| 43 | })), | 52 | return { |
| 53 | decorate: (engine, identity: EngineStackIdentity) => | ||
| 54 | rateLimiter.decorate( | ||
| 55 | identity.organizationId, | ||
| 56 | balanceGate.decorate( | ||
| 57 | identity.organizationId, | ||
| 58 | withExecutionUsageTracking(identity.organizationId, engine, (organizationId) => | ||
| 59 | Effect.runFork(autumn.trackExecution(organizationId)), | ||
| 60 | ), | ||
| 61 | ), | ||
| 62 | ), | ||
| 63 | }; | ||
| 64 | }), | ||
| 44 | ); | 65 | ); |
| 45 | 66 | ||
| 46 | /** | 67 | /** |