diff --git a/README.md b/README.md index 87b72c2f..a88f2413 100644 --- a/README.md +++ b/README.md @@ -440,6 +440,34 @@ config = SchematicConfig(flag_defaults={"some-flag-key": True}) client = Schematic("YOUR_API_KEY", config) ``` +You can also preflight a check: ask whether an action *would* be allowed, by +evaluating the flag against usage the action has not recorded yet. Pass +`usage` for any numeric condition, or `event_usage` when you know which event +subtype the usage lands on: + +```python +from schematic.client import CheckFlagOptions, EventUsage, Schematic + +client = Schematic("YOUR_API_KEY") + +allowed = client.check_flag( + "some-flag-key", + company={"id": "your-company-id"}, + options=CheckFlagOptions(usage=1000), +) + +allowed = client.check_flag( + "some-flag-key", + company={"id": "your-company-id"}, + options=CheckFlagOptions( + event_usage=EventUsage(event_subtype="inference_tokens", quantity=1000), + ), +) +``` + +Preflighted checks are never served from, or written to, the local flag check +cache: the answer is specific to the usage you simulated. + ### Offline Mode In development or testing environments, you may want to avoid making network requests to the Schematic API. You can run Schematic in offline mode by specifying the `offline` option; in this case, it does not matter what API key you specify: @@ -538,11 +566,170 @@ client = Schematic( ) ``` +## Credit reservations + +For features metered by credit burndown, such as inference tokens, `check()` +holds credits for the work you are about to do and `track_with_reservation()` +settles the hold with the actual usage. A *lease* is a tranche of credits the +SDK draws from the server up front; a *reservation* is one hold carved out of +it, sized to the upper bound of a single operation. In server mode there is no +lease: the server evaluates the flag and takes the hold in the same call. + +> Client mode needs [DataStream](#datastream), and on a multi-process +> deployment a shared Redis, so every process gates against the same lease +> balance. Without DataStream the SDK uses server mode. It is async only, like +> DataStream itself; the synchronous `Schematic` client always uses server mode. + +### Setup + +```python +import redis.asyncio as aioredis +from schematic.client import AsyncSchematic, AsyncSchematicConfig, CreditLeaseConfig + +redis_client = aioredis.from_url("redis://localhost:6379") + +config = AsyncSchematicConfig( + use_datastream=True, + credit_leases=CreditLeaseConfig( + redis_client=redis_client, # shared lease state; omit it to gate within one process + default_lease_size=10_000, # credits drawn per lease + default_lease_duration=300.0, # seconds a lease lives + default_reservation_ttl=60.0, # seconds a hold survives unsettled + ), +) +client = AsyncSchematic("YOUR_API_KEY", config) +``` + +`mode` defaults to `auto`, which picks client mode when DataStream is running +and server mode otherwise. When `datastream.company_cache` is a `RedisCache`, +its client backs lease state automatically, so `redis_client` only needs +setting to point leases at a different Redis. Without either, lease state stays +per-process and gates within that process alone, which the SDK warns about at +startup. + +### Server mode + +Server mode takes every hold over the API: one `check-and-reserve` call +evaluates the flag and holds the credits, and the settling event carries the +reservation ID. No lease, no Redis, no local state. It suits low-volume checks; +client mode suits high-throughput gating. Only `mode` and +`default_reservation_ttl` apply, and the SDK warns at startup when a +client-only option is set. + +```python +from schematic.client import CreditLeaseConfig, Schematic, SchematicConfig + +config = SchematicConfig( + credit_leases=CreditLeaseConfig( + default_reservation_ttl=60.0, # seconds the server holds credits, max 1 hour + ), +) +client = Schematic("YOUR_API_KEY", config) +``` + +### Checking and tracking + +Reserve the operation's upper bound, do the work, and report what it actually +used. The unspent slice comes back to the balance. + +```python +from schematic.client import CheckOptions + +result = await client.check( + "inference", + company={"id": "your-company-id"}, + options=CheckOptions( + usage=1000, # upper bound for this operation + event_subtype="inference_tokens", # the metered event + ), +) +if not result.allowed: + raise RuntimeError("credit balance exceeded") + +inference = await run_inference() + +# A check can allow without holding anything, for instance when the feature is +# not metered by credits, and that usage still has to be tracked. +if result.reservation is not None: + await client.track_with_reservation(result.reservation, inference.tokens_used) +else: + await client.track( + "inference_tokens", + company={"id": "your-company-id"}, + quantity=inference.tokens_used, + ) +``` + +If nothing settles a reservation, its hold is refunded at +`default_reservation_ttl`. A settle arriving after that still bills the server, +but no longer re-debits the local lease, so the local balance reads high until +the lease rolls over: size the TTL above the longest expected gap between +`check()` and `track_with_reservation()`. The settling event carries an +idempotency key derived from the reservation ID, so a retried or duplicated +settle is billed once. + +### Pre-warming + +A session's first check pays the lease acquire round trip. Warm the lease when +the user is identified instead: + +```python +from schematic.client import IdentifyOptions +from schematic.types import EventBodyIdentifyCompany + +await client.identify( + {"id": "your-user-id"}, + company=EventBodyIdentifyCompany(keys={"id": "your-company-id"}), + options=IdentifyOptions(prewarm=["credit-type-id"]), +) +``` + +Or call `await client.prewarm({"id": "your-company-id"}, ["credit-type-id"])` +directly. Both are no-ops in server mode, and neither raises. + +### Configuration + +All fields live on `CreditLeaseConfig`. Durations are seconds. Everything below +`default_reservation_ttl` steers client mode only. + +| Option | Default | Meaning | +|---|---|---| +| `mode` | `auto` | Where the hold lives: `client`, `server`, or `auto` (client when DataStream is running). | +| `default_reservation_ttl` | 60 | How long a hold survives unsettled. Capped at one hour in server mode; client mode keeps whatever you set, since the TTL only drives the local sweeper there. | +| `default_lease_duration` | 300 | Lease lifetime requested at acquire and extend. | +| `default_lease_size` | 10000 | Credits requested per acquire, and the minimum extend tranche. | +| `low_water_mark` | 0.25 | Remaining/granted ratio at or below which a background extend fires. | +| `sweep_interval` | 1 | How often expired holds are swept back to their leases. | +| `prewarm_resolve_timeout` | 5 | How long `prewarm` waits for a freshly identified company to surface. 0 skips the wait, and still warms a company already in the DataStream cache. | +| `redis_client` | the DataStream cache's client | Connected `redis.asyncio` client for lease and reservation state. | +| `redis_key_prefix` | `"schematic:"` | Key prefix for lease and reservation keys. Matches the Node SDK, so mixed fleets share leases. | +| `overrides` | none | Per-credit-type overrides of the four knobs above, keyed by credit type ID. | + +### When a check cannot gate + +A check that cannot gate, because the API is unreachable, Redis is down, or the +lease is exhausted, fails closed by default: `allowed` is False and no hold is +taken. Pass `on_acquire_failure="fail-open"` where letting traffic through +beats denying it. + +Fail-open does not skip the evaluation in client mode: the flag's rules still +run with the credit balance assumed sufficient, so plan targeting, overrides, +and every non-credit condition still apply, and a company that is not entitled +stays denied. Server mode has no local engine to re-run, so it returns your +default value (`CheckOptions.default_value`, else the client's flag default). +That default is False unless you set one, so fail-open in server mode denies +until you pass `default_value` or register a flag default. + +A 402 is different: the server knows the credits are not there, so the check +denies whatever `on_acquire_failure` says. + ## DataStream DataStream enables local flag evaluation by maintaining a WebSocket connection to Schematic and caching flag rules, company, and user data locally (or in a shared cache such as Redis). Flag checks are evaluated locally via a WASM rules engine, eliminating per-check network requests. -> **Async-only:** DataStream and Replicator Mode are only available on the `AsyncSchematic` client. The synchronous `Schematic` client does not support either feature — use `AsyncSchematic` (shown in all examples below) if you need them. +It also unlocks client-mode [credit reservations](#credit-reservations), where credit holds are carved out of a local lease instead of costing an API call per check. + +> **Async-only:** DataStream and Replicator Mode are only available on the `AsyncSchematic` client. The synchronous `Schematic` client does not support either feature; use `AsyncSchematic` (shown in all examples below) if you need them. ### Installation @@ -640,7 +827,7 @@ async def main(): asyncio.run(main()) ``` -`RedisCache` accepts a `prefix` argument (default `"schematic"`) if you need to namespace keys — this must match the prefix used by any other SDKs or the replicator writing to the same Redis instance. +`RedisCache` accepts a `prefix` argument (default `"schematic"`) if you need to namespace keys. It must match the prefix used by any other SDKs or the replicator writing to the same Redis instance. ### Replicator Mode diff --git a/conformance/README.md b/conformance/README.md new file mode 100644 index 00000000..0b0f4060 --- /dev/null +++ b/conformance/README.md @@ -0,0 +1,10 @@ +# Credit lease conformance suite + +`SPEC.md` and `vectors/*.json` are copied verbatim from schematic-node +(`conformance/` on `main`), the reference implementation for client-mode credit +leases. Do not edit them here: fix or extend them in schematic-node and copy the +result back, or the SDKs stop pinning the same behavior. + +`tests/conformance/test_vectors.py` is this repo's runner. The runner is the +only language-specific piece; every SDK reimplements it and must pass the same +vectors. diff --git a/conformance/SPEC.md b/conformance/SPEC.md new file mode 100644 index 00000000..6191f101 --- /dev/null +++ b/conformance/SPEC.md @@ -0,0 +1,446 @@ +# Credit lease & reservation semantics — conformance spec + +This document specifies the client-side credit lease/reservation semantics implemented by the +Schematic Node SDK (the reference implementation), in enough detail to reimplement them in another +language without reading the Node source. The machine-readable test vectors in +`conformance/vectors/*.json` pin the observable behavior; this spec explains the model, the +configuration knobs, and the invariants that cannot be expressed as deterministic vectors. + +Where this document and the vectors disagree, the vectors win — they are generated from the +reference implementation's behavior. + +- [Vector format](#vector-format) +- [Model overview](#model-overview) +- [State](#state) +- [Store operations](#store-operations) +- [Lease manager](#lease-manager) +- [Check flow](#check-flow) +- [Track / settle flow](#track--settle-flow) +- [Configuration knobs](#configuration-knobs) +- [Bounded-leak contract](#bounded-leak-contract) +- [Invariants not expressible as vectors](#invariants-not-expressible-as-vectors) + +## Vector format + +Each file in `conformance/vectors/` is a JSON document: + +```json +{ + "category": "reservation_lifecycle", + "vectors": [ + { + "name": "unique_snake_case_name", + "description": "What this vector pins and why.", + "backends": ["in_memory", "redis"], + "given": { + "config": { "lease_duration_ms": 300000, "reservation_ttl_ms": 60000, "lease_size": 1000, "low_water_mark": 0.25 }, + "leases": [ { "lease_id": "lse_1", "company_id": "co_1", "credit_type_id": "ct_1", "granted_amount": 1000, "expires_at_ms": 60000 } ] + }, + "operations": [ + { "op": "try_reserve", "company_id": "co_1", "credit_type_id": "ct_1", "credits": 100, "expect": { "balance": 900 } } + ] + } + ] +} +``` + +Rules: + +- All keys are `snake_case`. Vectors are plain JSON — no language-specific types. +- **Virtual clock.** The run starts at a fixed virtual instant `t0`. Every `*_at_ms` field is an + offset in milliseconds from `t0` (an absolute position on the virtual timeline, not relative to + the current operation). The `advance_clock` operation moves the clock forward; nothing else does. + Runners must execute vectors against a controllable clock (no wall time). +- `backends` restricts which store backends the vector runs against; when omitted, the vector must + pass against every backend the SDK ships (in-memory and Redis for Node). +- `given.leases` are installed via the store's `replace` operation at `t0` (each install must + return "written"). +- Assertions are attached per-operation via `expect`. Final-state assertions are expressed as + trailing read operations (`get_lease`, `reserved_credits`, `reservation_count`). +- `expect.balance` / `expect.consumed` use JSON `null` for the "no / refused" result. +- Reservation ids created by `check` operations are random; the vector names them via + `save_reservation_as` and later operations reference them with `handle`. + +### Operations + +Store-level (exercise the lease store and reservation store directly): + +| op | fields | expect | +| --- | --- | --- | +| `advance_clock` | `ms` | — | +| `replace_lease` | `lease_id`, `company_id`, `credit_type_id`, `granted_amount`, `expires_at_ms` | `written` (bool) | +| `drop_lease` | `company_id`, `credit_type_id` | — | +| `try_reserve` | `company_id`, `credit_type_id`, `credits` | `balance` (post-debit number, or `null`) | +| `refund_lease` | `company_id`, `credit_type_id`, `credits`, `pin_lease_id`? | — | +| `extend_lease` | `company_id`, `credit_type_id`, `granted_total`, `expires_at_ms`?, `pin_lease_id`? | — | +| `get_lease` | `company_id`, `credit_type_id` | `exists`, `lease_id`?, `granted_amount`?, `local_remaining_credits`? | +| `add_reservation` | `id`, `lease_id`, `company_id`, `credit_type_id`, `event_subtype`, `quantity_reserved`, `credits_reserved`, `consumption_rate`, `expires_at_ms` | — | +| `consume_reservation` | `id` or `handle`, `credits`, `crash_before_refund`? (bool, one-shot) | `consumed` (number or `null`), `throws`? | +| `get_reservation` | `id` or `handle` | `exists` | +| `reserved_credits` | `company_id`, `credit_type_id` | `total` | +| `reservation_count` | — | `count` | +| `sweep_expired` | — | `swept` | + +Manager-level (exercise the lease manager with a scripted wire client): + +| op | fields | expect | +| --- | --- | --- | +| `acquire_if_needed` | `company_id`, `credit_type_id`, `server`? ( `{ "lease": {...} }` or `{ "error": "..." }` ), `install_during_wire`? (lease installed into the store while the wire call is in flight, emulating a sibling pod winning the race) | `lease_id` (or `null`), `wire_acquires` (cumulative count), `last_acquire_requested_amount`?, `released_lease_ids` (cumulative) | +| `maybe_extend` | `company_id`, `credit_type_id`, `required_credits`?, `server`? ( `{ "lease": { "granted_total", "expires_at_ms" } }` or `{ "error": "..." }` ) | `wire_extends` (cumulative count), `last_extend_additional_amount`?, `last_extend_lease_id`? | +| `release_all_local_leases` | — (in-memory backend only) | `released_lease_ids`, `remaining_slots` | + +Flow-level (exercise the full check/track orchestration with a scripted rules engine): + +| op | fields | expect | +| --- | --- | --- | +| `check` | `flag_key`, `company` (`{ id, credit_balances }`), `usage`, `event_subtype`?, `on_acquire_failure`?, `engine` (array of scripted engine results, consumed in call order), `server`? (as above), `save_reservation_as`? | `allowed`, `reason`?, `err`?, `has_reservation`, `reservation`? (field subset), `fallback_called`?, `engine_calls`? (per-call `{ credit_balance, credit_cost?, event_usage? }`; `credit_balance` may be the string `"max_safe_integer"`) | +| `track` | `handle`, `actual_quantity` | `settled_locally`, `track` (`{ event, quantity, lease_id }`) | + +A scripted engine result is `{ "value": bool, "reason"?: string, "entitlement"?: { "value_type", +"credit_id"?, "consumption_rate"?, "event_subtype"?, "feature_id"?, "feature_key"? } }`. The engine +is an oracle: the vectors pin the *orchestration around* the rules engine (what it is called with, +and what the SDK does with its answer), not the engine itself — the engine is shared WASM across +SDKs and has its own tests. + +## Model overview + +Credit-metered features are gated client-side without a wire call per check. The SDK: + +1. **Leases** a tranche of credits from the server per `(company_id, credit_type_id)`. The server + pre-debits the company balance by the granted amount; the SDK tracks a local view of how much + of the tranche remains un-reserved (`local_remaining_credits`). +2. **Reserves** `usage x consumption_rate` credits from the lease at `check()` time, atomically + (check-and-debit). A successful, engine-approved check returns a *reservation handle*. +3. **Settles** the reservation at `track()` time with the actual usage: the actually-consumed + credits stay debited, the unspent slice is refunded to the lease, and a Track event bills the + server (the server is the source of truth for real consumption). + +Everything client-side is *local bookkeeping against the leased tranche*. The server reconciles: +an expired lease's unspent remainder is refunded to the company balance server-side, and Track +events (keyed by `lease_id`) drive the authoritative consumption. + +Leases and reservations both expire: + +- A **lease** past its expiry must be treated as *released* — its local balance is stale (the + server already refunded the remainder) and must never serve another reserve or be extended. +- A **reservation** past its TTL is swept: removed from the table and its full hold refunded to + the lease. Work that finishes after the sweep still bills the server (recovery emit) but does + not re-debit the local lease. + +## State + +**Lease slot** — at most one lease per `(company_id, credit_type_id)` key: + +| field | meaning | +| --- | --- | +| `lease_id` | Server-issued id. | +| `granted_amount` | Server-authoritative total granted to this lease (grows on extend). | +| `local_remaining_credits` | Local view: granted minus outstanding holds/consumption. Initialized to `granted_amount` on install. | +| `expires_at` | Expiry instant. Past it the lease is dead (see above). | + +**Reservation** — keyed by a unique id: + +| field | meaning | +| --- | --- | +| `id` | Unique (UUID in Node). | +| `lease_id` | The lease the hold was carved from. Pins refunds. | +| `company_id`, `credit_type_id` | Slot key. | +| `event_subtype` | Event the settle will bill as. | +| `quantity_reserved` | Caller-declared usage (event units). | +| `credits_reserved` | `quantity_reserved x consumption_rate`. | +| `consumption_rate` | Rate at reservation time. | +| `expires_at` | Reservation TTL deadline (sweep target). | +| `eval_ctx` | Company/user keys used at check time; threaded onto the Track event. | + +## Store operations + +These are the primitives both store backends (per-process in-memory; shared Redis) must implement +with identical observable semantics. Each mutation must be atomic per slot/reservation (see +[Invariants](#invariants-not-expressible-as-vectors)). + +### `replace(lease)` — install-if-not-live + +Install a fresh lease with `local_remaining_credits = granted_amount`, **only if** the slot is +empty or the existing lease is expired *and carries a different `lease_id`*. If a *live* lease +occupies the slot — even with a different `lease_id` (a sibling pod won the race) — leave it +untouched (its already-debited balance wins) and report "kept". If an *expired* lease with the +**same** `lease_id` occupies the slot (a stale acquire response for a lease the idempotent server +also handed to a racing sibling, which may since have extended it), do not rewrite it either: +rewriting would reset `local_remaining_credits` to the full grant and erase debits whose +reservations are still open. Reconcile it like `extend` instead — granted to the incoming total +(lower/equal totals are no-ops), expiry only forward, balance untouched — and report "kept". +Returns written/kept so the caller can run the redundant-lease release logic (see manager). + +### `try_reserve(company, credit, credits)` — atomic check-and-debit + +- Reject (return `null`, touch nothing) if: no lease in the slot; the lease is **expired**; the + remaining balance is `< credits`; or `credits` is not a finite non-negative number (NaN must + never reach the arithmetic — it slips through every comparison and would poison the balance + into approving everything). +- Otherwise debit and return the **post-debit balance** (so the caller can derive the pre-debit + figure as `returned + credits` without a racy follow-up read). +- Reserving down to exactly 0 is allowed. + +### `refund(company, credit, credits, pin_lease_id?)` + +Add credits back to `local_remaining_credits`, **clamped at `granted_amount`**. No-op if +`credits <= 0` or no lease is in the slot. When `pin_lease_id` is given, the refund applies +**only if** the slot still holds that lease: a hold carved out of expired lease A must never +inflate successor lease B — A's remainder (including this slice) was already refunded to the +company balance server-side when A expired, so crediting B would double-count. + +### `extend(company, credit, granted_total, new_expires_at?, pin_lease_id?)` — reconcile to total + +After a remote extend, reconcile the slot to the **server-authoritative total**: + +- Compute `delta = granted_total - stored granted_amount` **atomically against the currently + stored total** — never from a caller-held pre-wire-call read (two pods extending concurrently + from the same stale read would each apply a delta and mint phantom credits). If `delta > 0`, + set `granted_amount = granted_total` and add `delta` to `local_remaining_credits`. If + `delta <= 0` (a total a sibling already applied, or a stale lower total) it is a **no-op** — + applies converge in any order. +- Expiry only ever moves **forward**: `new_expires_at` is applied only if later than the stored + expiry, so an out-of-order apply cannot shorten a lease a sibling just extended. +- When `pin_lease_id` is given and the slot holds a different lease, drop the whole extend + (credits and expiry): the server granted the extension to the pinned lease; crediting a + successor would mint credits the server refunds with the pinned lease at its expiry. +- No-op if the slot is empty. + +### `drop(company, credit)` + +Remove the slot entry (after a remote release). Plain delete. + +### Reservation table: `add`, `get`, `consume`, `reserved_credits`, `sweep_expired` + +- `add(reservation)` — register. Idempotent on id. `add` does NOT debit the lease; the debit + already happened in `try_reserve` (see [ordering](#bounded-leak-contract)). +- `consume(id, credits_consumed)` — **exactly-once claim**: atomically remove the reservation + from the table; if it was already gone (swept, or consumed by a racing caller) return `null` + and touch nothing. On a successful claim, clamp `credits_consumed` to + `[0, credits_reserved]`, refund `credits_reserved - clamped` to the lease (pinned to the + reservation's `lease_id`), and return the clamped figure. The claim and the refund are two + steps; the claim is the arbiter (see bounded-leak contract). +- `reserved_credits(company, credit)` — sum of `credits_reserved` across open reservations for + the slot. A reservation counts iff it is still in the table, so + `local_remaining_credits + reserved_credits` stays exact between operations. +- `sweep_expired(now)` — remove every reservation with `expires_at <= now` and refund its full + hold to its lease (pinned to its `lease_id`; a stale-lease hold is dropped, not refunded). + Returns the number swept. Runs on a background interval (`sweep_interval_ms`) in production; + vectors call it explicitly. + +## Lease manager + +Owns the lease lifecycle against the server wire API (`acquire`, `extend`, `release`). + +### Acquire (`acquire_if_needed`) + +- If the slot holds a **live** lease, return it — no wire call. +- Otherwise call the server: `requested_amount = lease_size`, `expires_at = now + + lease_duration_ms`. An expired local entry is left in place for `replace` to overwrite + atomically (deleting it first would open a race window against sibling pods; every reader + re-guards on expiry anyway). +- On response, `replace` the slot. If `replace` kept an existing lease (a sibling won, or the + slot's expired row was reconciled in place): + - If the installed lease has a **different id** than the one the server handed us, ours is a + redundant hold nobody will draw on — release it (fire-and-forget; a failed release falls + back to server-side lease expiry). + - If the ids are the **same** (the server is idempotent for an active slot and handed the + racing acquire the sibling's lease back), release **nothing** — releasing would pull the + shared lease out from under every sibling. + - If the slot reads empty (expired in the gap), also release nothing. + - Either way, return whatever the slot now holds. +- Wire or store failure: return "no lease" (never throw) — the caller routes it through + fail-open/fail-closed. +- Per-process single-flight per slot: concurrent callers share one in-flight wire call + (best-effort; duplicates are absorbed by the idempotent server + `replace`). + +### Extend (`maybe_extend`) + +Triggered when EITHER: + +- `local_remaining_credits / max(granted_amount, 1) <= low_water_mark` (steady-state refresh), or +- the caller passes `required_credits` and `local_remaining_credits < required_credits` (a check + just failed a reserve of that size — extend opportunistically). + +Rules: + +- **Never extend an expired lease** — the server treats it as released; the right move is a fresh + acquire on the next check. +- Wire body: `additional_amount = max(lease_size, required_credits - local_remaining_credits)`. + Sizing to the shortfall matters: a single check needing more than `remaining + lease_size` + would otherwise fail its post-extend retry forever regardless of server balance. + `expires_at = now + lease_duration_ms`. +- On response, reconcile via the store's `extend` with the server's **total** and new expiry, + **pinned** to the extended lease's id. +- Failures resolve to "no lease" without throwing (often fire-and-forget). +- Per-process single-flight per slot, kept separate from acquire's (an in-flight extend must not + satisfy an acquire, or vice versa). + +### Release on close (`release_all_local_leases`) + +Only for a **per-process (in-memory) store**, whose leases are exclusively this process's: +release every live lease over the wire (returning the unspent remainder to the company balance +immediately) and drop it locally; **skip expired** leases (already swept server-side). A shared +(Redis) store must never do this — sibling pods still draw on those leases. Best-effort: +failures fall back to server-side expiry. + +## Check flow + +`check(eval_ctx, flag_key, { usage, event_subtype?, on_acquire_failure?, ... })` — the +lease-gated feature check. Fallback = the plain (non-lease) flag check, which has its own +degradation story; when the flow "falls back", no reservation is issued and no lease state is +touched beyond what already happened. + +Guards, in order: + +1. `usage` missing → plain check (lease path not requested). +2. `usage` not a finite non-negative number → resolve **statically** by `on_acquire_failure` + (deny for fail-closed; blanket allow for fail-open, reason `invalid_usage`). The value must + never reach the stores. +3. `usage == 0` → nothing to reserve; fall back to the plain check (no 0-credit reservation). +4. No datastream / cached flag / resolvable company (or named user) → fall back. + +Then: + +5. **Entitlement probe.** Run the rules engine once against the company's *real* balance — no + substitution, no credit-cost preflight (a preflight against the lease-depleted server balance + could fail the credit condition and hide the entitlement being probed for). Read the matched + entitlement's shape: + - Not credit-metered (`value_type != "credit"`: boolean/override grant, numeric allocation, + unlimited, or not entitled) → **fall back**, no lease traffic at all. + - Credit entitlement missing `credit_id`, a positive `consumption_rate`, or a resolvable + `event_subtype` (caller's explicit subtype wins over the entitlement's) → fall back. + - Probe error → fall back (it is a resolution step, not the gate). +6. `credit_cost = usage x consumption_rate`. +7. **Acquire** a lease for `(company, credit_id)`. Failure → [failure handling](#failure-handling) + with reason `lease_acquire_failed`. +8. **Reserve** `credit_cost` via `try_reserve`. On refusal, opportunistically + `maybe_extend(required_credits = credit_cost)` (awaited) and retry the reserve **once**. + Still refused → failure handling, reason `insufficient_lease_balance`. Store error → + failure handling, reason `lease_store_error`. +9. **Record the reservation** (TTL = `reservation_ttl_ms` from now) — *after* the debit, *before* + the engine gate. If persisting fails, undo the debit (claim-and-refund; direct refund if + nothing persisted; both pinned to the lease) and go to failure handling + (`lease_store_error`). If even the undo fails, accept the bounded leak. +10. **Engine gate.** Re-run the engine against a company snapshot whose + `credit_balances[credit_id]` is substituted with the **pre-reservation** local balance + (post-debit balance returned by `try_reserve` + `credit_cost` — exact as of the debit, no + read race), passing `credit_cost = { credit_id: credit_cost }` so the engine evaluates + `pre_reservation - credit_cost >= 0` — the same arithmetic `try_reserve` just enforced, plus + every non-credit rule (plan targeting, overrides). + - Engine **allows** → keep the hold; return `{ allowed: true, reservation }`. Fire-and-forget + a watermark-driven `maybe_extend`. + - Engine **denies** → cancel the reservation (claim + full refund) and return + `{ allowed: false }` with the engine's reason. + - Engine **errors** → cancel the reservation and resolve **statically** by mode (the engine + itself is down, so no fail-open re-evaluation is possible). + +### Failure handling + +Every can't-gate outcome (acquire failed, store unreachable, lease exhausted) funnels through the +configured `on_acquire_failure` mode (default **fail-closed**): + +- **fail-closed** → `{ allowed: false }`, reason = the failure reason. No reservation. +- **fail-open** → *err on the side of assuming the credits are there*, **not** blanket allow: + re-run the engine with the credit balance substituted to an effectively unlimited value + (`MAX_SAFE_INTEGER` in Node) and the caller's usage preflight threaded through. Plan + targeting, overrides, and every non-credit condition still apply — a company that is not + entitled stays **denied** even with the lease backend down. No reservation is issued either + way; `err` carries the failure reason. Only if that evaluation itself errors does the SDK + fall back to a blanket allow. + +## Track / settle flow + +`track_with_reservation(reservation, actual_quantity)` settles a reservation: + +1. `credits = actual_quantity x reservation.consumption_rate`. +2. `consume(reservation.id, credits)`: + - **Settled locally** (claim succeeded): the clamped consumed slice stays debited; the unspent + slice is refunded to the lease (pinned). + - **Not settled** (`null`: already swept after TTL, already consumed, or store unreachable): + local lease state is untouched — if the sweeper already refunded the full hold, nothing + re-debits the consumed slice, so the local balance reads **high** until the lease rolls + over. This is why `reservation_ttl_ms` should exceed the longest expected gap between + `check()` and `track_with_reservation()`. +3. Either way, emit the Track event built from the **caller-held handle** (not the store): + `event = event_subtype`, `quantity = actual_quantity` (the *unclamped* actual — the server is + the source of truth for real consumption; only local bookkeeping clamps to the reserved + amount), `lease_id = reservation.lease_id` (routes the server-side consumption through the + lease's sub-ledger instead of double-debiting the pre-debited grant), plus the reservation's + `eval_ctx` company/user and any caller traits. +4. The Track carries a deterministic idempotency key derived from the reservation id + (`"lease-reservation:" + reservation.id` in Node); the server dedupes by it for 24h, so a + recovery emit racing the normal emit, or an accidental double settle, collapses to one billed + event across pods and restarts. +5. Guard: a non-finite or negative `actual_quantity` skips the settle entirely (no store call, no + event) — the untouched reservation expires at its TTL and the sweeper refunds the full hold. + +## Configuration knobs + +| knob (vector key) | Node name | default | meaning | +| --- | --- | --- | --- | +| `lease_duration_ms` | `defaultLeaseDuration` | 300 000 (5 min) | Lease lifetime requested at acquire/extend (`expires_at = now + duration`). | +| `reservation_ttl_ms` | `defaultReservationTTL` | 60 000 (60 s) | Reservation lifetime; the sweep deadline. Size above the longest expected check→track gap. | +| `lease_size` | `defaultLeaseSize` | 10 000 | Credits requested per acquire, and the minimum extend tranche. | +| `low_water_mark` | `lowWaterMark` | 0.25 | Remaining/granted ratio at or below which a background extend is kicked off. | +| `sweep_interval_ms` | `sweepIntervalMs` | 1 000 | Expired-reservation sweep cadence. | +| — | `onAcquireFailure` | `fail-closed` | Per-check failure mode (see check flow). | + +Per-credit-type overrides of the first four are supported (keyed by credit type id); resolution is +override → client config → default. + +## Bounded-leak contract + +The flow deliberately orders its two-step transitions so that a process crash between steps leaks +*locally held credits* (which the server reclaims at lease expiry) rather than enabling a +double-spend. The invariant direction is always: **the debit/claim is durable first; the +record/refund may be lost.** + +| # | crash window | what leaks | bound | reclaimed by | must NOT happen | +| --- | --- | --- | --- | --- | --- | +| 1 | after `try_reserve` (debit), before `add` (record) | the debited hold — invisible to the reservation table, so the sweeper can never refund it | `credits_reserved` of that one check | lease expiry: the expired balance is never served again, and the server refunds the whole grant; the successor lease installs at full grant | a reservation record without a debit (a later consume would refund credits never held → double-spend). Vectors pin that the debit lands strictly before the record. | +| 2 | inside `consume`: after the claim, before the refund | the unspent slice of that reservation | `credits_reserved` of that one reservation | lease expiry (same mechanism) | a double refund: the claim is exactly-once, so a retried settle or a sweeper finds nothing to claim and refunds nothing | +| 3 | (Redis only) after the claim, before index cleanup | nothing (bookkeeping only): the per-slot reserved-credits index transiently over-counts | one index field | the sweeper reconciles the orphaned index entry — **without refunding** (without the claimed record, exactly-once cannot be arbitrated across racing sweepers) | a refund driven by an index entry alone | + +Additional pinned properties: + +- A leak never survives its lease: after lease expiry the stale balance is refused + (`try_reserve → null`) and a successor lease restores the full grant. +- A retried check after a window-1 crash settles independently: its own slice refunds exactly + once; the leaked slice never refunds. +- A late retried settle after a window-2 crash (even after a successor lease is installed) + refunds nothing into the successor. + +## Invariants not expressible as vectors + +These hold in the reference implementation but need concurrency, wall clocks, or non-JSON values +to demonstrate; ports must uphold them and should test them natively. + +1. **Per-slot atomicity.** `replace`, `try_reserve`, `refund`, `extend` are atomic per lease + slot; `consume`'s claim is atomic per reservation. In-memory: per-key mutual exclusion. Redis: + single-key Lua scripts (single-key keeps them Redis-Cluster-safe; the refund to the lease hash + is deliberately a separate single-key step, never a multi-key script — see leak window 2/3). +2. **Server-clock expiry (shared backend).** With a shared store, lease expiry must be decided + against the *store's* clock (Redis `TIME`), not the calling process's — pods with skewed + clocks must agree on liveness. Backend rows carry a TTL grace window past `expires_at` + (60 s lease / 30 s reservation in Node) so the sweeper can still read them; expired-but-not- + evicted rows must still refuse reserves. +3. **NaN/precision guards.** Non-finite or negative amounts are rejected at every boundary + (`usage`, `try_reserve`, `actual_quantity`) — JSON cannot encode NaN, so vectors only cover + the negative case. Fractional credit amounts are legal throughout (rates like 0.1); Redis + stores balances as strings to avoid integer truncation. +4. **Single-flight.** Per-process, per-slot single-flight for acquire and for extend, tracked + separately. Best-effort only: duplicate wire calls are safe (idempotent server + keep-first + `replace` + reconcile-to-total `extend`). +5. **Concurrent cross-pod extends converge.** Two pods extending from the same stale read must + not double-count — guaranteed by reconcile-to-total computed inside the store (the sequential + out-of-order-totals vector pins the arithmetic; the concurrent schedule needs a race). +6. **Idempotent billing.** The Track idempotency key is deterministic from the reservation id; + the server dedupes for 24h. Double settles and recovery emits collapse to one billed event. +7. **Background sweep loop.** `start_sweep`/`stop` run `sweep_expired` on an interval; timers + must not keep the process alive. Vectors call `sweep_expired` explicitly instead. +8. **Fire-and-forget never rejects.** `acquire_if_needed`, `maybe_extend`, and release paths + resolve (to "no lease") on failure rather than rejecting — they are often unawaited. +9. **Offline/unconfigured degradation.** Lease config absent → `check` is a plain flag check; + `track_with_reservation` on an unconfigured client still emits the billing event with the + `lease_id` and idempotency key intact. diff --git a/conformance/vectors/check-flow.json b/conformance/vectors/check-flow.json new file mode 100644 index 00000000..eff99903 --- /dev/null +++ b/conformance/vectors/check-flow.json @@ -0,0 +1,527 @@ +{ + "category": "check_flow", + "vectors": [ + { + "name": "check_happy_path_gates_on_pre_reservation_balance", + "description": "A lease-bearing check: probe against the real balance (no substitution, no credit cost), reserve usage x rate from the lease, then gate the engine on the PRE-reservation local balance with credit_cost — the same arithmetic the atomic reserve just enforced. The hold sticks only because the engine allowed.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 10, + "event_subtype": "inference_tokens", + "save_reservation_as": "r1", + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 10, + "event_subtype": "inference_tokens" + } + }, + { "value": true, "reason": "ok" } + ], + "expect": { + "allowed": true, + "has_reservation": true, + "reservation": { + "lease_id": "lse_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10 + }, + "engine_calls": [{ "credit_balance": 5000 }, { "credit_balance": 1000, "credit_cost": 100 }] + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 900 } + }, + { "op": "reserved_credits", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "total": 100 } } + ] + }, + { + "name": "check_denied_by_engine_cancels_the_hold", + "description": "When the gate evaluation denies, the reservation made before the eval is cancelled: claimed and fully refunded, leaving no hold and no reservation.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 10, + "event_subtype": "inference_tokens", + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 10, + "event_subtype": "inference_tokens" + } + }, + { "value": false, "reason": "denied_by_targeting" } + ], + "expect": { "allowed": false, "reason": "denied_by_targeting", "has_reservation": false } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 1000 } + }, + { "op": "reserved_credits", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "total": 0 } }, + { "op": "reservation_count", "expect": { "count": 0 } } + ] + }, + { + "name": "check_acquire_failure_fail_closed_denies", + "description": "fail-closed (the default): when no lease can be acquired the check denies outright with the failure reason; no reservation is issued.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + } + }, + "operations": [ + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 10, + "event_subtype": "inference_tokens", + "on_acquire_failure": "fail-closed", + "server": { "acquire": { "error": "wire down" } }, + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 10, + "event_subtype": "inference_tokens" + } + } + ], + "expect": { + "allowed": false, + "reason": "lease_acquire_failed", + "err": "lease_acquire_failed", + "has_reservation": false + } + }, + { "op": "reservation_count", "expect": { "count": 0 } } + ] + }, + { + "name": "check_acquire_failure_fail_open_reevaluates", + "description": "fail-open is NOT blanket allow: the engine re-runs with the credit balance substituted to an effectively unlimited value and the caller's usage preflight threaded through, so non-credit rules still apply. Here they pass, so the check allows — with the failure recorded in err and no reservation.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + } + }, + "operations": [ + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 10, + "event_subtype": "inference_tokens", + "on_acquire_failure": "fail-open", + "server": { "acquire": { "error": "wire down" } }, + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 10, + "event_subtype": "inference_tokens" + } + }, + { "value": true, "reason": "evaluated" } + ], + "expect": { + "allowed": true, + "reason": "evaluated (lease_acquire_failed_fail_open)", + "err": "lease_acquire_failed", + "has_reservation": false, + "engine_calls": [ + { "credit_balance": 5000 }, + { + "credit_balance": "max_safe_integer", + "event_usage": { "event_subtype": "inference_tokens", "quantity": 10 } + } + ] + } + }, + { "op": "reservation_count", "expect": { "count": 0 } } + ] + }, + { + "name": "check_fail_open_still_denies_when_rules_deny", + "description": "fail-open with a denying rules evaluation stays denied: substituting an unlimited balance only bypasses the credit gate, never plan targeting or overrides.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + } + }, + "operations": [ + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 10, + "event_subtype": "inference_tokens", + "on_acquire_failure": "fail-open", + "server": { "acquire": { "error": "wire down" } }, + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 10, + "event_subtype": "inference_tokens" + } + }, + { "value": false, "reason": "not_targeted" } + ], + "expect": { + "allowed": false, + "reason": "not_targeted (lease_acquire_failed_fail_open)", + "err": "lease_acquire_failed", + "has_reservation": false + } + } + ] + }, + { + "name": "check_insufficient_lease_extends_and_retries", + "description": "A reserve refusal triggers an awaited opportunistic extend sized to cover the request (required_credits = credit_cost), then exactly one reserve retry. On success the check proceeds normally, gating on the post-extend pre-reservation balance.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 950, + "expect": { "balance": 50 } + }, + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 10, + "event_subtype": "inference_tokens", + "server": { "extend": { "lease": { "granted_total": 2000, "expires_at_ms": 600000 } } }, + "save_reservation_as": "r1", + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 10, + "event_subtype": "inference_tokens" + } + }, + { "value": true, "reason": "ok" } + ], + "expect": { + "allowed": true, + "has_reservation": true, + "wire_extends": 1, + "last_extend_additional_amount": 1000, + "engine_calls": [{ "credit_balance": 5000 }, { "credit_balance": 1050, "credit_cost": 100 }] + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 2000, "local_remaining_credits": 950 } + } + ] + }, + { + "name": "check_insufficient_lease_after_failed_extend_resolves_by_mode", + "description": "When the opportunistic extend fails and the retry is still refused, the check resolves through the failure mode (fail-closed here) with reason insufficient_lease_balance, leaving the lease balance untouched.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 950, + "expect": { "balance": 50 } + }, + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 10, + "event_subtype": "inference_tokens", + "on_acquire_failure": "fail-closed", + "server": { "extend": { "error": "wire down" } }, + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 10, + "event_subtype": "inference_tokens" + } + } + ], + "expect": { + "allowed": false, + "reason": "insufficient_lease_balance", + "err": "insufficient_lease_balance", + "has_reservation": false + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 50 } + }, + { "op": "reservation_count", "expect": { "count": 0 } } + ] + }, + { + "name": "check_falls_back_without_usable_credit_entitlement", + "description": "A non-credit matched entitlement (boolean/override/numeric/unlimited/not entitled) or an incomplete credit entitlement (no positive consumption rate) defers to the plain check: no lease traffic, no reservation.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + } + }, + "operations": [ + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": {} }, + "usage": 10, + "event_subtype": "inference_tokens", + "engine": [{ "value": true, "reason": "probe", "entitlement": { "value_type": "boolean" } }], + "expect": { "fallback_called": true, "reason": "fallback", "has_reservation": false } + }, + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 10, + "event_subtype": "inference_tokens", + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 0, + "event_subtype": "inference_tokens" + } + } + ], + "expect": { "fallback_called": true, "reason": "fallback", "has_reservation": false } + }, + { "op": "reservation_count", "expect": { "count": 0 } } + ] + }, + { + "name": "check_zero_usage_falls_back", + "description": "usage = 0 means nothing to reserve: the check defers to the plain (preflight-threaded) check instead of issuing a no-op 0-credit reservation.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 0, + "event_subtype": "inference_tokens", + "engine": [], + "expect": { "fallback_called": true, "reason": "fallback", "has_reservation": false } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 1000 } + } + ] + }, + { + "name": "check_invalid_usage_resolves_statically_by_mode", + "description": "A negative (or non-finite) usage must never reach the stores; the check resolves statically by mode without any engine evaluation: deny for fail-closed, blanket allow for fail-open, reason invalid_usage either way.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": -5, + "event_subtype": "inference_tokens", + "on_acquire_failure": "fail-closed", + "engine": [], + "expect": { + "allowed": false, + "reason": "invalid_usage", + "err": "invalid_usage", + "has_reservation": false, + "engine_calls": [] + } + }, + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": -5, + "event_subtype": "inference_tokens", + "on_acquire_failure": "fail-open", + "engine": [], + "expect": { + "allowed": true, + "reason": "invalid_usage_fail_open", + "err": "invalid_usage", + "has_reservation": false, + "engine_calls": [] + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 1000 } + } + ] + } + ] +} diff --git a/conformance/vectors/crash-windows.json b/conformance/vectors/crash-windows.json new file mode 100644 index 00000000..c8ee05b9 --- /dev/null +++ b/conformance/vectors/crash-windows.json @@ -0,0 +1,271 @@ +{ + "category": "crash_window", + "vectors": [ + { + "name": "debit_without_record_leaks_bounded", + "description": "Crash window 1 (debit-then-add): the atomic debit landed but the reservation record never did. The leak is exactly the reserved amount; the sweeper can never refund a hold that was never recorded.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { "op": "reserved_credits", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "total": 0 } }, + { "op": "advance_clock", "ms": 55000 }, + { "op": "sweep_expired", "expect": { "swept": 0 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 900 } + } + ] + }, + { + "name": "debit_leak_reclaimed_at_lease_expiry", + "description": "Crash window 1 recovery: the leaked balance is never served after lease expiry, and the successor lease installs at the full grant — the leak does not outlive the lease.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { "op": "advance_clock", "ms": 60001 }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 1, + "expect": { "balance": null } + }, + { + "op": "replace_lease", + "lease_id": "lse_2", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 180000, + "expect": { "written": true } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 1000 } + } + ] + }, + { + "name": "retried_check_after_debit_leak_settles_once", + "description": "A retry after a window-1 crash is a fresh check with a fresh reservation: its unspent slice refunds exactly once; the leaked slice never refunds — not on a repeat consume, not on a sweep.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 3600000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 800 } + }, + { + "op": "add_reservation", + "id": "res_retry", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { "op": "consume_reservation", "id": "res_retry", "credits": 40, "expect": { "consumed": 40 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 860 } + }, + { "op": "consume_reservation", "id": "res_retry", "credits": 40, "expect": { "consumed": null } }, + { "op": "advance_clock", "ms": 70000 }, + { "op": "sweep_expired", "expect": { "swept": 0 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 860 } + } + ] + }, + { + "name": "crash_before_refund_claim_is_durable", + "description": "Crash window 2 (consume-then-refund): the claim survives the crash, so the reservation is gone everywhere and nothing can double-spend; the unspent slice's refund is lost, bounded by credits_reserved. A retried settle neither re-claims nor double-refunds, and the sweeper cannot refund a claimed reservation.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 3600000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { + "op": "consume_reservation", + "id": "res_1", + "credits": 30, + "crash_before_refund": true, + "expect": { "throws": true } + }, + { "op": "get_reservation", "id": "res_1", "expect": { "exists": false } }, + { "op": "reserved_credits", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "total": 0 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 900 } + }, + { "op": "consume_reservation", "id": "res_1", "credits": 30, "expect": { "consumed": null } }, + { "op": "advance_clock", "ms": 70000 }, + { "op": "sweep_expired", "expect": { "swept": 0 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 900 } + } + ] + }, + { + "name": "crash_before_refund_reclaimed_at_lease_expiry", + "description": "Crash window 2 recovery: after the lease expires and a successor takes the slot at full grant, a very late retried settle of the crashed reservation must not leak the lost refund into the successor.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { + "op": "consume_reservation", + "id": "res_1", + "credits": 30, + "crash_before_refund": true, + "expect": { "throws": true } + }, + { "op": "advance_clock", "ms": 60001 }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 1, + "expect": { "balance": null } + }, + { + "op": "replace_lease", + "lease_id": "lse_2", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 180000, + "expect": { "written": true } + }, + { "op": "consume_reservation", "id": "res_1", "credits": 0, "expect": { "consumed": null } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "lease_id": "lse_2", "local_remaining_credits": 1000 } + } + ] + } + ] +} diff --git a/conformance/vectors/expiry.json b/conformance/vectors/expiry.json new file mode 100644 index 00000000..e1f7e2cc --- /dev/null +++ b/conformance/vectors/expiry.json @@ -0,0 +1,198 @@ +{ + "category": "expiry", + "vectors": [ + { + "name": "expired_lease_never_serves_reserves", + "description": "Past its expiry a lease's balance is stale (the server refunded the grant): reserves are refused even with ample local balance.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { "op": "advance_clock", "ms": 60001 }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 1, + "expect": { "balance": null } + } + ] + }, + { + "name": "successor_after_expiry_restores_full_grant", + "description": "A successor lease installed over an expired slot starts at its full grant — nothing from the expired lease (debits or leaks) carries over.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 700, + "expect": { "balance": 300 } + }, + { "op": "advance_clock", "ms": 60001 }, + { + "op": "replace_lease", + "lease_id": "lse_2", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 180000, + "expect": { "written": true } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "lease_id": "lse_2", "local_remaining_credits": 1000 } + } + ] + }, + { + "name": "sweep_refunds_expired_holds_only", + "description": "The sweeper refunds an expired reservation's full hold to its lease and leaves live reservations untouched.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 300, + "expect": { "balance": 700 } + }, + { + "op": "add_reservation", + "id": "res_short", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 10000 + }, + { + "op": "add_reservation", + "id": "res_long", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 20, + "credits_reserved": 200, + "consumption_rate": 10, + "expires_at_ms": 200000 + }, + { "op": "sweep_expired", "expect": { "swept": 0 } }, + { "op": "advance_clock", "ms": 10001 }, + { "op": "sweep_expired", "expect": { "swept": 1 } }, + { "op": "get_reservation", "id": "res_short", "expect": { "exists": false } }, + { "op": "get_reservation", "id": "res_long", "expect": { "exists": true } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 800 } + }, + { "op": "reserved_credits", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "total": 200 } } + ] + }, + { + "name": "sweep_of_stale_lease_hold_does_not_inflate_successor", + "description": "Sweeping (or consuming) a reservation carved from an expired lease refunds nothing into the successor lease occupying the slot: the hold is pinned to its originating lease_id.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "add_reservation", + "id": "res_stale", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 90000 + }, + { "op": "advance_clock", "ms": 60001 }, + { + "op": "replace_lease", + "lease_id": "lse_2", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000, + "expect": { "written": true } + }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 200, + "expect": { "balance": 800 } + }, + { "op": "advance_clock", "ms": 30000 }, + { "op": "sweep_expired", "expect": { "swept": 1 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "lease_id": "lse_2", "local_remaining_credits": 800 } + } + ] + } + ] +} diff --git a/conformance/vectors/lease-lifecycle.json b/conformance/vectors/lease-lifecycle.json new file mode 100644 index 00000000..085b5faa --- /dev/null +++ b/conformance/vectors/lease-lifecycle.json @@ -0,0 +1,447 @@ +{ + "category": "lease_lifecycle", + "vectors": [ + { + "name": "replace_installs_full_grant", + "description": "A fresh install initializes local_remaining_credits to the full granted amount.", + "operations": [ + { + "op": "replace_lease", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000, + "expect": { "written": true } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { + "exists": true, + "lease_id": "lse_1", + "granted_amount": 1000, + "local_remaining_credits": 1000 + } + } + ] + }, + { + "name": "replace_keeps_live_lease_even_with_different_id", + "description": "A live lease occupying the slot wins over any replace — even one carrying a different lease_id (a sibling raced this acquire). Its already-debited balance is preserved.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 400, + "expect": { "balance": 600 } + }, + { + "op": "replace_lease", + "lease_id": "lse_2", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 5000, + "expires_at_ms": 120000, + "expect": { "written": false } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { + "exists": true, + "lease_id": "lse_1", + "granted_amount": 1000, + "local_remaining_credits": 600 + } + } + ] + }, + { + "name": "replace_reconciles_expired_slot_with_same_id", + "description": "A stale acquire response for the SAME lease landing over its own expired local row must not reinstall it: that would reset local_remaining_credits and erase debits whose reservations are still open. The row is reconciled like an extend (granted to total, expiry forward, balance untouched) and reported as kept.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 400, + "expect": { "balance": 600 } + }, + { "op": "advance_clock", "ms": 60001 }, + { + "op": "replace_lease", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1500, + "expires_at_ms": 120000, + "expect": { "written": false } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { + "exists": true, + "lease_id": "lse_1", + "granted_amount": 1500, + "local_remaining_credits": 1100 + } + }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 1000 } + } + ] + }, + { + "name": "replace_overwrites_expired_lease", + "description": "An expired lease does not block the slot: replace overwrites it atomically and restores the full new grant.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 400, + "expect": { "balance": 600 } + }, + { "op": "advance_clock", "ms": 60001 }, + { + "op": "replace_lease", + "lease_id": "lse_2", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 180000, + "expect": { "written": true } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "lease_id": "lse_2", "local_remaining_credits": 1000 } + } + ] + }, + { + "name": "try_reserve_insufficient_and_boundary", + "description": "A reserve larger than the remaining balance is refused and touches nothing; reserving down to exactly zero is allowed; negative amounts are always refused.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 100, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 101, + "expect": { "balance": null } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 100 } + }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": -1, + "expect": { "balance": null } + }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 0 } + }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 1, + "expect": { "balance": null } + }, + { + "op": "try_reserve", + "company_id": "co_9", + "credit_type_id": "ct_1", + "credits": 1, + "expect": { "balance": null } + } + ] + }, + { + "name": "refund_clamped_at_granted_amount", + "description": "A refund can never push the local balance above the granted amount.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { "op": "refund_lease", "company_id": "co_1", "credit_type_id": "ct_1", "credits": 500 }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 1000 } + } + ] + }, + { + "name": "refund_pinned_to_lease_id_dropped_on_successor", + "description": "A refund pinned to an expired lease's id must not inflate the successor lease occupying the slot — the expired lease's remainder was already returned to the company balance server-side. An unpinned refund still applies.", + "given": { + "leases": [ + { + "lease_id": "lse_2", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 200, + "expect": { "balance": 800 } + }, + { + "op": "refund_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "pin_lease_id": "lse_1" + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 800 } + }, + { + "op": "refund_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "pin_lease_id": "lse_2" + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 900 } + } + ] + }, + { + "name": "extend_reconciles_to_total_and_converges", + "description": "Extend applies the server-authoritative TOTAL: the delta is computed against the stored total, so a repeated or stale-lower total is a no-op and out-of-order applies converge without minting credits.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 800, + "expect": { "balance": 200 } + }, + { + "op": "extend_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_total": 3000, + "expires_at_ms": 300000 + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 3000, "local_remaining_credits": 2200 } + }, + { + "op": "extend_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_total": 3000, + "expires_at_ms": 300000 + }, + { + "op": "extend_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_total": 2000, + "expires_at_ms": 300000 + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 3000, "local_remaining_credits": 2200 } + } + ] + }, + { + "name": "extend_expiry_only_moves_forward", + "description": "An extend carrying an earlier expiry must not shorten the lease: after an extend to a later expiry, a stale out-of-order apply with an earlier expiry leaves the lease live past that earlier instant.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "extend_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_total": 2000, + "expires_at_ms": 120000 + }, + { + "op": "extend_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_total": 2000, + "expires_at_ms": 30000 + }, + { "op": "advance_clock", "ms": 90000 }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 1900 } + } + ] + }, + { + "name": "extend_pinned_to_lease_id_dropped_on_successor", + "description": "An extend pinned to a lease the slot no longer holds is dropped entirely — crediting the successor would mint credits the server granted to the expired lease.", + "given": { + "leases": [ + { + "lease_id": "lse_2", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "extend_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_total": 5000, + "expires_at_ms": 600000, + "pin_lease_id": "lse_1" + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 1000, "local_remaining_credits": 1000 } + }, + { + "op": "extend_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_total": 2000, + "expires_at_ms": 600000, + "pin_lease_id": "lse_2" + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 2000, "local_remaining_credits": 2000 } + } + ] + } + ] +} diff --git a/conformance/vectors/lease-manager.json b/conformance/vectors/lease-manager.json new file mode 100644 index 00000000..79ab1afa --- /dev/null +++ b/conformance/vectors/lease-manager.json @@ -0,0 +1,417 @@ +{ + "category": "lease_manager", + "vectors": [ + { + "name": "acquire_installs_tranche_and_reuses_live_lease", + "description": "First acquire requests lease_size from the server and installs the response at full grant; a second acquire while the lease is live makes no wire call.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + } + }, + "operations": [ + { + "op": "acquire_if_needed", + "company_id": "co_1", + "credit_type_id": "ct_1", + "server": { + "lease": { "lease_id": "lse_1", "granted_amount": 1000, "expires_at_ms": 300000 } + }, + "expect": { + "lease_id": "lse_1", + "wire_acquires": 1, + "last_acquire_requested_amount": 1000, + "released_lease_ids": [] + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "lease_id": "lse_1", "local_remaining_credits": 1000 } + }, + { + "op": "acquire_if_needed", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "lease_id": "lse_1", "wire_acquires": 1 } + } + ] + }, + { + "name": "acquire_replaces_expired_slot_without_release", + "description": "An expired slot triggers a fresh acquire that supplants the stale entry in place; the redundant-lease release path must not fire (replace wrote, it did not keep a live lease).", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_stale", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { "op": "advance_clock", "ms": 60001 }, + { + "op": "acquire_if_needed", + "company_id": "co_1", + "credit_type_id": "ct_1", + "server": { + "lease": { "lease_id": "lse_fresh", "granted_amount": 1000, "expires_at_ms": 360000 } + }, + "expect": { "lease_id": "lse_fresh", "wire_acquires": 1, "released_lease_ids": [] } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "lease_id": "lse_fresh", "local_remaining_credits": 1000 } + } + ] + }, + { + "name": "lost_acquire_race_different_id_releases_redundant_lease", + "description": "A sibling installs a live lease while this acquire's wire call is in flight. The installed lease (with its debited balance) wins; the lease the server minted for the loser is redundant and gets released so it is not orphaned against the company balance.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + } + }, + "operations": [ + { + "op": "acquire_if_needed", + "company_id": "co_1", + "credit_type_id": "ct_1", + "server": { + "lease": { "lease_id": "lse_loser", "granted_amount": 1000, "expires_at_ms": 300000 } + }, + "install_during_wire": { + "lease_id": "lse_winner", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + }, + "expect": { "lease_id": "lse_winner", "wire_acquires": 1, "released_lease_ids": ["lse_loser"] } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "lease_id": "lse_winner" } + } + ] + }, + { + "name": "lost_acquire_race_same_id_releases_nothing", + "description": "The server is idempotent for an active slot: a racing acquire is handed back the SAME lease the sibling installed. There is nothing to release — releasing would pull the shared lease out from under every sibling.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + } + }, + "operations": [ + { + "op": "acquire_if_needed", + "company_id": "co_1", + "credit_type_id": "ct_1", + "server": { + "lease": { "lease_id": "lse_shared", "granted_amount": 1000, "expires_at_ms": 300000 } + }, + "install_during_wire": { + "lease_id": "lse_shared", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + }, + "expect": { "lease_id": "lse_shared", "wire_acquires": 1, "released_lease_ids": [] } + } + ] + }, + { + "name": "extend_triggered_at_low_water_mark_requests_tranche", + "description": "At or below the low-water-mark ratio a steady-state extend fires, requesting the configured tranche (lease_size) and reconciling the local row to the server total.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 800, + "expect": { "balance": 200 } + }, + { + "op": "maybe_extend", + "company_id": "co_1", + "credit_type_id": "ct_1", + "server": { "lease": { "granted_total": 2000, "expires_at_ms": 600000 } }, + "expect": { + "wire_extends": 1, + "last_extend_additional_amount": 1000, + "last_extend_lease_id": "lse_1" + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 2000, "local_remaining_credits": 1200 } + } + ] + }, + { + "name": "extend_triggered_by_required_credits_above_watermark", + "description": "Above the watermark no steady-state extend fires; a required_credits hint larger than the local remaining triggers one anyway (a check just failed a reserve of that size).", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "maybe_extend", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "wire_extends": 0 } + }, + { + "op": "maybe_extend", + "company_id": "co_1", + "credit_type_id": "ct_1", + "required_credits": 1500, + "server": { "lease": { "granted_total": 2000, "expires_at_ms": 600000 } }, + "expect": { "wire_extends": 1, "last_extend_additional_amount": 1000 } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 2000, "local_remaining_credits": 1900 } + } + ] + }, + { + "name": "extend_sized_to_shortfall_when_larger_than_tranche", + "description": "additional_amount = max(lease_size, required_credits - local_remaining): a single request larger than remaining + tranche must extend by the shortfall, or its post-extend retry would fail forever regardless of server balance.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "maybe_extend", + "company_id": "co_1", + "credit_type_id": "ct_1", + "required_credits": 5000, + "server": { "lease": { "granted_total": 5100, "expires_at_ms": 600000 } }, + "expect": { "wire_extends": 1, "last_extend_additional_amount": 4100 } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 5100, "local_remaining_credits": 5000 } + } + ] + }, + { + "name": "never_extend_an_expired_lease", + "description": "An expired lease is released as far as the server is concerned — the only correct move is a fresh acquire, never an extend, no matter how depleted the balance.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_old", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 900, + "expect": { "balance": 100 } + }, + { "op": "advance_clock", "ms": 60001 }, + { + "op": "maybe_extend", + "company_id": "co_1", + "credit_type_id": "ct_1", + "required_credits": 1500, + "expect": { "wire_extends": 0 } + } + ] + }, + { + "name": "wire_failures_resolve_to_no_lease_without_state_changes", + "description": "A failed acquire yields no lease and installs nothing; a failed extend leaves the local row untouched. Neither throws (both are routed through fail-open/fail-closed by callers).", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + } + }, + "operations": [ + { + "op": "acquire_if_needed", + "company_id": "co_1", + "credit_type_id": "ct_1", + "server": { "error": "wire down" }, + "expect": { "lease_id": null, "wire_acquires": 1 } + }, + { "op": "get_lease", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "exists": false } }, + { + "op": "replace_lease", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000, + "expect": { "written": true } + }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 800, + "expect": { "balance": 200 } + }, + { + "op": "maybe_extend", + "company_id": "co_1", + "credit_type_id": "ct_1", + "server": { "error": "wire down" }, + "expect": { "wire_extends": 1 } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "granted_amount": 1000, "local_remaining_credits": 200 } + } + ] + }, + { + "name": "release_all_releases_live_and_skips_expired", + "description": "On close, a per-process store releases its live leases over the wire (returning remainders immediately) and drops them locally; expired leases are skipped — the server already swept them. Only valid for an exclusively-owned (in-memory) store; a shared backend must never enumerate-and-release.", + "backends": ["in_memory"], + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_live", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + }, + { + "lease_id": "lse_expired", + "company_id": "co_2", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 30000 + } + ] + }, + "operations": [ + { "op": "advance_clock", "ms": 30001 }, + { + "op": "release_all_local_leases", + "expect": { "released_lease_ids": ["lse_live"] } + }, + { "op": "get_lease", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "exists": false } } + ] + } + ] +} diff --git a/conformance/vectors/reservation-lifecycle.json b/conformance/vectors/reservation-lifecycle.json new file mode 100644 index 00000000..eea6b50e --- /dev/null +++ b/conformance/vectors/reservation-lifecycle.json @@ -0,0 +1,353 @@ +{ + "category": "reservation_lifecycle", + "vectors": [ + { + "name": "consume_exact_usage_no_refund", + "description": "Consuming exactly the reserved amount removes the reservation and refunds nothing.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { "op": "consume_reservation", "id": "res_1", "credits": 100, "expect": { "consumed": 100 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 900 } + }, + { "op": "get_reservation", "id": "res_1", "expect": { "exists": false } } + ] + }, + { + "name": "consume_under_reserved_refunds_unspent", + "description": "Consuming less than reserved refunds the unspent slice to the lease in the same step.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { "op": "consume_reservation", "id": "res_1", "credits": 30, "expect": { "consumed": 30 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 970 } + } + ] + }, + { + "name": "consume_over_reserved_clamps_to_hold", + "description": "Local consumption is clamped to credits_reserved: over-use consumes the full hold, refunds nothing, and never debits the lease beyond the reservation. (The billed Track quantity is NOT clamped — see track-settle vectors.)", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { "op": "consume_reservation", "id": "res_1", "credits": 999, "expect": { "consumed": 100 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 900 } + } + ] + }, + { + "name": "consume_zero_cancels_with_full_refund", + "description": "Consuming 0 credits acts as a cancel: the full hold is refunded. Negative consumption clamps to 0 the same way.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { "op": "consume_reservation", "id": "res_1", "credits": 0, "expect": { "consumed": 0 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 1000 } + } + ] + }, + { + "name": "consume_is_exactly_once", + "description": "A missing reservation and a second consume of the same id both return null and refund nothing — the claim is the exactly-once arbiter.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { "op": "consume_reservation", "id": "res_missing", "credits": 10, "expect": { "consumed": null } }, + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 100, + "expect": { "balance": 900 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { "op": "consume_reservation", "id": "res_1", "credits": 30, "expect": { "consumed": 30 } }, + { "op": "consume_reservation", "id": "res_1", "credits": 30, "expect": { "consumed": null } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 970 } + } + ] + }, + { + "name": "reserved_credits_sums_open_holds_per_slot", + "description": "reserved_credits sums credits_reserved across open reservations for the exact (company, credit) slot only, and a hold stops counting the moment it is consumed.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 350, + "expect": { "balance": 650 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { + "op": "add_reservation", + "id": "res_2", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 25, + "credits_reserved": 250, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { + "op": "add_reservation", + "id": "res_other_credit", + "lease_id": "lse_9", + "company_id": "co_1", + "credit_type_id": "ct_2", + "event_subtype": "inference_tokens", + "quantity_reserved": 99, + "credits_reserved": 999, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { + "op": "add_reservation", + "id": "res_other_company", + "lease_id": "lse_8", + "company_id": "co_2", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 99, + "credits_reserved": 999, + "consumption_rate": 10, + "expires_at_ms": 60000 + }, + { + "op": "reserved_credits", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "total": 350 } + }, + { + "op": "reserved_credits", + "company_id": "co_1", + "credit_type_id": "ct_2", + "expect": { "total": 999 } + }, + { "op": "reserved_credits", "company_id": "co_9", "credit_type_id": "ct_1", "expect": { "total": 0 } }, + { "op": "consume_reservation", "id": "res_1", "credits": 40, "expect": { "consumed": 40 } }, + { "op": "reserved_credits", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "total": 250 } } + ] + }, + { + "name": "fractional_credit_amounts_are_exact", + "description": "Fractional consumption rates produce fractional holds; reserve, consume, and refund arithmetic must not truncate.", + "given": { + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 10, + "expires_at_ms": 60000 + } + ] + }, + "operations": [ + { + "op": "try_reserve", + "company_id": "co_1", + "credit_type_id": "ct_1", + "credits": 2.5, + "expect": { "balance": 7.5 } + }, + { + "op": "add_reservation", + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 25, + "credits_reserved": 2.5, + "consumption_rate": 0.1, + "expires_at_ms": 60000 + }, + { "op": "consume_reservation", "id": "res_1", "credits": 1.5, "expect": { "consumed": 1.5 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 8.5 } + } + ] + } + ] +} diff --git a/conformance/vectors/track-settle.json b/conformance/vectors/track-settle.json new file mode 100644 index 00000000..1da71a09 --- /dev/null +++ b/conformance/vectors/track-settle.json @@ -0,0 +1,194 @@ +{ + "category": "track_settle", + "vectors": [ + { + "name": "track_underuse_settles_and_refunds_unspent", + "description": "Settling with less than the reserved usage consumes actual x rate, refunds the unspent slice to the lease, and emits a Track billing the ACTUAL quantity keyed to the lease.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 10, + "event_subtype": "inference_tokens", + "save_reservation_as": "r1", + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 10, + "event_subtype": "inference_tokens" + } + }, + { "value": true, "reason": "ok" } + ], + "expect": { "allowed": true, "has_reservation": true } + }, + { + "op": "track", + "handle": "r1", + "actual_quantity": 4, + "expect": { + "settled_locally": true, + "track": { "event": "inference_tokens", "quantity": 4, "lease_id": "lse_1" } + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 960 } + }, + { "op": "reserved_credits", "company_id": "co_1", "credit_type_id": "ct_1", "expect": { "total": 0 } } + ] + }, + { + "name": "track_overuse_bills_actual_but_clamps_local_debit", + "description": "Actual usage above the reservation: the LOCAL settle clamps consumption to the reserved hold (the lease is never debited past the reservation), but the Track event bills the unclamped actual quantity — the server is the source of truth for real consumption.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 10, + "event_subtype": "inference_tokens", + "save_reservation_as": "r1", + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 10, + "event_subtype": "inference_tokens" + } + }, + { "value": true, "reason": "ok" } + ], + "expect": { "allowed": true, "has_reservation": true } + }, + { + "op": "track", + "handle": "r1", + "actual_quantity": 25, + "expect": { + "settled_locally": true, + "track": { "event": "inference_tokens", "quantity": 25, "lease_id": "lse_1" } + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 900 } + } + ] + }, + { + "name": "track_after_sweep_is_a_recovery_emit", + "description": "Work outliving the reservation TTL: the sweeper already refunded the full hold, so the late settle does not touch the lease (the local balance reads high until rollover) — but the Track is still emitted so the server bills the actual usage. Server-side idempotency (a deterministic key derived from the reservation id) is what keeps a racing normal emit from double-billing.", + "given": { + "config": { + "lease_duration_ms": 300000, + "reservation_ttl_ms": 60000, + "lease_size": 1000, + "low_water_mark": 0.25 + }, + "leases": [ + { + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "granted_amount": 1000, + "expires_at_ms": 300000 + } + ] + }, + "operations": [ + { + "op": "check", + "flag_key": "inference", + "company": { "id": "co_1", "credit_balances": { "ct_1": 5000 } }, + "usage": 10, + "event_subtype": "inference_tokens", + "save_reservation_as": "r1", + "engine": [ + { + "value": true, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "ct_1", + "consumption_rate": 10, + "event_subtype": "inference_tokens" + } + }, + { "value": true, "reason": "ok" } + ], + "expect": { "allowed": true, "has_reservation": true } + }, + { "op": "advance_clock", "ms": 60001 }, + { "op": "sweep_expired", "expect": { "swept": 1 } }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 1000 } + }, + { + "op": "track", + "handle": "r1", + "actual_quantity": 4, + "expect": { + "settled_locally": false, + "track": { "event": "inference_tokens", "quantity": 4, "lease_id": "lse_1" } + } + }, + { + "op": "get_lease", + "company_id": "co_1", + "credit_type_id": "ct_1", + "expect": { "exists": true, "local_remaining_credits": 1000 } + } + ] + } + ] +} diff --git a/poetry.lock b/poetry.lock index 416786d8..59e6adf8 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.3.2 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. [[package]] name = "aiohappyeyeballs" @@ -6,12 +6,10 @@ version = "2.7.1" description = "Happy Eyeballs for asyncio" optional = false python-versions = ">=3.10" -groups = ["main", "dev"] files = [ {file = "aiohappyeyeballs-2.7.1-py3-none-any.whl", hash = "sha256:9243213661e29250eb41368e5daa826fc017156c3b8a11440826b2e3ed376472"}, {file = "aiohappyeyeballs-2.7.1.tar.gz", hash = "sha256:065665c041c42a5938ed220bdcd7230f22527fbec085e1853d2402c8a3615d9d"}, ] -markers = {main = "extra == \"aiohttp\""} [[package]] name = "aiohttp" @@ -19,7 +17,6 @@ version = "3.14.3" description = "Async http client/server framework (asyncio)" optional = false python-versions = ">=3.10" -groups = ["main", "dev"] files = [ {file = "aiohttp-3.14.3-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:eb0495d778817619273c108784292be161a924b9f5ae5cbbc70a2caa6838250b"}, {file = "aiohttp-3.14.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:c3c200cf9757edd785051dc699c7ecbec22110dbfcb3fefc7a9f9695eda8ea7a"}, @@ -141,7 +138,6 @@ files = [ {file = "aiohttp-3.14.3-cp314-cp314t-win_arm64.whl", hash = "sha256:2e1161602f45a54de2ce0905243a95f58cb42dcd378402f3697f5e0b21e9d2e7"}, {file = "aiohttp-3.14.3.tar.gz", hash = "sha256:9491196535a88924a60afd5b5f434b5b203b6cc616250878dbdb223a8f7844bc"}, ] -markers = {main = "extra == \"aiohttp\""} [package.dependencies] aiohappyeyeballs = ">=2.5.0" @@ -155,7 +151,7 @@ typing_extensions = {version = ">=4.4", markers = "python_version < \"3.13\""} yarl = ">=1.17.0,<2.0" [package.extras] -speedups = ["Brotli (>=1.2) ; platform_python_implementation == \"CPython\" and sys_platform != \"android\" and sys_platform != \"ios\"", "aiodns (>=3.3.0) ; sys_platform != \"android\" and sys_platform != \"ios\"", "backports.zstd ; platform_python_implementation == \"CPython\" and python_version < \"3.14\" and sys_platform != \"android\" and sys_platform != \"ios\"", "brotlicffi (>=1.2) ; platform_python_implementation != \"CPython\""] +speedups = ["Brotli (>=1.2)", "aiodns (>=3.3.0)", "backports.zstd", "brotlicffi (>=1.2)"] [[package]] name = "aiosignal" @@ -163,12 +159,10 @@ version = "1.4.0" description = "aiosignal: a list of registered asynchronous callbacks" optional = false python-versions = ">=3.9" -groups = ["main", "dev"] files = [ {file = "aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e"}, {file = "aiosignal-1.4.0.tar.gz", hash = "sha256:f47eecd9468083c2029cc99945502cb7708b082c232f9aca65da147157b251c7"}, ] -markers = {main = "extra == \"aiohttp\""} [package.dependencies] frozenlist = ">=1.1.0" @@ -180,7 +174,6 @@ version = "0.8.0" description = "Reusable constraint types to use with typing.Annotated" optional = false python-versions = ">=3.10" -groups = ["main"] files = [ {file = "annotated_types-0.8.0-py3-none-any.whl", hash = "sha256:f072f4d804ea359e4eaf198b1af7a8b0943881a87f31bb764f8bf219bb9419e0"}, {file = "annotated_types-0.8.0.tar.gz", hash = "sha256:13b2beaad985e05e2d6407ee4c4f35590b11f8d693a258a561055cac8f64cab7"}, @@ -192,7 +185,6 @@ version = "4.15.0" description = "High-level concurrency and networking framework on top of asyncio or Trio" optional = false python-versions = ">=3.10" -groups = ["main"] files = [ {file = "anyio-4.15.0-py3-none-any.whl", hash = "sha256:7ecd9937369ffce8bba0b5ccb9b3a9507b101b0ed50256aecfbab27e6c2acb99"}, {file = "anyio-4.15.0.tar.gz", hash = "sha256:b5c620ed540725e2579c31b17bb995b3bf02c9281c9cace04c7d186380bab85e"}, @@ -212,12 +204,10 @@ version = "5.0.1" description = "Timeout context manager for asyncio programs" optional = false python-versions = ">=3.8" -groups = ["main", "dev"] files = [ {file = "async_timeout-5.0.1-py3-none-any.whl", hash = "sha256:39e3809566ff85354557ec2398b55e096c8364bacac9405a7a1fa429e77fe76c"}, {file = "async_timeout-5.0.1.tar.gz", hash = "sha256:d9321a7a3d5a6a5e187e824d2fa0793ce379a202935782d555d6e9d2735677d3"}, ] -markers = {main = "extra == \"aiohttp\" and python_version == \"3.10\"", dev = "python_full_version < \"3.11.3\""} [[package]] name = "attrs" @@ -225,12 +215,10 @@ version = "26.1.0" description = "Classes Without Boilerplate" optional = false python-versions = ">=3.9" -groups = ["main", "dev"] files = [ {file = "attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309"}, {file = "attrs-26.1.0.tar.gz", hash = "sha256:d03ceb89cb322a8fd706d4fb91940737b6642aa36998fe130a9bc96c985eff32"}, ] -markers = {main = "extra == \"aiohttp\""} [[package]] name = "certifi" @@ -238,7 +226,6 @@ version = "2026.7.22" description = "Python package for providing Mozilla's CA Bundle." optional = false python-versions = ">=3.7" -groups = ["main"] files = [ {file = "certifi-2026.7.22-py3-none-any.whl", hash = "sha256:62f22742b58a1a33014a2b6b706588a8d7e2a88ae7bd1a6ebe8c992928483775"}, {file = "certifi-2026.7.22.tar.gz", hash = "sha256:741e2c3b351ddf169a738da9f2c048608ff7f2c5cc02f1ebc6b118bb090d5d55"}, @@ -250,8 +237,6 @@ version = "0.4.6" description = "Cross-platform colored terminal text." optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,>=2.7" -groups = ["dev"] -markers = "sys_platform == \"win32\"" files = [ {file = "colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6"}, {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"}, @@ -263,8 +248,6 @@ version = "1.3.1" description = "Backport of PEP 654 (exception groups)" optional = false python-versions = ">=3.7" -groups = ["main", "dev"] -markers = "python_version == \"3.10\"" files = [ {file = "exceptiongroup-1.3.1-py3-none-any.whl", hash = "sha256:a7a39a3bd276781e98394987d3a5701d0c4edffb633bb7a5144577f82c773598"}, {file = "exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219"}, @@ -282,7 +265,6 @@ version = "2.1.2" description = "execnet: rapid multi-Python deployment" optional = false python-versions = ">=3.8" -groups = ["dev"] files = [ {file = "execnet-2.1.2-py3-none-any.whl", hash = "sha256:67fba928dd5a544b783f6056f449e5e3931a5c378b128bc18501f7ea79e296ec"}, {file = "execnet-2.1.2.tar.gz", hash = "sha256:63d83bfdd9a23e35b9c6a3261412324f964c2ec8dcd8d3c6916ee9373e0befcd"}, @@ -291,13 +273,38 @@ files = [ [package.extras] testing = ["hatch", "pre-commit", "pytest", "tox"] +[[package]] +name = "fakeredis" +version = "2.38.0" +description = "Python implementation of redis API, can be used for testing purposes." +optional = false +python-versions = ">=3.8" +files = [ + {file = "fakeredis-2.38.0-py3-none-any.whl", hash = "sha256:d9fb0518c4eaa35f1f2c94df6b4a4c97ff3ca9f43d6cc8112317a9037d244301"}, + {file = "fakeredis-2.38.0.tar.gz", hash = "sha256:d2abfd24652f86501044499bf08c9d639db050f695eefc06b8b8b6f0bb24dbd6"}, +] + +[package.dependencies] +lupa = {version = ">=2.1", optional = true, markers = "extra == \"lua\""} +redis = ">=4.3" +sortedcontainers = ">=2" +typing-extensions = {version = ">=4.7", markers = "python_version < \"3.11\""} + +[package.extras] +bf = ["pyprobables (>=0.6)"] +cf = ["pyprobables (>=0.6)"] +json = ["jsonpath-ng (>=1.6)"] +lua = ["lupa (>=2.1)"] +probabilistic = ["pyprobables (>=0.6)"] +valkey = ["valkey (>=6)"] +vectorset = ["jsonpath-ng (>=1.6)", "numpy (>=2.4.0)"] + [[package]] name = "frozenlist" version = "1.8.0" description = "A list-like structure which implements collections.abc.MutableSequence" optional = false python-versions = ">=3.9" -groups = ["main", "dev"] files = [ {file = "frozenlist-1.8.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:b37f6d31b3dcea7deb5e9696e529a6aa4a898adc33db82da12e4c60a7c4d2011"}, {file = "frozenlist-1.8.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:ef2b7b394f208233e471abc541cc6991f907ffd47dc72584acee3147899d6565"}, @@ -430,7 +437,6 @@ files = [ {file = "frozenlist-1.8.0-py3-none-any.whl", hash = "sha256:0c18a16eab41e82c295618a77502e17b195883241c563b00f0aa5106fc4eaa0d"}, {file = "frozenlist-1.8.0.tar.gz", hash = "sha256:3ede829ed8d842f6cd48fc7081d7a41001a56f1f38603f9d49bf3020d59a31ad"}, ] -markers = {main = "extra == \"aiohttp\""} [[package]] name = "h11" @@ -438,7 +444,6 @@ version = "0.16.0" description = "A pure-Python, bring-your-own-I/O implementation of HTTP/1.1" optional = false python-versions = ">=3.8" -groups = ["main"] files = [ {file = "h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86"}, {file = "h11-0.16.0.tar.gz", hash = "sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1"}, @@ -450,7 +455,6 @@ version = "1.0.9" description = "A minimal low-level HTTP client." optional = false python-versions = ">=3.8" -groups = ["main"] files = [ {file = "httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55"}, {file = "httpcore-1.0.9.tar.gz", hash = "sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8"}, @@ -472,7 +476,6 @@ version = "0.28.1" description = "The next generation HTTP client." optional = false python-versions = ">=3.8" -groups = ["main"] files = [ {file = "httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad"}, {file = "httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc"}, @@ -485,7 +488,7 @@ httpcore = "==1.*" idna = "*" [package.extras] -brotli = ["brotli ; platform_python_implementation == \"CPython\"", "brotlicffi ; platform_python_implementation != \"CPython\""] +brotli = ["brotli", "brotlicffi"] cli = ["click (==8.*)", "pygments (==2.*)", "rich (>=10,<14)"] http2 = ["h2 (>=3,<5)"] socks = ["socksio (==1.*)"] @@ -497,8 +500,6 @@ version = "0.1.12" description = "Aiohttp transport for HTTPX" optional = true python-versions = ">=3.8" -groups = ["main"] -markers = "extra == \"aiohttp\"" files = [ {file = "httpx_aiohttp-0.1.12-py3-none-any.whl", hash = "sha256:5b0eac39a7f360fa7867a60bcb46bb1024eada9c01cbfecdb54dc1edb3fb7141"}, {file = "httpx_aiohttp-0.1.12.tar.gz", hash = "sha256:81feec51fd82c0ecfa0e9aaf1b1a6c2591260d5e2bcbeb7eb0277a78e610df2c"}, @@ -514,7 +515,6 @@ version = "3.19" description = "Internationalized Domain Names in Applications (IDNA)" optional = false python-versions = ">=3.9" -groups = ["main", "dev"] files = [ {file = "idna-3.19-py3-none-any.whl", hash = "sha256:815e7be7a7806d54abb586dc943addc79e8b2ee16915059658cbeff4b1b43bf4"}, {file = "idna-3.19.tar.gz", hash = "sha256:5e0811a4383b21dc5838069f801c4fb62113b7447663d2530d2bd6e77b49bf15"}, @@ -529,19 +529,93 @@ version = "2.3.0" description = "brain-dead simple config-ini parsing" optional = false python-versions = ">=3.10" -groups = ["dev"] files = [ {file = "iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12"}, {file = "iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730"}, ] +[[package]] +name = "lupa" +version = "2.8" +description = "Python wrapper around Lua and LuaJIT" +optional = false +python-versions = ">=3.8" +files = [ + {file = "lupa-2.8-cp310-abi3-win32.whl", hash = "sha256:c2a5fd15dc62374e1661a55f01744c9ec1c56f291ba4a0749d3af2174556e78f"}, + {file = "lupa-2.8-cp310-abi3-win_arm64.whl", hash = "sha256:9e304fb1c50cf23fd8882afbe1aa87525ef8a72667bcab3b37b2bbb2bc542269"}, + {file = "lupa-2.8-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:97bd01e90b8031e56a5fd5bb70605aea09f1dba675c1140308a52780f93d06f1"}, + {file = "lupa-2.8-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0b5ebe1a13c45767919c86750b84fe2da9f6288b6f3cea4ce7660bb2abc9d921"}, + {file = "lupa-2.8-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:097e7d0f1719a88020b67c82e05d53d7973c166952393afcecfd8434c7e19a15"}, + {file = "lupa-2.8-cp310-cp310-win_amd64.whl", hash = "sha256:7bb223ee8f72d0dc076b0d65296ee72f1c69450f9d2fed5315f7707d98c4a03d"}, + {file = "lupa-2.8-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:b12e43c1fb787189dfc28cd604aef0baa2cb95e27da19498d520361d0ace070a"}, + {file = "lupa-2.8-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f6f603391dffb256e36a79fd2044084d5f4b8a0a4c0e5ad291cd3ab3aaf1fd0a"}, + {file = "lupa-2.8-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9f6f41c91366e7d0d474f87d81c1274af861f40812bf729c9f97ab4c8f3c7ac8"}, + {file = "lupa-2.8-cp311-cp311-win_amd64.whl", hash = "sha256:f5a6af145b0ea818f01d27bfe2583a4b538570bef61d22c8773e0eccf011234c"}, + {file = "lupa-2.8-cp312-abi3-macosx_10_13_x86_64.whl", hash = "sha256:f4342f4de76ae7ce2ab0672d36003bdb7e1a33252f293b569298ddd792e70e33"}, + {file = "lupa-2.8-cp312-abi3-manylinux2010_i686.manylinux_2_12_i686.manylinux_2_28_i686.whl", hash = "sha256:4203fa1659315e939a5304e75001b8cc14234fb3cbb3ed86c049b0cc5d90fcee"}, + {file = "lupa-2.8-cp312-abi3-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:81f2d843ce668b653146c007467570210ae44be51dac6926666c51d49536f307"}, + {file = "lupa-2.8-cp312-abi3-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:d3d0cde2c77588d1c60875a4f34f059513476c6e1775351897195b51e0f3df08"}, + {file = "lupa-2.8-cp312-abi3-manylinux_2_34_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:9e0d11b8f3a8dac6413f704fef7161d048bb10c58bdac6cbffa5e60efa56e9a3"}, + {file = "lupa-2.8-cp312-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:54cff414f21f8cd8c6be4aae52541f3b9cd39602b59e3a3db9b5c9f9f674ff18"}, + {file = "lupa-2.8-cp312-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:24b4d8af5558e549b70daf1547f5c1c1d664ecea9fc790f83efe5d75e9a93797"}, + {file = "lupa-2.8-cp312-abi3-musllinux_1_2_i686.whl", hash = "sha256:ce86dff1ee7f7cf45f5622065ae991949dd7bb1703581cbc58a630137bb7ccf9"}, + {file = "lupa-2.8-cp312-abi3-musllinux_1_2_ppc64le.whl", hash = "sha256:f4d01b2a08c70bbb883a9e082b6b36b89121ed5910b710f1ba11c73295ff4fba"}, + {file = "lupa-2.8-cp312-abi3-musllinux_1_2_riscv64.whl", hash = "sha256:7f210d5a8353e510ea1199c42cf3cbdd630553bf2bc8fb4c00fea06fdec7c798"}, + {file = "lupa-2.8-cp312-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:4f81a02806e7c7ad26d8c6fa222c8bef1b0c1b124347c879be880b41339d41e4"}, + {file = "lupa-2.8-cp312-abi3-win32.whl", hash = "sha256:360056453a7a4eaa4ac5a204c31a5a014b1eb2ee5490603234d2ba831684f1f2"}, + {file = "lupa-2.8-cp312-abi3-win_arm64.whl", hash = "sha256:1628371c6592a6d5650497a9e31fb2bb3a7e9883c1f301d1111265e484045af9"}, + {file = "lupa-2.8-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:450650f91c48c2415b0d59ab3abfcfda3b6efb5b858205f4d4bda8ad141fa529"}, + {file = "lupa-2.8-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:27044f3363047f946b3d3aab9157cbd172b3538ada9ec1baef43432bf7d03a78"}, + {file = "lupa-2.8-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8cf4f064a0e5531afce2d7d750120c10c10f9529139af6ca6150d13151034398"}, + {file = "lupa-2.8-cp312-cp312-win_amd64.whl", hash = "sha256:281bedc5deb92d31e649a3552edd662449365a635904fa4d5cb4509c7245e34e"}, + {file = "lupa-2.8-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:45fc9da0145ecb0083ef5ff9975116cc784bd0258bdc2bd131ba15483ce18398"}, + {file = "lupa-2.8-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:58e18afed57955b41130e269c78f53d4123ab86e236b53816f4cbffa25cb5d30"}, + {file = "lupa-2.8-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fc47f536ac13a79cef47d29a2b205576a22841f042a2bcec1676b95806e7706a"}, + {file = "lupa-2.8-cp313-cp313-win_amd64.whl", hash = "sha256:ce9404c661dbac65cc9bed351ad45e797af93d30d70be309a3fa8209ac86d93b"}, + {file = "lupa-2.8-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:348c3f8ecabb6324dcbc05c2740d762ef8fcec7b06c79e45262ab97a217684e3"}, + {file = "lupa-2.8-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:951496471056061598a7d1729a6cdf48d662fec777a9f2d8aa5a1e62fd30e5a5"}, + {file = "lupa-2.8-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a591b9947ca347b41a63370e121d6e2b1458fe6dde9ae065029ec10a37f25ff4"}, + {file = "lupa-2.8-cp314-cp314-win_amd64.whl", hash = "sha256:3903c9cf628dae2f56405503247b77a61a3a61bd2dda470e336950c74776d55d"}, + {file = "lupa-2.8-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:f711a8ab0486b9ac6fdda94a22ddcfbc9f0d4a27e3a8cf1bf79c6e48b33017c1"}, + {file = "lupa-2.8-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:dc51250e76367a3e27fcd01dc769b9bfcbbc34f48df48dde53d6af6e75b7eaa5"}, + {file = "lupa-2.8-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f8a22088a552828958603323f0a5c4b3e11e03b75d0bf4c965ef879de9b60a8d"}, + {file = "lupa-2.8-cp314-cp314t-win32.whl", hash = "sha256:4f7c553c1d8cfffbe85d81daef730d12cae4b6002d457542914da0ac8a1145b3"}, + {file = "lupa-2.8-cp314-cp314t-win_amd64.whl", hash = "sha256:d8766aff03a78c80ad2d188a8bdb216de5ec838359cd87e05bbdfa56394a6105"}, + {file = "lupa-2.8-cp314-cp314t-win_arm64.whl", hash = "sha256:91d622777febda3ab1bed1d45295f2f32a4680c7b3d7caf8c669998ed5c44118"}, + {file = "lupa-2.8-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:81b283bfb13cc43fa4910fc98ec110ab861bcb39680f48b266f99d6e3be1049e"}, + {file = "lupa-2.8-cp38-cp38-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5caf45d15d424cee52fd67341e96e2b1dde0658ae90eb156ac56aa0d8330bc38"}, + {file = "lupa-2.8-cp38-cp38-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:33e7e5aebca64b154b0a1679caf79e19254ff37bba51e87abab6848f97cb2de1"}, + {file = "lupa-2.8-cp38-cp38-win32.whl", hash = "sha256:e8d4f4dd4acf4a0e42adc6b1ad220e1c86fe3028402c2f78bd0728a6d241bbe9"}, + {file = "lupa-2.8-cp38-cp38-win_amd64.whl", hash = "sha256:1ac2b1ec7504e6148cba1bc35ac36c74d18a0ca6d367ffe7e78a3773c2694c0e"}, + {file = "lupa-2.8-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:b036738282a5acd2e71fdddb317c9df8b87c1673aa57f403d05fcc2be8abc4ba"}, + {file = "lupa-2.8-cp39-abi3-manylinux2010_i686.manylinux_2_12_i686.manylinux_2_28_i686.whl", hash = "sha256:ac6b6e8d0e617e26a98cbb44880bcd75de5d32b3ad7b3b3793583909292b47ed"}, + {file = "lupa-2.8-cp39-abi3-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:ba3a7dd839f90c3d2e53bebe3c192b1f3f9fd720a6781256405123211fd0dce6"}, + {file = "lupa-2.8-cp39-abi3-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:d7edb13a7a5250b5c6c22d1495d9e842b5c9fc5081c8fe6b5efe2112fe3e41f9"}, + {file = "lupa-2.8-cp39-abi3-manylinux_2_34_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:891f72e0bffbed1e4175f975aeb2a083956586a100066525e1be485f617f7b25"}, + {file = "lupa-2.8-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:a295f87b5b7ebbfd5191932e8cb0e51df3c7769101ac6b6c7d7c9fb27bfd1307"}, + {file = "lupa-2.8-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:4fe5d7a810b64ea8511eb885fc8cdde042ee5ff7b7d08ae78f32449756acb177"}, + {file = "lupa-2.8-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:bfc470012ef66ad064c7bd77416af03a3452ef630b04b9012595ea13f2e54518"}, + {file = "lupa-2.8-cp39-abi3-musllinux_1_2_ppc64le.whl", hash = "sha256:250e035fdaffe8c87093e3ebc206ac29a26131b1568ea711d780c26001ce96e7"}, + {file = "lupa-2.8-cp39-abi3-musllinux_1_2_riscv64.whl", hash = "sha256:b9bddb09acfffb4f828f790f444b11dc0cca591afea1a244d9329eea2d20c003"}, + {file = "lupa-2.8-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:2e64acbbd47e9b82a64405a39e0d2b36a5a7dad8ab41c0f3437f572f7d282ba3"}, + {file = "lupa-2.8-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:f6ddca4774d5ca451768a95e378a3aa041076e29f4613b8562f8e98efb6690fd"}, + {file = "lupa-2.8-cp39-cp39-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3ffcfd8e19f943ad459136b3f60f085ae4948f024192a93ca4b4ac3023ec88d8"}, + {file = "lupa-2.8-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9f3f3955f65f9fde2dc6eda3041ccd394cf54d4bf083f0cdf6feb3d58e5f38d3"}, + {file = "lupa-2.8-cp39-cp39-win32.whl", hash = "sha256:9e76e45057cfcaa20ee3422c2289a91f9d51783d020da3570ee226de8f6e71cd"}, + {file = "lupa-2.8-cp39-cp39-win_amd64.whl", hash = "sha256:6fbcc9911f05c67affbd225fc024268e61e98a18ad1b1c2aed6c8796e4056554"}, + {file = "lupa-2.8-cp39-cp39-win_arm64.whl", hash = "sha256:6c817d5421094507662e5f8feb8cd1e154c10879921c06079b6063be9d8f33c5"}, + {file = "lupa-2.8-pp311-pypy311_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:32e4e5103bbddcdd2458fb2ccae6c8ba11c9997c711d7e379e0d45551d109c76"}, + {file = "lupa-2.8-pp311-pypy311_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7667001804657496dee9feced2daae5000b4604a3218dd8e6b7b754982ba88b8"}, + {file = "lupa-2.8-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:86f6f668966965b15247dc32d064cfe7be67b71e584ccfacbe2f637575296878"}, + {file = "lupa-2.8.tar.gz", hash = "sha256:d8022641b9ec8ecf2c5ecbe9f47e5a70e0b87c4b5ae921b92cb02a638e0acd08"}, +] + [[package]] name = "multidict" version = "6.7.1" description = "multidict implementation" optional = false python-versions = ">=3.9" -groups = ["main", "dev"] files = [ {file = "multidict-6.7.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:c93c3db7ea657dd4637d57e74ab73de31bccefe144d3d4ce370052035bc85fb5"}, {file = "multidict-6.7.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:974e72a2474600827abaeda71af0c53d9ebbc3c2eb7da37b37d7829ae31232d8"}, @@ -690,7 +764,6 @@ files = [ {file = "multidict-6.7.1-py3-none-any.whl", hash = "sha256:55d97cc6dae627efa6a6e548885712d4864b81110ac76fa4e534c03819fa4a56"}, {file = "multidict-6.7.1.tar.gz", hash = "sha256:ec6652a1bee61c53a3e5776b6049172c53b6aaba34f18c9ad04f82712bac623d"}, ] -markers = {main = "extra == \"aiohttp\""} [package.dependencies] typing-extensions = {version = ">=4.1.0", markers = "python_version < \"3.11\""} @@ -701,7 +774,6 @@ version = "1.13.0" description = "Optional static typing for Python" optional = false python-versions = ">=3.8" -groups = ["dev"] files = [ {file = "mypy-1.13.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:6607e0f1dd1fb7f0aca14d936d13fd19eba5e17e1cd2a14f808fa5f8f6d8f60a"}, {file = "mypy-1.13.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8a21be69bd26fa81b1f80a61ee7ab05b076c674d9b18fb56239d72e21d9f4c80"}, @@ -755,7 +827,6 @@ version = "1.1.0" description = "Type system extensions for programs checked with the mypy type checker." optional = false python-versions = ">=3.8" -groups = ["dev"] files = [ {file = "mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505"}, {file = "mypy_extensions-1.1.0.tar.gz", hash = "sha256:52e68efc3284861e772bbcd66823fde5ae21fd2fdb51c62a211403730b916558"}, @@ -767,7 +838,6 @@ version = "26.3" description = "Core utilities for Python packages" optional = false python-versions = ">=3.9" -groups = ["dev"] files = [ {file = "packaging-26.3-py3-none-any.whl", hash = "sha256:d7193f7c8e4e93f444fde0262bf90af30e16fa0ad0ad44cb553c87339b23cd1c"}, {file = "packaging-26.3.tar.gz", hash = "sha256:94edc256424af38762eb31306eed28beb9f0efc50a8837492c9d6fd6004aed79"}, @@ -779,7 +849,6 @@ version = "1.6.0" description = "plugin and hook calling mechanisms for python" optional = false python-versions = ">=3.9" -groups = ["dev"] files = [ {file = "pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746"}, {file = "pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3"}, @@ -795,7 +864,6 @@ version = "0.5.2" description = "Accelerated property cache" optional = false python-versions = ">=3.10" -groups = ["main", "dev"] files = [ {file = "propcache-0.5.2-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:d5a81be28596d6559f6131ef33e10200de6e17643b3c74ce03f9eb103be6ae8b"}, {file = "propcache-0.5.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:29cbaac5ea0212663e6845e04b5e188d5a6ae6dd919810ac835bf1d3b42c3f4c"}, @@ -919,7 +987,6 @@ files = [ {file = "propcache-0.5.2-py3-none-any.whl", hash = "sha256:be1ddfcbb376e3de5d2e2db1d58d6d67463e6b4f9f040c000de8e300295465fe"}, {file = "propcache-0.5.2.tar.gz", hash = "sha256:01c4fc7480cd0598bb4b57022df55b9ca296da7fc5a8760bd8451a7e63a7d427"}, ] -markers = {main = "extra == \"aiohttp\""} [[package]] name = "pydantic" @@ -927,7 +994,6 @@ version = "2.13.5" description = "Data validation using Python type hints" optional = false python-versions = ">=3.9" -groups = ["main"] files = [ {file = "pydantic-2.13.5-py3-none-any.whl", hash = "sha256:346a034f080da3755d8e9cb5e00e8b07de1d39e4f6e2c87d8ab7cafa0b269a73"}, {file = "pydantic-2.13.5.tar.gz", hash = "sha256:51a9c5f7b2f8e636f04c6cada605d9b6a3bf1348fdf945a3d8869b19bba0ee08"}, @@ -941,7 +1007,7 @@ typing-inspection = ">=0.4.2" [package.extras] email = ["email-validator (>=2.0.0)"] -timezone = ["tzdata ; python_version >= \"3.9\" and platform_system == \"Windows\""] +timezone = ["tzdata"] [[package]] name = "pydantic-core" @@ -949,7 +1015,6 @@ version = "2.46.5" description = "Core functionality for Pydantic validation and serialization" optional = false python-versions = ">=3.9" -groups = ["main"] files = [ {file = "pydantic_core-2.46.5-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:657b40d6240c0a7b6a64b30f22d1e3aa631c7e846c621b0c0f6d1d75e2e15ea6"}, {file = "pydantic_core-2.46.5-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:ecb42011e12ee19cafbc312887cbf3546959fe02fbad44f272d4be5baa997615"}, @@ -1082,7 +1147,6 @@ version = "2.13.0" description = "JSON Web Token implementation in Python" optional = false python-versions = ">=3.9" -groups = ["dev"] files = [ {file = "pyjwt-2.13.0-py3-none-any.whl", hash = "sha256:66adcc2aff09b3f1bbd95fc1e1577df8ac8723c978552fd43304c8a290ac5728"}, {file = "pyjwt-2.13.0.tar.gz", hash = "sha256:41571c89ca91598c79e8ef18a2d07367d4810fbbd6f637794879baf1b7703423"}, @@ -1100,7 +1164,6 @@ version = "7.4.4" description = "pytest: simple powerful testing with Python" optional = false python-versions = ">=3.7" -groups = ["dev"] files = [ {file = "pytest-7.4.4-py3-none-any.whl", hash = "sha256:b090cdf5ed60bf4c45261be03239c2c1c22df034fbffe691abe93cd80cea01d8"}, {file = "pytest-7.4.4.tar.gz", hash = "sha256:2cf0005922c6ace4a3e2ec8b4080eb0d9753fdc93107415332f50ce9e7994280"}, @@ -1123,7 +1186,6 @@ version = "0.23.8" description = "Pytest support for asyncio" optional = false python-versions = ">=3.8" -groups = ["dev"] files = [ {file = "pytest_asyncio-0.23.8-py3-none-any.whl", hash = "sha256:50265d892689a5faefb84df80819d1ecef566eb3549cf915dfb33569359d1ce2"}, {file = "pytest_asyncio-0.23.8.tar.gz", hash = "sha256:759b10b33a6dc61cce40a8bd5205e302978bbbcc00e279a8b61d9a6a3c82e4d3"}, @@ -1142,7 +1204,6 @@ version = "3.8.0" description = "pytest xdist plugin for distributed testing, most importantly across multiple CPUs" optional = false python-versions = ">=3.9" -groups = ["dev"] files = [ {file = "pytest_xdist-3.8.0-py3-none-any.whl", hash = "sha256:202ca578cfeb7370784a8c33d6d05bc6e13b4f25b5053c30a152269fd10f0b88"}, {file = "pytest_xdist-3.8.0.tar.gz", hash = "sha256:7e578125ec9bc6050861aa93f2d59f1d8d085595d6551c2c90b6f4fad8d3a9f1"}, @@ -1163,7 +1224,6 @@ version = "2.9.0.post0" description = "Extensions to the standard Python datetime module" optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,>=2.7" -groups = ["dev"] files = [ {file = "python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3"}, {file = "python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427"}, @@ -1178,7 +1238,6 @@ version = "5.3.1" description = "Python client for Redis database and key-value store" optional = false python-versions = ">=3.8" -groups = ["dev"] files = [ {file = "redis-5.3.1-py3-none-any.whl", hash = "sha256:dc1909bd24669cc31b5f67a039700b16ec30571096c5f1f0d9d2324bff31af97"}, {file = "redis-5.3.1.tar.gz", hash = "sha256:ca49577a531ea64039b5a36db3d6cd1a0c7a60c34124d46924a45b956e8cf14c"}, @@ -1198,7 +1257,6 @@ version = "0.11.5" description = "An extremely fast Python linter and code formatter, written in Rust." optional = false python-versions = ">=3.7" -groups = ["dev"] files = [ {file = "ruff-0.11.5-py3-none-linux_armv6l.whl", hash = "sha256:2561294e108eb648e50f210671cc56aee590fb6167b594144401532138c66c7b"}, {file = "ruff-0.11.5-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:ac12884b9e005c12d0bd121f56ccf8033e1614f736f766c118ad60780882a077"}, @@ -1226,20 +1284,28 @@ version = "1.17.0" description = "Python 2 and 3 compatibility utilities" optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,>=2.7" -groups = ["dev"] files = [ {file = "six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274"}, {file = "six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81"}, ] +[[package]] +name = "sortedcontainers" +version = "2.4.0" +description = "Sorted Containers -- Sorted List, Sorted Dict, Sorted Set" +optional = false +python-versions = "*" +files = [ + {file = "sortedcontainers-2.4.0-py2.py3-none-any.whl", hash = "sha256:a163dcaede0f1c021485e957a39245190e74249897e2ae4b2aa38595db237ee0"}, + {file = "sortedcontainers-2.4.0.tar.gz", hash = "sha256:25caa5a06cc30b6b83d11423433f65d1f9d76c4c6a0c90e3379eaa43b9bfdb88"}, +] + [[package]] name = "tomli" version = "2.4.1" description = "A lil' TOML parser" optional = false python-versions = ">=3.8" -groups = ["dev"] -markers = "python_version == \"3.10\"" files = [ {file = "tomli-2.4.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:f8f0fc26ec2cc2b965b7a3b87cd19c5c6b8c5e5f436b984e85f486d652285c30"}, {file = "tomli-2.4.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:4ab97e64ccda8756376892c53a72bd1f964e519c77236368527f758fbc36a53a"}, @@ -1296,7 +1362,6 @@ version = "2.9.0.20260807" description = "Typing stubs for python-dateutil" optional = false python-versions = ">=3.10" -groups = ["dev"] files = [ {file = "types_python_dateutil-2.9.0.20260807-py3-none-any.whl", hash = "sha256:54aa3707350ed7a9cc0776fd2f6739679d6967d11b40150985e81edcb86df4db"}, {file = "types_python_dateutil-2.9.0.20260807.tar.gz", hash = "sha256:e0b8a90d464c8684c66b7b8e4556d9074afdddcc56ca45323f0987134f9e7034"}, @@ -1308,7 +1373,6 @@ version = "4.16.0" description = "Backported and Experimental Type Hints for Python 3.9+" optional = false python-versions = ">=3.9" -groups = ["main", "dev"] files = [ {file = "typing_extensions-4.16.0-py3-none-any.whl", hash = "sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8"}, {file = "typing_extensions-4.16.0.tar.gz", hash = "sha256:dc983d19a509c94dba722ee6abd33940f7c05a89e243c47e907eb4db6f1a43e5"}, @@ -1320,7 +1384,6 @@ version = "0.4.4" description = "Runtime typing introspection tools" optional = false python-versions = ">=3.10" -groups = ["main"] files = [ {file = "typing_inspection-0.4.4-py3-none-any.whl", hash = "sha256:65b8397ba37ccbce054456aaccddfc91e6e3083c92824df348d96ca832f3f147"}, {file = "typing_inspection-0.4.4.tar.gz", hash = "sha256:547274fa6b0a561ccf549cc9524b999a578e737d015d8709d021f9d0d13bea47"}, @@ -1335,8 +1398,6 @@ version = "48.0.0" description = "A WebAssembly runtime powered by Wasmtime" optional = true python-versions = ">=3.9" -groups = ["main"] -markers = "extra == \"datastream\" or extra == \"rulesengine\"" files = [ {file = "wasmtime-48.0.0-py3-none-android_26_arm64_v8a.whl", hash = "sha256:a55abf132fe238b843a963c68cd1a30d8f686c1bc75d8fbf042d8b7a1d51ee36"}, {file = "wasmtime-48.0.0-py3-none-android_26_x86_64.whl", hash = "sha256:d8e94276ff6c0c5ce73ee16ccbacb00b3512a4b3a664749380705d81ee06a23c"}, @@ -1361,8 +1422,6 @@ version = "16.1.1" description = "An implementation of the WebSocket Protocol (RFC 6455 & 7692)" optional = true python-versions = ">=3.10" -groups = ["main"] -markers = "extra == \"datastream\"" files = [ {file = "websockets-16.1.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:49ae99bdfcae803a885c926bf14f886196e84925395bb3f568fef5c0f0979d7d"}, {file = "websockets-16.1.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:5bfd1ac19b1b9986a9c95a82d5e23a391ebb09e12c34d7be6094b86efcc35731"}, @@ -1481,7 +1540,6 @@ version = "1.24.5" description = "Yet another URL library" optional = false python-versions = ">=3.10" -groups = ["main", "dev"] files = [ {file = "yarl-1.24.5-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:88f50c94e21a0a7f14042c015b0eba1881af78562e7bf007e0033e624da59750"}, {file = "yarl-1.24.5-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:6efbccc3d7f75d5b03105172a8dc86d82ba4da86817952529dd93185f4a88be2"}, @@ -1588,7 +1646,6 @@ files = [ {file = "yarl-1.24.5-py3-none-any.whl", hash = "sha256:a33700d13d9b7d84fd10947b09ff69fb9a792e519c8cb9764a3ca70baa6c23a7"}, {file = "yarl-1.24.5.tar.gz", hash = "sha256:e81b83143bee16329c23db3c1b2d82b29892fcbcb849186d2f6e98a5abe9a57f"}, ] -markers = {main = "extra == \"aiohttp\""} [package.dependencies] idna = ">=2.0" @@ -1601,6 +1658,6 @@ datastream = ["wasmtime", "websockets"] rulesengine = ["wasmtime"] [metadata] -lock-version = "2.1" +lock-version = "2.0" python-versions = "^3.10" -content-hash = "d12f2b0714d21d258fab37327fffd7eed95cbda40524b341fc8e0479072c5342" +content-hash = "b97aece3e401780396eb371b240c4c82f59799ccece668843b9f37dd9b0af950" diff --git a/pyproject.toml b/pyproject.toml index 6a88bccb..21cebff5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,6 +62,7 @@ types-python-dateutil = "^2.9.0.20240316" ruff = "==0.11.5" aiohttp = ">=3.14.0,<4" redis = ">=5.0,<6" +fakeredis = {extras = ["lua"], version = "^2.38.0"} [tool.pytest.ini_options] testpaths = [ "tests" ] diff --git a/src/schematic/cache/redis.py b/src/schematic/cache/redis.py index 895199e1..68c085a5 100644 --- a/src/schematic/cache/redis.py +++ b/src/schematic/cache/redis.py @@ -35,6 +35,16 @@ def __init__( self._prefix = prefix self._default_ttl_ms = default_ttl_ms + @property + def client(self) -> Any: + """The Redis client this cache was built on, so other parts of the SDK + (credit leases) can share the connection instead of opening a second.""" + return self._client + + @property + def prefix(self) -> str: + return self._prefix + def _prefixed(self, key: str) -> str: return f"{self._prefix}:{key}" diff --git a/src/schematic/client.py b/src/schematic/client.py index adf07b9b..45d0bcf4 100644 --- a/src/schematic/client.py +++ b/src/schematic/client.py @@ -1,18 +1,53 @@ +import asyncio import atexit import datetime as dt import logging +import math +import time from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable, Dict, List, Literal, Optional, Union import httpx from .base_client import AsyncBaseSchematic, BaseSchematic -from .cache import DEFAULT_CACHE_SIZE, DEFAULT_CACHE_TTL, AsyncCacheProvider, CacheProvider, LocalCache +from .cache import ( + DEFAULT_CACHE_SIZE, + DEFAULT_CACHE_TTL, + AsyncCacheProvider, + CacheProvider, + LocalCache, + RedisCache, +) +from .core.api_error import ApiError +from .core.request_options import RequestOptions from .datastream import DataStreamClient, DataStreamClientOptions +from .errors import PaymentRequiredError from .event_buffer import AsyncEventBuffer, EventBuffer from .event_capture import AsyncEventCaptureClient, EventCaptureClient from .http_client import AsyncOfflineHTTPClient, OfflineHTTPClient +from .leases import ( + DEFAULT_LEASE_DURATION, + DEFAULT_PREWARM_RESOLVE_TIMEOUT, + CreditCheckDeps, + CreditsWireClient, + InMemoryLeaseStore, + InMemoryReservationStore, + LeaseConfig, + LeaseConfigOverride, + LeaseManager, + LeaseStore, + RedisLeaseStore, + RedisReservationStore, + ReservationStore, + check_with_lease, + consume_reservation_and_build_event, +) +from .leases import build_reservation_track_event as _build_reservation_track_event +from .leases import is_valid_quantity as _is_valid_quantity +from .leases import settled_quantity as _settled_quantity +from .leases.redis_lease_store import DEFAULT_KEY_PREFIX as DEFAULT_LEASE_KEY_PREFIX from .logging import DEFAULT_LOG_LEVEL, LogLevel, get_default_logger from .types import ( + CheckAndReserveFlagResponseData, CheckFlagRequestBody, CheckFlagResponseData, CreateEventRequestBody, @@ -22,6 +57,9 @@ EventBodyIdentifyCompany, EventBodyTrack, FeatureEntitlement, + FlagCheckReservationResponseData, + PreflightEventUsageRequestBody, + PreflightRequestBody, RulesengineCheckFlagResult, ) @@ -32,6 +70,52 @@ REASON_FLAG_NOT_FOUND = "Flag not found - using default value" REASON_ERROR = "Error occurred - using default value" +# Prefix of the deterministic idempotency key carried by the track event that +# settles a reservation, so a duplicate or retried settle is dropped +# server-side instead of billing the usage twice. +RESERVATION_TRACK_IDEMPOTENCY_PREFIX = "lease-reservation:" + +# How long a server-side hold lives when the caller configures no TTL. In +# seconds, like every other duration on this client. +DEFAULT_RESERVATION_TTL = 60.0 + +# How often prewarm re-asks DataStream for a company it is waiting on. +PREWARM_POLL_INTERVAL = 0.1 + +# The longest hold the server will take. A longer configured TTL is clamped to +# it, rather than sent and rejected on every check. +MAX_RESERVATION_TTL = 3600.0 +# Room for the two clocks to disagree. expires_at is computed here and measured +# against the server's own clock, so a TTL sitting exactly on the cap is +# rejected whenever this process runs even slightly ahead. +RESERVATION_TTL_SKEW_ALLOWANCE = 60.0 + +# What the API calls a denial for want of credits. The 402 branch and the +# plain 200-with-value-false branch both report it, so callers have one string +# to match on. +INSUFFICIENT_CREDITS_REASON = "Insufficient credits" + +# Where a credit hold lives for a check() that passes usage. +# - "server": one check-and-reserve API call per check; the server evaluates +# the flag and takes the hold in the same round trip. +# - "client": a lease drawn from the server per company and credit type, with +# each check's hold carved out of it locally. Needs DataStream, so only +# AsyncSchematic can run it. +# - "auto" (default): "client" when DataStream is running, "server" otherwise. +CreditLeaseMode = Literal["client", "server", "auto"] + +# What a check does when it cannot gate: deny ("fail-closed"), or fall back to +# the caller's default value ("fail-open"). +OnAcquireFailure = Literal["fail-closed", "fail-open"] + + +@dataclass +class EventUsage: + """Usage of one event subtype, for preflighting a flag check.""" + + event_subtype: str + quantity: int + @dataclass class CheckFlagOptions: @@ -39,6 +123,401 @@ class CheckFlagOptions: default_value: Optional[Union[bool, Callable[[], bool]]] = None timeout: Optional[float] = None + # Preflight fields: hypothetical usage the flag is evaluated against, so a + # caller can ask "would this action be allowed?" before performing it. + # They mirror the API's PreflightRequestBody. + # + # Quantity applied to any numeric condition met while evaluating the flag. + usage: Optional[int] = None + # Usage of one specific event subtype. Preferred over `usage` when the + # subtype is known, since it only moves conditions measuring that subtype. + event_usage: Optional[EventUsage] = None + # Cost in credits, keyed by credit ID, for callers that already computed + # it. Takes precedence over usage and event_usage for the same credit. + credit_cost: Optional[Dict[str, float]] = None + + +@dataclass +class CreditLeaseConfig: + """Opt in to credit-gated checks (``check`` / ``track_with_reservation``). + + Leave it unset and ``check`` is a plain flag check that holds nothing. + Every duration is in seconds. The knobs below the first two steer client + mode only, and server mode warns at construction when one is set. + """ + + # Where the hold lives. "server" takes it over the check-and-reserve API; + # "client" carves it out of a local lease over DataStream; "auto" picks + # client when DataStream is running and server otherwise. + mode: CreditLeaseMode = "auto" + # How long a hold survives unsettled. Size it above the longest expected + # gap between check() and track_with_reservation(). Anything above the + # server's one hour cap is clamped, in server mode. Client mode keeps it: + # the TTL only drives the local sweeper there. + default_reservation_ttl: float = DEFAULT_RESERVATION_TTL + # Lease lifetime requested at acquire and extend. Default 5 minutes. + default_lease_duration: Optional[float] = None + # Credits requested per acquire, and the minimum extend tranche. Default 10000. + default_lease_size: Optional[float] = None + # Remaining/granted ratio at or below which a background extend fires. Default 0.25. + low_water_mark: Optional[float] = None + # How often expired reservations are swept back to their leases. Default 1 second. + sweep_interval: Optional[float] = None + # How long prewarm() waits for a freshly identified company to surface + # over DataStream. Default 5 seconds; 0 skips the wait. + prewarm_resolve_timeout: Optional[float] = None + # A connected redis.asyncio client for lease and reservation state. Without + # one the SDK reuses the DataStream company cache's Redis, and failing that + # keeps lease state per-process, which gates within one process only. + redis_client: Optional[Any] = None + # Key prefix for lease and reservation keys. Defaults to the DataStream + # cache's prefix when its Redis is reused, else "schematic:", which is what + # the Node SDK uses, so mixed fleets share the same leases. + redis_key_prefix: Optional[str] = None + # Per-credit-type overrides of the four resolvable knobs, keyed by credit + # type ID. Client mode only. + overrides: Optional[Dict[str, LeaseConfigOverride]] = None + + +@dataclass +class CheckOptions: + """Options accepted by ``check``.""" + + # Units of the feature this operation will consume, as any finite + # non-negative number. The check holds usage * consumption_rate credits. + # A check takes at most one hold. + usage: Optional[float] = None + # Event subtype the usage applies to, e.g. "inference_tokens". Needed only + # when the flag meters more than one event. + event_subtype: Optional[str] = None + # What to do when the check cannot gate (API error, unreachable server). + on_acquire_failure: OnAcquireFailure = "fail-closed" + default_value: Optional[Union[bool, Callable[[], bool]]] = None + # Per-check timeout for the API calls this check makes, in seconds. + timeout: Optional[float] = None + + +@dataclass +class Reservation: + """Handle returned by ``check`` when a credit hold was taken. + + Pass it to ``track_with_reservation`` when the work completes. + """ + + id: str + # The lease the hold draws from. Server mode has no lease, so this mirrors + # `id` and the field stays populated for code that reads it. + lease_id: str + mode: Literal["client", "server"] + company_id: str + credit_type_id: str + # Event subtype the settling track event is recorded under. + event_subtype: str + quantity_reserved: float + credits_reserved: float + consumption_rate: float + # When the unspent hold is refunded if nothing settles it. + expires_at: dt.datetime + # Evaluation context the hold was issued for, so the settling track event + # attributes the usage to the same company and user. + company: Optional[Dict[str, str]] = None + user: Optional[Dict[str, str]] = None + + +@dataclass +class CheckResult: + """Result of ``check``.""" + + # Whether the caller may proceed. + allowed: bool + # The flag's boolean value; `allowed` mirrors it outside the credit paths. + value: bool + reason: str + flag_key: str + reservation: Optional[Reservation] = None + entitlement: Optional[FeatureEntitlement] = None + flag_id: Optional[str] = None + error: Optional[str] = None + + +@dataclass +class TrackWithReservationOptions: + """Extras accepted by ``track_with_reservation``.""" + + traits: Optional[Dict[str, Any]] = None + + +def _build_preflight(options: Optional[CheckFlagOptions]) -> Optional[PreflightRequestBody]: + """Build the preflight body for a flag check, or None when the caller set + no preflight field.""" + if options is None: + return None + if options.usage is None and options.event_usage is None and options.credit_cost is None: + return None + return PreflightRequestBody( + credit_cost=options.credit_cost, + event_usage=( + PreflightEventUsageRequestBody( + event_subtype=options.event_usage.event_subtype, + quantity=options.event_usage.quantity, + ) + if options.event_usage is not None + else None + ), + usage=options.usage, + ) + + +def _preflight_quantity(usage: float) -> int: + """Cast a usage onto the integer the preflight body carries. + + A hold can be sized from a fractional usage, but the API's preflight usage + is an integer. A preflight asks an upper-bound question ("would this action + be allowed?"), so a fraction rounds up: the check must not pass on less + usage than the operation is about to record. + """ + return int(usage) if float(usage).is_integer() else math.ceil(usage) + + +def _check_options_to_flag_options(options: Optional[CheckOptions]) -> Optional[CheckFlagOptions]: + """Map credit-aware check options onto plain flag check options. + + With an event subtype the usage goes out as the event_usage pair so the + engine matches it to that subtype's condition; without one it goes out as + the generic usage knob. A usage no hold could be sized from is dropped + instead of threaded, since it would only make the preflight unserializable. + """ + if options is None: + return None + flag_options = CheckFlagOptions(default_value=options.default_value, timeout=options.timeout) + if options.usage is not None and _is_valid_quantity(options.usage): + quantity = _preflight_quantity(options.usage) + if options.event_subtype is not None: + flag_options.event_usage = EventUsage(event_subtype=options.event_subtype, quantity=quantity) + else: + flag_options.usage = quantity + return flag_options + + +# Options that only steer the client-mode lease plumbing, so server mode would +# quietly ignore them. +_CLIENT_ONLY_LEASE_OPTIONS = ( + "default_lease_duration", + "default_lease_size", + "low_water_mark", + "sweep_interval", + "prewarm_resolve_timeout", + "redis_client", + "redis_key_prefix", + "overrides", +) + + +def _mode_uses_leases(mode: CreditLeaseMode, datastream_enabled: bool) -> bool: + """Whether a configured mode wants the local lease plumbing built.""" + if mode == "server": + return False + if mode == "client": + return True + return datastream_enabled + + +def _warn_credit_lease_config( + logger: logging.Logger, + credit_leases: CreditLeaseConfig, + offline: bool, + *, + supports_client_mode: bool, + datastream_enabled: bool = False, +) -> None: + """Say once, at construction, when the configured credit leases will not + gate the way the caller asked.""" + if offline: + logger.warning( + "credit_leases is configured but the client is offline; check() returns flag defaults " + "and holds no credits." + ) + return + # The client-only knobs are ignored wherever the mode lands on server, not + # only where the caller named it: 'auto' without DataStream lands there + # too, and on the sync client it always does. + if not _mode_uses_leases(credit_leases.mode, supports_client_mode and datastream_enabled): + ignored = [name for name in _CLIENT_ONLY_LEASE_OPTIONS if getattr(credit_leases, name) is not None] + if ignored: + logger.warning( + f"credit_leases resolves to server mode, so {', '.join(ignored)} will be ignored; those " + "options only apply to client mode, where leases are carved up locally over DataStream." + ) + if credit_leases.mode == "server": + return + if not supports_client_mode: + if credit_leases.mode == "client": + logger.warning( + "credit_leases.mode is 'client', which needs DataStream, and DataStream is only available on " + "AsyncSchematic; check() falls back to a plain, ungated flag check. Use 'server' (or the 'auto' " + "default) to gate on credits from this client." + ) + return + if datastream_enabled: + return + if credit_leases.mode == "client": + logger.warning( + "credit_leases.mode is 'client' but DataStream is not enabled; check() falls back to plain flag " + "checks with no credit gating. Set use_datastream=True to gate on local leases." + ) + else: + # Not a misconfiguration: auto without DataStream is the server-mode + # default, which gates over the API instead. + logger.info( + "credit_leases is configured and DataStream is not enabled, so credit holds are taken in server " + "mode, one check-and-reserve call per check. Set use_datastream=True for client-side leases." + ) + + +def _resolve_reservation_ttl(logger: logging.Logger, credit_leases: Optional[CreditLeaseConfig]) -> float: + """How long this client asks the server to hold credits for, clamped to + what the server will grant. + + Client mode is exempt: its TTL never reaches the server, it only tells the + local sweeper when to refund an unsettled hold, and the caller may well + want one that outlives an hour. + """ + if credit_leases is None: + return DEFAULT_RESERVATION_TTL + ttl = credit_leases.default_reservation_ttl + if credit_leases.mode == "client": + return ttl + effective = MAX_RESERVATION_TTL - RESERVATION_TTL_SKEW_ALLOWANCE + if ttl > effective: + logger.warning( + f"credit_leases.default_reservation_ttl of {ttl}s is above the server's one hour cap; " + f"server-mode holds will be clamped to {effective}s" + ) + return effective + return ttl + + +def _reservation_request_kwargs(options: CheckOptions) -> Dict[str, Any]: + """Preflight body and request options for a check-and-reserve call, with + the preflight omitted when the caller set nothing. + + Retries are always off. The default policy re-sends on 408, 429 and 5xx, + and this request carries no idempotency key, so a 502 arriving after the + server committed the hold would take a second one against the same + balance. A caller that wants the call retried can retry the check. + """ + kwargs: Dict[str, Any] = {} + preflight = _build_preflight(_check_options_to_flag_options(options)) + if preflight is not None: + kwargs["preflight"] = preflight + request_options: RequestOptions = {"max_retries": 0} + if options.timeout is not None: + request_options["timeout"] = options.timeout + kwargs["request_options"] = request_options + return kwargs + + +def _is_payment_required(error: Exception) -> bool: + """Whether an error is the server saying the credits are not there. + + The generated features client has no 402 branch, unlike the credits one, so + a real 402 from check-and-reserve arrives as the base ApiError carrying the + status code rather than as PaymentRequiredError. + """ + if isinstance(error, PaymentRequiredError): + return True + return isinstance(error, ApiError) and error.status_code == 402 + + +def _payment_required_message(error: Exception) -> str: + """The server's own explanation for a 402, when the body carries one.""" + body = getattr(error, "body", None) + message = body.get("error") if isinstance(body, dict) else getattr(body, "error", None) + if isinstance(message, str) and message: + return message + return str(error) + + +def _payment_required_result(flag_key: str, error: Exception) -> CheckResult: + """A 402 is the server's answer, not a failure to answer: it knows the + credits are not there. Deny whatever on_acquire_failure says.""" + return CheckResult( + allowed=False, + value=False, + reason=INSUFFICIENT_CREDITS_REASON, + flag_key=flag_key, + error=_payment_required_message(error), + ) + + +def _reservation_check_result(flag_key: str, data: CheckAndReserveFlagResponseData) -> CheckResult: + """The flag verdict a check-and-reserve response carries, before any hold + is attached to it.""" + return CheckResult( + allowed=data.value, + value=data.value, + reason=data.reason, + flag_key=data.flag or flag_key, + entitlement=data.entitlement, + flag_id=data.flag_id, + error=data.error, + ) + + +def _server_reservation( + held: FlagCheckReservationResponseData, + event_subtype: str, + company: Optional[Dict[str, str]], + user: Optional[Dict[str, str]], +) -> Reservation: + """The caller's handle on a hold the server took.""" + return Reservation( + id=held.id, + lease_id=held.id, + mode="server", + company_id=held.company_id, + credit_type_id=held.credit_type_id, + event_subtype=event_subtype, + quantity_reserved=held.quantity_reserved, + credits_reserved=held.credits_reserved, + consumption_rate=held.consumption_rate, + expires_at=held.expires_at, + company=company, + user=user, + ) + + +def _server_failure_result( + flag_key: str, options: CheckOptions, reason: str, default_value: bool, +) -> CheckResult: + """Resolve a check that could not gate. ``fail-closed`` denies; + ``fail-open`` returns the caller's resolved default, since server mode has + no local engine to re-run with the balance assumed sufficient.""" + if options.on_acquire_failure == "fail-closed": + return CheckResult(allowed=False, value=False, reason=reason, flag_key=flag_key, error=reason) + return CheckResult( + allowed=default_value, + value=default_value, + reason=f"{reason}_fail_open", + flag_key=flag_key, + error=reason, + ) + + +def _missing_event_subtype_result(options: CheckOptions, result: CheckResult) -> CheckResult: + """Resolve a check whose hold was released because nothing names the event + the settling track event would carry. + + Fail-open means assume the credits are there, and the server has already + evaluated the flag and said yes; only the settle is impossible, so its + verdict stands. Fail-closed denies, as it does for any check it cannot gate. + """ + reason = "missing_event_subtype" + if options.on_acquire_failure == "fail-closed": + return CheckResult(allowed=False, value=False, reason=reason, flag_key=result.flag_key, error=reason) + result.reservation = None + result.error = reason + return result @dataclass @@ -75,6 +554,10 @@ class IdentifyOptions: # Client-supplied dedupe key. Duplicate events with the same key # (scoped to the environment) are dropped server-side for 24 hours. idempotency_key: Optional[str] = None + # Credit type IDs to warm leases for once the identify is enqueued, so the + # session's first check() does not pay the acquire round trip. Honored by + # AsyncSchematic in client mode; ignored everywhere else. + prewarm: Optional[List[str]] = None def _event_options_to_kwargs( @@ -128,6 +611,7 @@ class SchematicConfig: offline: bool = False timeout: Optional[float] = None cache_providers: Optional[List[CacheProvider[CheckFlagResponseData]]] = None + credit_leases: Optional[CreditLeaseConfig] = None class Schematic(BaseSchematic): @@ -161,6 +645,12 @@ def __init__(self, api_key: str, config: Optional[SchematicConfig] = None): else [LocalCache[CheckFlagResponseData](DEFAULT_CACHE_SIZE, DEFAULT_CACHE_TTL)] ) self.offline = config.offline + self._credit_leases = config.credit_leases + self._reservation_ttl = _resolve_reservation_ttl(self.logger, config.credit_leases) + if config.credit_leases is not None: + _warn_credit_lease_config( + self.logger, config.credit_leases, self.offline, supports_client_mode=False, + ) atexit.register(self.shutdown) @@ -309,23 +799,171 @@ def _check_flag_via_api( options: Optional[CheckFlagOptions] = None, ) -> CheckFlagResponseData: try: + preflight = _build_preflight(options) cache_key = _build_cache_key(flag_key, company, user) - cached_value = self._safe_cache_get(cache_key) - if cached_value is not None: - return cached_value - - resp = self.features.check_flag(flag_key, company=company, user=user) + # The cache is keyed by flag, company and user, and a preflighted + # check asks a different question ("would this action be allowed?") + # than the plain one, so it can neither be answered from the cache + # nor written to it. + if preflight is None: + cached_value = self._safe_cache_get(cache_key) + if cached_value is not None: + return cached_value + + preflight_kwargs: Dict[str, Any] = {} if preflight is None else {"preflight": preflight} + resp = self.features.check_flag(flag_key, company=company, user=user, **preflight_kwargs) if resp is None or resp.data is None or resp.data.value is None: return self._default_response(flag_key, options, REASON_FLAG_NOT_FOUND) - self._safe_cache_set(cache_key, resp.data) + if preflight is None: + self._safe_cache_set(cache_key, resp.data) return resp.data except Exception as e: self.logger.error(e) return self._default_response(flag_key, options, f"{REASON_ERROR}: {e}") + def _effective_lease_mode(self) -> Optional[Literal["client", "server"]]: + """Which mode a check with usage resolves to on this client. + + None means no credit gating at all. Client mode rides on DataStream, + which this SDK offers on AsyncSchematic alone, so it resolves to + nothing here and the check stays plain. + """ + if self._credit_leases is None or self.offline: + return None + if self._credit_leases.mode == "client": + return None + return "server" + + def check( + self, + flag_key: str, + company: Optional[Dict[str, str]] = None, + user: Optional[Dict[str, str]] = None, + options: Optional[CheckOptions] = None, + ) -> CheckResult: + """Credit-aware feature check. + + When ``credit_leases`` is configured and the caller passes ``usage``, + this gates the check on the company's credit balance and returns a + reservation handle on success; pass that handle to + ``track_with_reservation`` when the work completes. + + Without either it falls through to a plain flag check and returns + ``allowed = value`` with no reservation. The caller's preflight + (``usage`` / ``event_subtype``) is still threaded through that plain + check, so the verdict accounts for the usage about to be recorded, just + without holding anything. + """ + mode = self._effective_lease_mode() + if options is None or options.usage is None or mode is None: + return self._check_fallback(flag_key, company, user, options) + if mode == "server": + return self._check_with_server_reservation(flag_key, company, user, options) + return self._check_fallback(flag_key, company, user, options) + + def _check_fallback( + self, + flag_key: str, + company: Optional[Dict[str, str]], + user: Optional[Dict[str, str]], + options: Optional[CheckOptions], + ) -> CheckResult: + resp = self.check_flag_with_entitlement( + flag_key, company=company, user=user, options=_check_options_to_flag_options(options), + ) + return CheckResult( + allowed=resp.value, + value=resp.value, + reason=resp.reason, + flag_key=resp.flag or flag_key, + entitlement=resp.entitlement, + flag_id=resp.flag_id, + error=resp.error, + ) + + def _check_with_server_reservation( + self, + flag_key: str, + company: Optional[Dict[str, str]], + user: Optional[Dict[str, str]], + options: CheckOptions, + ) -> CheckResult: + """Gate one check on the server: a single check-and-reserve call + evaluates the flag against the company's real balance and takes the + hold in the same round trip. + + ``fail-open`` here returns the caller's default rather than re-running + the rules with an assumed-sufficient balance, since the call that would + have answered is the one that failed and there is no local engine to + fall back on. No flag_check event is enqueued: the server logs the + check, the same way the plain REST path does. + """ + def failure(reason: str) -> CheckResult: + return _server_failure_result( + flag_key, options, reason, self._resolve_default(flag_key, _check_options_to_flag_options(options)), + ) + + if not _is_valid_quantity(options.usage): + self.logger.error( + f"Server reservation: invalid usage {options.usage!r} for flag {flag_key}; " + "must be a finite, non-negative number" + ) + return failure("invalid_usage") + + if options.usage == 0: + self.logger.debug( + f"Server reservation: usage is 0 for flag {flag_key}, nothing to hold, using a plain check" + ) + return self._check_fallback(flag_key, company, user, options) + + try: + resp = self.features.check_and_reserve_flag( + flag_key, + company=company, + user=user, + quantity=options.usage, + expires_at=dt.datetime.now(dt.timezone.utc) + dt.timedelta(seconds=self._reservation_ttl), + **_reservation_request_kwargs(options), + ) + data = resp.data + except Exception as e: + if _is_payment_required(e): + return _payment_required_result(flag_key, e) + self.logger.error(f"Server reservation: check-and-reserve for flag {flag_key} failed: {e}") + return failure("server_reservation_failed") + + result = _reservation_check_result(flag_key, data) + + # No hold comes back when the flag denied, the credits were short, or + # the feature is not credit-metered. Nothing to release either way. + held = data.reservation + if not data.value or held is None: + return result + + # The settling track event is named by the event subtype; the caller's + # wins, otherwise the server names it on the hold. With neither, the + # hold could never be settled, so release it now instead of parking + # the credits until the TTL. + event_subtype = options.event_subtype or held.event_subtype + if not event_subtype: + self.logger.error( + f"Server reservation: reservation {held.id} for flag {flag_key} names no event subtype; " + "releasing it, since it could never be settled" + ) + try: + self.credits.release_credit_reservation(held.id) + except Exception as e: + self.logger.warning( + f"Server reservation: failed to release {held.id} ({e}); its hold is refunded when it expires" + ) + return _missing_event_subtype_result(options, result) + + result.reservation = _server_reservation(held, event_subtype, company, user) + return result + def identify( self, keys: Dict[str, str], @@ -366,6 +1004,46 @@ def track( options=options, ) + def track_with_reservation( + self, + reservation: Optional[Reservation], + actual_quantity: float, + options: Optional[TrackWithReservationOptions] = None, + ) -> None: + """Settle a reservation issued by ``check`` with the actual usage. + + The track event carries the reservation ID, and the server settles the + hold, refunding the unspent slice, when it processes the event. The + event's idempotency key is derived from the reservation ID, so a + duplicate or retried settle is dropped server-side rather than billed + twice. + """ + if self.offline: + return + # A check can allow without taking a hold, so a caller that settles + # whatever check() handed back can land here with nothing to settle. + # The usage still has to be recorded, but only a plain track() can. + if reservation is None: + self.logger.error( + "track_with_reservation: no reservation to settle; the check allowed without taking a hold. " + "Report the usage with track() instead" + ) + return + # A quantity the server cannot bill must reach neither the event nor + # the hold: skip the settle and let the hold refund itself at its TTL. + if not _is_valid_quantity(actual_quantity): + self.logger.error( + f"track_with_reservation: invalid actual_quantity {actual_quantity!r} for reservation " + f"{reservation.id}; must be a finite, non-negative number. Skipping the settle, the hold is " + "refunded at its TTL" + ) + return + self._enqueue_event( + "track", + _build_reservation_track_event(reservation, _settled_quantity(actual_quantity), options), + options=TrackOptions(idempotency_key=f"{RESERVATION_TRACK_IDEMPOTENCY_PREFIX}{reservation.id}"), + ) + def _enqueue_event( self, event_type: str, @@ -418,6 +1096,7 @@ class AsyncSchematicConfig: cache_providers: Optional[List[CacheProvider[CheckFlagResponseData]]] = None use_datastream: bool = False datastream: Optional[DataStreamConfig] = None + credit_leases: Optional[CreditLeaseConfig] = None class AsyncSchematic(AsyncBaseSchematic): @@ -482,6 +1161,16 @@ def __init__(self, api_key: str, config: Optional[AsyncSchematicConfig] = None): self.offline = config.offline self._shutdown_requested = False self._is_shutting_down = False + self._credit_leases = config.credit_leases + self._reservation_ttl = _resolve_reservation_ttl(self.logger, config.credit_leases) + # Client-mode plumbing, built below once DataStream is wired so that + # "auto" can resolve against it. Server mode builds none of it. + self._lease_store: Optional[LeaseStore] = None + self._reservations: Optional[ReservationStore] = None + self._lease_manager: Optional[LeaseManager] = None + self._lease_backend_shared = False + self._prewarm_resolve_timeout = DEFAULT_PREWARM_RESOLVE_TIMEOUT + self._background_tasks: set = set() # DataStream client self._datastream_client: Optional[DataStreamClient] = None @@ -514,6 +1203,19 @@ def __init__(self, api_key: str, config: Optional[AsyncSchematicConfig] = None): self._datastream_client = DataStreamClient(ds_opts) + if config.credit_leases is not None: + _warn_credit_lease_config( + self.logger, + config.credit_leases, + self.offline, + supports_client_mode=True, + datastream_enabled=self._datastream_client is not None, + ) + if not self.offline and _mode_uses_leases( + config.credit_leases.mode, self._datastream_client is not None + ): + self._build_lease_plumbing(config.credit_leases, config.datastream) + self._initialized = True async def __aenter__(self): @@ -523,6 +1225,61 @@ async def __aenter__(self): async def __aexit__(self, exc_type, exc_val, exc_tb): await self.shutdown() + def _build_lease_plumbing( + self, credit_leases: CreditLeaseConfig, datastream: Optional[DataStreamConfig], + ) -> None: + """Build the lease store, the reservation table, and their manager. + + Lease state belongs in a shared cache so gating holds across pods. An + explicit redis_client wins; otherwise the DataStream company cache's + Redis is reused, so an existing setup backs leases with no second + client to wire up. + """ + redis_client = credit_leases.redis_client + key_prefix = credit_leases.redis_key_prefix + company_cache = datastream.company_cache if datastream is not None else None + if redis_client is None and isinstance(company_cache, RedisCache): + redis_client = company_cache.client + if key_prefix is None: + # RedisCache joins its prefix to a key with a colon, so lease + # keys land in the same namespace as the cached entities. + key_prefix = f"{company_cache.prefix}:" + self.logger.debug( + "credit_leases: reusing the DataStream cache's Redis client for lease and reservation state" + ) + prefix = key_prefix or DEFAULT_LEASE_KEY_PREFIX + if redis_client is not None: + self._lease_backend_shared = True + self._lease_store = RedisLeaseStore( + redis_client, + key_prefix=prefix, + default_lease_duration=credit_leases.default_lease_duration or DEFAULT_LEASE_DURATION, + ) + self._reservations = RedisReservationStore(redis_client, self._lease_store, key_prefix=prefix) + else: + self.logger.warning( + "credit_leases is enabled without a shared Redis backend, so lease and reservation state stays " + "per-process and gating holds within this process only. Set credit_leases.redis_client (or give " + "datastream.company_cache a RedisCache) so leases gate across every SDK instance." + ) + self._lease_store = InMemoryLeaseStore() + self._reservations = InMemoryReservationStore(self._lease_store) + self._lease_manager = LeaseManager( + CreditsWireClient(self.credits), + self._lease_store, + reservation_store=self._reservations, + config=LeaseConfig( + lease_duration=credit_leases.default_lease_duration, + reservation_ttl=credit_leases.default_reservation_ttl, + lease_size=credit_leases.default_lease_size, + low_water_mark=credit_leases.low_water_mark, + sweep_interval=credit_leases.sweep_interval, + overrides=credit_leases.overrides or {}, + ), + ) + if credit_leases.prewarm_resolve_timeout is not None: + self._prewarm_resolve_timeout = credit_leases.prewarm_resolve_timeout + async def _start_datastream(self) -> None: if self._datastream_client is not None: try: @@ -530,6 +1287,11 @@ async def _start_datastream(self) -> None: except Exception as e: self.logger.error(f"Failed to start DataStream client: {e}") self._datastream_client = None + return + # The sweeper needs a running loop, and it has nothing to sweep + # until checks can reserve, which is once DataStream is up. + if self._lease_manager is not None: + self._lease_manager.start_sweep() async def initialize(self) -> None: await self._start_datastream() @@ -564,6 +1326,7 @@ async def check_flag_with_entitlement( resp = await ds.check_flag( CheckFlagRequestBody(company=company, user=user), flag_key, + options=options, ) await self._enqueue_flag_check_event(flag_key, resp, company, user) return self._ds_result_to_response(flag_key, resp, options) @@ -750,23 +1513,288 @@ async def _check_flag_via_api( options: Optional[CheckFlagOptions] = None, ) -> CheckFlagResponseData: try: + preflight = _build_preflight(options) cache_key = _build_cache_key(flag_key, company, user) - cached_value = self._safe_cache_get(cache_key) - if cached_value is not None: - return cached_value - - resp = await self.features.check_flag(flag_key, company=company, user=user) + # The cache is keyed by flag, company and user, and a preflighted + # check asks a different question ("would this action be allowed?") + # than the plain one, so it can neither be answered from the cache + # nor written to it. + if preflight is None: + cached_value = self._safe_cache_get(cache_key) + if cached_value is not None: + return cached_value + + preflight_kwargs: Dict[str, Any] = {} if preflight is None else {"preflight": preflight} + resp = await self.features.check_flag(flag_key, company=company, user=user, **preflight_kwargs) if resp is None or resp.data is None or resp.data.value is None: return self._default_response(flag_key, options, REASON_FLAG_NOT_FOUND) - self._safe_cache_set(cache_key, resp.data) + if preflight is None: + self._safe_cache_set(cache_key, resp.data) return resp.data except Exception as e: self.logger.error(e) return self._default_response(flag_key, options, f"{REASON_ERROR}: {e}") + def _effective_lease_mode(self) -> Optional[Literal["client", "server"]]: + """Which mode a check with usage resolves to right now. + + None means no credit gating at all. "auto" resolves per check rather + than once at construction, so a DataStream whose start() failed, which + clears the client, falls to server mode instead of leaving every check + ungated. + """ + if self._credit_leases is None or self.offline: + return None + mode = self._credit_leases.mode + if mode == "server": + return "server" + plumbing_ready = ( + self._lease_manager is not None and self._lease_store is not None and self._reservations is not None + ) + if mode == "client": + return "client" if plumbing_ready else None + return "client" if self._datastream_client is not None and plumbing_ready else "server" + + async def check( + self, + flag_key: str, + company: Optional[Dict[str, str]] = None, + user: Optional[Dict[str, str]] = None, + options: Optional[CheckOptions] = None, + ) -> CheckResult: + """Credit-aware feature check. + + When ``credit_leases`` is configured and the caller passes ``usage``, + this gates the check on the company's credit balance and returns a + reservation handle on success; pass that handle to + ``track_with_reservation`` when the work completes. + + Without either it falls through to a plain flag check and returns + ``allowed = value`` with no reservation. The caller's preflight + (``usage`` / ``event_subtype``) is still threaded through that plain + check, so the verdict accounts for the usage about to be recorded, just + without holding anything. + """ + mode = self._effective_lease_mode() + if options is None or options.usage is None or mode is None: + return await self._check_fallback(flag_key, company, user, options) + if mode == "server": + return await self._check_with_server_reservation(flag_key, company, user, options) + return await self._check_with_lease(flag_key, company, user, options) + + async def _check_with_lease( + self, + flag_key: str, + company: Optional[Dict[str, str]], + user: Optional[Dict[str, str]], + options: CheckOptions, + ) -> CheckResult: + """Gate one check on a hold carved out of a local credit lease.""" + lease_store, reservations, manager = self._lease_store, self._reservations, self._lease_manager + if lease_store is None or reservations is None or manager is None: + return await self._check_fallback(flag_key, company, user, options) + + async def fallback() -> CheckResult: + return await self._check_fallback(flag_key, company, user, options) + + async def enqueue_flag_check(body: EventBodyFlagCheck) -> None: + await self._enqueue_event("flag_check", body) + + return await check_with_lease( + CreditCheckDeps( + datastream=self._datastream_client, + lease_store=lease_store, + reservations=reservations, + manager=manager, + logger=self.logger, + enqueue_flag_check=enqueue_flag_check, + ), + flag_key, + company, + user, + options, + fallback, + ) + + async def prewarm(self, company: Dict[str, str], credit_type_ids: List[str]) -> None: + """Acquire a lease per credit type up front, so a session's first + check() does not pay the acquire round trip. + + Best effort: failures are logged, never raised. When the company keys + carry no id, this fetches the company over DataStream, waiting up to + ``credit_leases.prewarm_resolve_timeout`` for it to surface, which + covers a company the server has only just ingested. + """ + if self._lease_manager is None: + self.logger.debug( + "prewarm is a no-op in server mode; there is no local lease to warm" + if self._effective_lease_mode() == "server" + else "prewarm called but client-mode credit leases are not configured" + ) + return + if not company: + self.logger.debug("prewarm needs company keys") + return + company_id = await self._resolve_company_id_with_wait(company) + if not company_id: + self.logger.debug( + f"prewarm: company {company} did not resolve within {self._prewarm_resolve_timeout}s; " + "the first check() acquires instead" + ) + return + await asyncio.gather(*(self._prewarm_one(company_id, credit_type_id) for credit_type_id in credit_type_ids)) + + async def _prewarm_one(self, company_id: str, credit_type_id: str) -> None: + manager = self._lease_manager + if manager is None: + return + try: + await manager.acquire_if_needed(company_id, credit_type_id) + except Exception as e: + self.logger.warning(f"prewarm: failed to acquire a lease for {credit_type_id}: {e}") + + async def _resolve_company_id_with_wait(self, company: Dict[str, str]) -> Optional[str]: + """Resolve company keys to an ID, waiting for the company to surface. + + identify does not push a company into the DataStream cache, since + companies are only streamed on request, so this fetches (cache first, + then over the socket) rather than watching an empty cache. The fetch + also primes the cache, so the first real check() takes the lease path. + A prewarm_resolve_timeout of 0 keeps the cache lookup and skips the + wait, so an already-seen company still warms. + """ + company_id = company.get("id") + if company_id: + return company_id + datastream = self._datastream_client + if datastream is None: + return None + # An earlier check or prewarm may already have cached this company, and + # that answer costs nothing. + try: + cached = await datastream.get_cached_company(company) + if cached is not None and cached.id: + return cached.id + except Exception as e: + self.logger.debug(f"prewarm: DataStream company cache lookup failed ({e})") + if self._prewarm_resolve_timeout <= 0: + return None + deadline = time.monotonic() + self._prewarm_resolve_timeout + while True: + try: + resolved = await datastream.get_company(company) + if resolved is not None and resolved.id: + return resolved.id + except Exception as e: + # Expected while the socket is still connecting, and while the + # server has yet to ingest a preceding identify. + self.logger.debug(f"prewarm: DataStream company fetch failed ({e})") + if time.monotonic() >= deadline: + return None + await asyncio.sleep(PREWARM_POLL_INTERVAL) + + async def _check_fallback( + self, + flag_key: str, + company: Optional[Dict[str, str]], + user: Optional[Dict[str, str]], + options: Optional[CheckOptions], + ) -> CheckResult: + resp = await self.check_flag_with_entitlement( + flag_key, company=company, user=user, options=_check_options_to_flag_options(options), + ) + return CheckResult( + allowed=resp.value, + value=resp.value, + reason=resp.reason, + flag_key=resp.flag or flag_key, + entitlement=resp.entitlement, + flag_id=resp.flag_id, + error=resp.error, + ) + + async def _check_with_server_reservation( + self, + flag_key: str, + company: Optional[Dict[str, str]], + user: Optional[Dict[str, str]], + options: CheckOptions, + ) -> CheckResult: + """Gate one check on the server: a single check-and-reserve call + evaluates the flag against the company's real balance and takes the + hold in the same round trip. + + ``fail-open`` here returns the caller's default rather than re-running + the rules with an assumed-sufficient balance, since the call that would + have answered is the one that failed. No flag_check event is enqueued: + the server logs the check, the same way the plain REST path does. + """ + def failure(reason: str) -> CheckResult: + return _server_failure_result( + flag_key, options, reason, self._resolve_default(flag_key, _check_options_to_flag_options(options)), + ) + + if not _is_valid_quantity(options.usage): + self.logger.error( + f"Server reservation: invalid usage {options.usage!r} for flag {flag_key}; " + "must be a finite, non-negative number" + ) + return failure("invalid_usage") + + if options.usage == 0: + self.logger.debug( + f"Server reservation: usage is 0 for flag {flag_key}, nothing to hold, using a plain check" + ) + return await self._check_fallback(flag_key, company, user, options) + + try: + resp = await self.features.check_and_reserve_flag( + flag_key, + company=company, + user=user, + quantity=options.usage, + expires_at=dt.datetime.now(dt.timezone.utc) + dt.timedelta(seconds=self._reservation_ttl), + **_reservation_request_kwargs(options), + ) + data = resp.data + except Exception as e: + if _is_payment_required(e): + return _payment_required_result(flag_key, e) + self.logger.error(f"Server reservation: check-and-reserve for flag {flag_key} failed: {e}") + return failure("server_reservation_failed") + + result = _reservation_check_result(flag_key, data) + + # No hold comes back when the flag denied, the credits were short, or + # the feature is not credit-metered. Nothing to release either way. + held = data.reservation + if not data.value or held is None: + return result + + # The settling track event is named by the event subtype; the caller's + # wins, otherwise the server names it on the hold. With neither, the + # hold could never be settled, so release it now instead of parking + # the credits until the TTL. + event_subtype = options.event_subtype or held.event_subtype + if not event_subtype: + self.logger.error( + f"Server reservation: reservation {held.id} for flag {flag_key} names no event subtype; " + "releasing it, since it could never be settled" + ) + try: + await self.credits.release_credit_reservation(held.id) + except Exception as e: + self.logger.warning( + f"Server reservation: failed to release {held.id} ({e}); its hold is refunded when it expires" + ) + return _missing_event_subtype_result(options, result) + + result.reservation = _server_reservation(held, event_subtype, company, user) + return result + async def identify( self, keys: Dict[str, str], @@ -785,6 +1813,37 @@ async def identify( ), options=options, ) + if options is not None and options.prewarm: + company_keys = company.keys if company is not None else None + if company_keys: + # Push the identify out before warming. Left in the buffer it + # would wait a whole flush period, and the prewarm's company + # resolution polls a server that has not seen the company yet. + try: + await self.event_buffer.flush() + except Exception as e: + self.logger.debug(f"identify: flushing before prewarm failed: {e}") + self._spawn_prewarm(company_keys, options.prewarm) + else: + self.logger.debug("identify: prewarm needs company keys on the identify event") + + def _spawn_prewarm(self, company: Dict[str, str], credit_type_ids: List[str]) -> None: + """Warm the leases behind an identify without making the caller wait.""" + + async def run() -> None: + try: + await self.prewarm(company, credit_type_ids) + except Exception as e: + self.logger.warning(f"identify prewarm failed: {e}") + + try: + task = asyncio.ensure_future(run()) + except RuntimeError: + self.logger.debug("identify: no running event loop, skipping prewarm") + return + # Held so the loop does not collect the task mid-flight. + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) async def track( self, @@ -808,6 +1867,11 @@ async def track( ) # Update company metrics in DataStream if available and connected + await self._update_company_metrics(company, event, quantity) + + async def _update_company_metrics( + self, company: Optional[Dict[str, str]], event: str, quantity: Optional[int], + ) -> None: ds = self._get_datastream() if company and ds is not None and ds.is_connected(): try: @@ -819,6 +1883,97 @@ async def track( except Exception as e: self.logger.error(f"Failed to update company metrics: {e}") + async def track_with_reservation( + self, + reservation: Optional[Reservation], + actual_quantity: float, + options: Optional[TrackWithReservationOptions] = None, + ) -> None: + """Settle a reservation issued by ``check`` with the actual usage. + + A server-mode hold settles by id: the track event carries it and the + server refunds the unspent slice when it processes the event. A + client-mode hold is consumed against its local lease first, and the + event carries the lease id so the server bills through the lease's + sub-ledger rather than decrementing the pre-debited grant again. + + Either way the event's idempotency key is derived from the reservation + ID, so a duplicate or retried settle is dropped server-side rather than + billed twice. + """ + if self.offline: + return + # A check can allow without taking a hold, so a caller that settles + # whatever check() handed back can land here with nothing to settle. + # The usage still has to be recorded, but only a plain track() can. + if reservation is None: + self.logger.error( + "track_with_reservation: no reservation to settle; the check allowed without taking a hold. " + "Report the usage with track() instead" + ) + return + # A quantity the server cannot bill must reach neither the event nor + # the hold: skip the settle and let the hold refund itself at its TTL. + if not _is_valid_quantity(actual_quantity): + self.logger.error( + f"track_with_reservation: invalid actual_quantity {actual_quantity!r} for reservation " + f"{reservation.id}; must be a finite, non-negative number. Skipping the settle, the hold is " + "refunded at its TTL" + ) + return + quantity = _settled_quantity(actual_quantity) + if reservation.mode == "server": + event = _build_reservation_track_event(reservation, quantity, options) + else: + event = await self._settle_client_reservation(reservation, actual_quantity, quantity, options) + await self._enqueue_event( + "track", + event, + options=TrackOptions(idempotency_key=f"{RESERVATION_TRACK_IDEMPOTENCY_PREFIX}{reservation.id}"), + ) + # The settled usage counts toward the company's metrics like any other + # track event, so a locally cached company stays consistent with it. + await self._update_company_metrics(reservation.company, reservation.event_subtype, quantity) + + async def _settle_client_reservation( + self, + reservation: Reservation, + actual_quantity: float, + quantity: int, + options: Optional[TrackWithReservationOptions], + ) -> EventBodyTrack: + """Consume a client-mode hold locally and hand back the event that bills it. + + The server is the source of truth for real consumption, so a settle + that cannot run locally still emits: the event's idempotency key keeps + the retry from billing twice. + """ + if self._reservations is None: + # The handle came from a lease-configured client, so the event + # still needs its lease id and dedupe key even though this client + # holds nothing to settle. + self.logger.warning( + "track_with_reservation: client-mode credit leases are not configured here, " + "emitting an unsettled track" + ) + return _build_reservation_track_event(reservation, quantity, options) + try: + outcome = await consume_reservation_and_build_event( + self._reservations, reservation, actual_quantity, options, + ) + except Exception as e: + self.logger.warning( + f"track_with_reservation: failed to settle reservation {reservation.id} locally ({e}), " + "emitting the track anyway" + ) + return _build_reservation_track_event(reservation, quantity, options) + if not outcome.settled_locally: + self.logger.debug( + f"track_with_reservation: reservation {reservation.id} was not settled locally (swept at its " + "TTL, already settled, or the store is unreachable); the track is keyed for server-side dedupe" + ) + return outcome.track + async def _enqueue_event( self, event_type: str, @@ -878,6 +2033,15 @@ async def shutdown(self) -> None: self.logger.info("Shutting down AsyncSchematic...") try: + if self._lease_manager is not None: + self._lease_manager.stop() + if not self._lease_backend_shared: + # Per-process leases have no sibling drawing on them, so + # releasing hands the unspent remainder back to the company + # balance now instead of at expiry. A shared lease must + # survive this process's shutdown, or the release pulls the + # grant out from under the pods still drawing on it. + await self._lease_manager.release_all_local_leases() if self._datastream_client is not None: try: await self._datastream_client.close() diff --git a/src/schematic/datastream/datastream_client.py b/src/schematic/datastream/datastream_client.py index 90faa8e3..2578cb96 100644 --- a/src/schematic/datastream/datastream_client.py +++ b/src/schematic/datastream/datastream_client.py @@ -1,23 +1,35 @@ from __future__ import annotations import asyncio -import httpx import logging import typing - from dataclasses import dataclass from typing import Any, Callable, Dict, List, Optional +import httpx +from ..cache import AsyncCacheProvider, AsyncLocalCache from ..types.check_flag_request_body import CheckFlagRequestBody from ..types.rulesengine_check_flag_result import RulesengineCheckFlagResult from ..types.rulesengine_company import RulesengineCompany from ..types.rulesengine_flag import RulesengineFlag from ..types.rulesengine_user import RulesengineUser -from ..cache import AsyncCacheProvider, AsyncLocalCache from .merge import partial_company, partial_user from .rules_engine import RulesEngineClient -from .types import DataStreamBaseReq, DataStreamReq, DataStreamResp, EntityType, KeyConflictError, MessageType, RulesEngineError -from .websocket_client import MAX_MESSAGE_SIZE, ClientOptions as WSClientOptions, DatastreamWSClient +from .types import ( + DataStreamBaseReq, + DataStreamReq, + DataStreamResp, + EntityType, + KeyConflictError, + MessageType, + RulesEngineError, +) +from .websocket_client import MAX_MESSAGE_SIZE, DatastreamWSClient +from .websocket_client import ClientOptions as WSClientOptions + +if typing.TYPE_CHECKING: + # Imported for typing only: the client module imports this package. + from ..client import CheckFlagOptions _hints_cache: Dict[type, Dict[str, Any]] = {} @@ -384,8 +396,13 @@ async def check_flag( self, eval_ctx: CheckFlagRequestBody, flag_key: str, + options: Optional["CheckFlagOptions"] = None, ) -> RulesengineCheckFlagResult: - """Evaluate a flag for a company and/or user context.""" + """Evaluate a flag for a company and/or user context. + + ``options`` carries the caller's preflight (hypothetical usage) into + the local evaluation. + """ flag = await self.get_flag(flag_key) if flag is None: raise RuntimeError(f"Flag not found: {flag_key}") @@ -415,11 +432,11 @@ async def check_flag( # Replicator mode — evaluate with whatever is cached if self._replicator_mode: - return self._evaluate_flag(flag, cached_company, cached_user) + return self._evaluate_flag(flag, cached_company, cached_user, options) # If we have all required entities cached, evaluate immediately if (not needs_company or cached_company) and (not needs_user or cached_user): - return self._evaluate_flag(flag, cached_company, cached_user) + return self._evaluate_flag(flag, cached_company, cached_user, options) if not self.is_connected(): raise RuntimeError("Datastream not connected and required entities not in cache") @@ -437,7 +454,31 @@ async def check_flag( tasks.append(_resolved(cached_user)) results: list = await asyncio.gather(*tasks) - return self._evaluate_flag(flag, results[0], results[1]) + return self._evaluate_flag(flag, results[0], results[1], options) + + def evaluate_flag( + self, + flag: RulesengineFlag, + company: Optional[RulesengineCompany], + user: Optional[RulesengineUser], + options: Optional["CheckFlagOptions"] = None, + ) -> RulesengineCheckFlagResult: + """Evaluate a flag against entities the caller already holds. + + ``check_flag`` resolves its entities from the caches first; this runs + the engine on the snapshots it is handed, which is what lets the credit + lease path gate against a substituted balance. + """ + return self._evaluate_flag(flag, company, user, options) + + async def get_cached_company(self, keys: Dict[str, str]) -> Optional[RulesengineCompany]: + """The cached company for these keys, without asking the server. + + ``get_company`` falls through to a socket round trip on a miss; this + answers only from what is already local, for callers that would rather + move on than wait. + """ + return await self._get_company_from_cache(keys) async def update_company_metrics(self, keys: Dict[str, str], event: str, quantity: int) -> None: """Update company metrics locally in cache (for track events).""" @@ -1007,6 +1048,7 @@ def _evaluate_flag( flag: RulesengineFlag, company: Optional[RulesengineCompany], user: Optional[RulesengineUser], + options: Optional["CheckFlagOptions"] = None, ) -> RulesengineCheckFlagResult: """Evaluate a flag with the local rules engine. @@ -1019,7 +1061,7 @@ def _evaluate_flag( raise RulesEngineError(f"Rules engine not initialized (flag {flag.key})") try: - return self._rules_engine.check_flag(flag, company, user) + return self._rules_engine.check_flag(flag, company, user, options) except Exception as exc: self._logger.warning("Rules engine evaluation failed for flag %s: %s", flag.key, exc) raise RulesEngineError(f"Rules engine evaluation failed for flag {flag.key}: {exc}") from exc diff --git a/src/schematic/datastream/rules_engine.py b/src/schematic/datastream/rules_engine.py index 5759b714..6085f013 100644 --- a/src/schematic/datastream/rules_engine.py +++ b/src/schematic/datastream/rules_engine.py @@ -5,13 +5,17 @@ import re import time from pathlib import Path -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Dict, Optional from ..types.rulesengine_check_flag_result import RulesengineCheckFlagResult from ..types.rulesengine_company import RulesengineCompany from ..types.rulesengine_flag import RulesengineFlag from ..types.rulesengine_user import RulesengineUser +if TYPE_CHECKING: + # Imported for typing only: the client module imports this package. + from ..client import CheckFlagOptions + logger = logging.getLogger(__name__) _CAMEL_RE = re.compile(r"([A-Z])") @@ -49,6 +53,23 @@ def _strip_none(obj: Any) -> Any: return obj +def _engine_options(options: "CheckFlagOptions") -> Dict[str, Any]: + """Build the engine's preflight options block, in the snake_case shape its + serde struct expects, with unset fields dropped.""" + event_usage = options.event_usage + return _strip_none( + { + "credit_cost": options.credit_cost, + "event_usage": ( + {"event_subtype": event_usage.event_subtype, "quantity": event_usage.quantity} + if event_usage is not None + else None + ), + "usage": options.usage, + } + ) + + # Path to the WASM binary shipped alongside this module _WASM_PATH = Path(__file__).parent / "wasm" / "rulesengine.wasm" @@ -149,20 +170,29 @@ def check_flag( flag: RulesengineFlag, company: Optional[RulesengineCompany] = None, user: Optional[RulesengineUser] = None, + options: Optional["CheckFlagOptions"] = None, ) -> RulesengineCheckFlagResult: """Evaluate a flag using the WASM rules engine. Accepts Fern-generated Pydantic models (or plain dicts). Serialises them into a single JSON envelope, passes it to the WASM module, and returns a ``RulesengineCheckFlagResult``. + + ``options`` carries the caller's preflight (hypothetical usage) into + the evaluation. The key is left off the envelope entirely when there + is no preflight, so envelopes for plain checks are unchanged. """ self._ensure_initialized() - envelope = { + envelope: dict[str, Any] = { "flag": _strip_none(flag.model_dump(exclude_none=True, mode="json")), "company": _strip_none(company.model_dump(exclude_none=True, mode="json")) if company else None, "user": _strip_none(user.model_dump(exclude_none=True, mode="json")) if user else None, } + if options is not None: + engine_options = _engine_options(options) + if engine_options: + envelope["options"] = engine_options result_json = self._call_wasm(json.dumps(envelope)) result_data = _deep_camel_to_snake(json.loads(result_json)) diff --git a/src/schematic/event_buffer.py b/src/schematic/event_buffer.py index ebcd5d7c..49f3c4c2 100644 --- a/src/schematic/event_buffer.py +++ b/src/schematic/event_buffer.py @@ -121,6 +121,11 @@ def push(self, event: CreateEventRequestBody): with self.lock: self.events.append(event) + def flush(self): + """Send everything buffered right now, rather than waiting for the + period to elapse.""" + self._flush() + def stop(self): try: self.stopped = True @@ -233,6 +238,11 @@ async def push(self, event: CreateEventRequestBody): async with self.lock: self.events.append(event) + async def flush(self): + """Send everything buffered right now, rather than waiting for the + period to elapse.""" + await self._flush() + async def stop(self): try: self.stopped = True diff --git a/src/schematic/leases/__init__.py b/src/schematic/leases/__init__.py new file mode 100644 index 00000000..6dcb0832 --- /dev/null +++ b/src/schematic/leases/__init__.py @@ -0,0 +1,73 @@ +"""Client-mode credit leases: local lease slots, reservations, and their manager. + +Async only. Client-mode leases ride on DataStream, which this SDK offers on +``AsyncSchematic`` alone, so there is no synchronous variant. + +The in-memory stores gate within one process; the Redis stores gate across +pods and share their key layout with the Node SDK, so mixed fleets agree on +every lease. +""" + +from .check import CheckDataStream, CreditCheckDeps, check_with_lease +from .lease_manager import CreditsWireClient, LeaseGrant, LeaseManager, LeaseWireClient +from .lease_store import InMemoryLeaseStore, LeaseStore, lease_key +from .redis_lease_store import RedisLeaseStore +from .redis_reservation_store import RedisReservationStore +from .reservation_store import InMemoryReservationStore, ReservationStore +from .track import ( + ReservationConsumeResult, + build_reservation_track_event, + consume_reservation_and_build_event, + settled_quantity, +) +from .types import ( + DEFAULT_LEASE_DURATION, + DEFAULT_LEASE_SIZE, + DEFAULT_LOW_WATER_MARK, + DEFAULT_PREWARM_RESOLVE_TIMEOUT, + DEFAULT_RESERVATION_TTL, + DEFAULT_SWEEP_INTERVAL, + Clock, + LeaseConfig, + LeaseConfigOverride, + LeaseState, + ReservationRecord, + ResolvedLeaseConfig, + is_valid_quantity, + resolve_lease_config, +) + +__all__ = [ + "CheckDataStream", + "Clock", + "CreditCheckDeps", + "CreditsWireClient", + "DEFAULT_LEASE_DURATION", + "DEFAULT_LEASE_SIZE", + "DEFAULT_LOW_WATER_MARK", + "DEFAULT_PREWARM_RESOLVE_TIMEOUT", + "DEFAULT_RESERVATION_TTL", + "DEFAULT_SWEEP_INTERVAL", + "InMemoryLeaseStore", + "InMemoryReservationStore", + "LeaseConfig", + "LeaseConfigOverride", + "LeaseGrant", + "LeaseManager", + "LeaseState", + "LeaseStore", + "LeaseWireClient", + "RedisLeaseStore", + "RedisReservationStore", + "ReservationConsumeResult", + "ReservationRecord", + "ReservationStore", + "ResolvedLeaseConfig", + "build_reservation_track_event", + "check_with_lease", + "consume_reservation_and_build_event", + "is_valid_quantity", + "lease_key", + "resolve_lease_config", + "settled_quantity", +] diff --git a/src/schematic/leases/check.py b/src/schematic/leases/check.py new file mode 100644 index 00000000..5793cfa1 --- /dev/null +++ b/src/schematic/leases/check.py @@ -0,0 +1,526 @@ +"""The client-mode check flow: probe the entitlement, lease, reserve, gate. + +One ``check()`` with usage runs the rules engine twice. The first run is a +probe against the company's real balance that names the credit being metered; +the second gates the call against the lease's local balance, after the credits +have already been debited. The conformance vectors in ``conformance/vectors`` +pin every step, and ``conformance/SPEC.md`` explains why each one is ordered +the way it is. +""" + +from __future__ import annotations + +import logging +import time +import uuid +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, Optional, Protocol + +from ..types import ( + EventBodyFlagCheck, + RulesengineCheckFlagResult, + RulesengineCompany, + RulesengineFlag, + RulesengineUser, +) +from .lease_manager import LeaseManager +from .lease_store import LeaseStore +from .reservation_store import ReservationStore +from .types import Clock, ReservationRecord, is_valid_quantity + +if TYPE_CHECKING: + from ..client import CheckFlagOptions, CheckOptions, CheckResult, Reservation + +logger = logging.getLogger(__name__) + +# The balance the fail-open evaluation substitutes for the metered credit: +# large enough that the credit gate always passes, and still exact as a JSON +# number, so the engine reads back what the SDK sent. This is Node's +# Number.MAX_SAFE_INTEGER, which the vectors name "max_safe_integer". +FAIL_OPEN_BALANCE = 2**53 - 1 + + +class CheckDataStream(Protocol): + """The slice of ``DataStreamClient`` a lease-bearing check touches. + + Narrow on purpose: it keeps this module off the DataStream client's wider + surface, and lets the conformance runner script the flow without a socket. + """ + + async def get_flag(self, flag_key: str) -> Optional[RulesengineFlag]: ... + + async def get_company(self, keys: Dict[str, str]) -> RulesengineCompany: ... + + async def get_user(self, keys: Dict[str, str]) -> Any: ... + + def evaluate_flag( + self, flag: Any, company: Any, user: Any, options: Any = None + ) -> RulesengineCheckFlagResult: ... + + +@dataclass +class CreditCheckDeps: + """Everything a lease-bearing check draws on, gathered by the client.""" + + datastream: Optional[CheckDataStream] + lease_store: LeaseStore + reservations: ReservationStore + manager: LeaseManager + logger: logging.Logger + # Report a flag_check event for a check this module resolved itself. The + # plain check paths enqueue one per check, so without this a lease-gated + # check would be invisible to flag-check analytics and company last-seen. + # Fallback exits do not call it: the plain check they delegate to reports + # its own. + enqueue_flag_check: Callable[[EventBodyFlagCheck], Awaitable[None]] + clock: Clock = time.time + + +async def check_with_lease( + deps: CreditCheckDeps, + flag_key: str, + company: Optional[Dict[str, str]], + user: Optional[Dict[str, str]], + options: "CheckOptions", + fallback: Callable[[], Awaitable["CheckResult"]], +) -> "CheckResult": + """Gate one check against a local lease, returning a hold when it allows. + + ``fallback`` is the plain flag check. Every step that cannot resolve a + credit to meter defers to it, since the plain check has its own degradation + story and issues no hold. A step that *can* resolve the credit but cannot + gate on it goes through ``on_acquire_failure`` instead. + """ + log = deps.logger + mode = options.on_acquire_failure or "fail-closed" + usage = options.usage + + # A malformed usage must never reach the stores, and the caller asked for a + # contract for exactly this case, so resolve it through that rather than + # letting it surface as an opaque reserve failure. + if usage is None or not is_valid_quantity(usage): + log.error( + f"Lease check: invalid usage {usage!r} for flag {flag_key}; must be a finite, non-negative number" + ) + return await _emit_flag_check( + deps, company, user, _static_failure_result(mode, flag_key, "invalid_usage", None) + ) + + # Nothing to reserve. The plain check still carries the preflight, so every + # rule evaluates normally; a 0-credit handle would only be a no-op. + if usage == 0: + log.debug(f"Lease check: usage is 0 for flag {flag_key}, nothing to reserve, using a plain check") + return await fallback() + + datastream = deps.datastream + if datastream is None: + log.debug("Lease check: no DataStream, using a plain check") + return await fallback() + + flag = await _load_flag(datastream, flag_key, log) + if flag is None: + log.debug(f"Lease check: no cached flag for {flag_key}, using a plain check") + return await fallback() + + if not company: + log.debug("Lease check: no company keys, using a plain check") + return await fallback() + + # Resolve company and user the way a plain DataStream check does: cache + # first, then a live fetch. Evaluating without an entity the caller named + # would silently skip its targeted rules and overrides, so a miss defers to + # the plain check instead. + resolved_company = await _load_company(datastream, company, log) + if resolved_company is None: + return await fallback() + + resolved_user: Optional[RulesengineUser] = None + if user: + resolved_user = await _load_user(datastream, user, log) + if resolved_user is None: + return await fallback() + + # Entitlement-first resolution. The probe runs against the real balance + # with no preflight: applying a credit cost to a lease-depleted server + # balance could fail the credit condition, drop the engine to a + # lower-priority rule, and hide the entitlement being looked for. + try: + probe = datastream.evaluate_flag(flag, resolved_company, resolved_user, None) + except Exception as err: + # A probe failure is a resolution miss, not the gate, and no hold + # exists yet to cancel. + log.warning(f"Lease check: entitlement probe failed for flag {flag_key} ({err}), using a plain check") + return await fallback() + + entitlement = probe.entitlement + if entitlement is None or entitlement.value_type != "credit": + value_type = entitlement.value_type if entitlement is not None else "" + log.debug( + f"Lease check: flag {flag_key} matched a non-credit entitlement (value_type={value_type}), " + "using a plain check, no reservation" + ) + return await fallback() + + credit_id = entitlement.credit_id + consumption_rate = entitlement.consumption_rate or 0.0 + # The caller's subtype wins; otherwise the entitlement names the metered + # event. A credit entitlement with neither a resolvable subtype nor a + # positive rate can never be billed, so it is not gateable. + event_subtype = options.event_subtype or entitlement.event_subtype + if not credit_id or consumption_rate <= 0 or not event_subtype: + log.debug( + f"Lease check: flag {flag_key} has an incomplete credit entitlement " + f"(credit_id={credit_id or ''}, consumption_rate={consumption_rate}, " + f"event_subtype={event_subtype or ''}), using a plain check" + ) + return await fallback() + + credit_cost = usage * consumption_rate + + async def failure(reason: str) -> "CheckResult": + result = await _handle_lease_failure( + datastream=datastream, + log=log, + mode=mode, + flag_key=flag_key, + reason=reason, + flag=flag, + company=resolved_company, + user=resolved_user, + credit_id=credit_id, + options=options, + ) + return await _emit_flag_check( + deps, + company, + user, + result, + company_id=resolved_company.id, + user_id=resolved_user.id if resolved_user is not None else None, + ) + + # The caller's per-check timeout governs the lease wire calls, the same way + # it governs the plain check's. + lease = await deps.manager.acquire_if_needed(resolved_company.id, credit_id, options.timeout) + if lease is None: + return await failure("lease_acquire_failed") + + # try_reserve is the atomic gate: check and debit in one step, returning + # the post-debit balance so the pre-debit figure needs no second read. + try: + post_reserve_balance = await deps.lease_store.try_reserve(resolved_company.id, credit_id, credit_cost) + if post_reserve_balance is None: + # Pass the cost as required_credits so a single large request + # extends even while the ratio sits above the water mark. + await deps.manager.maybe_extend(resolved_company.id, credit_id, credit_cost, options.timeout) + post_reserve_balance = await deps.lease_store.try_reserve(resolved_company.id, credit_id, credit_cost) + except Exception as err: + log.error(f"Lease check: reserve against {resolved_company.id}/{credit_id} failed: {err}") + return await failure("lease_store_error") + if post_reserve_balance is None: + return await failure("insufficient_lease_balance") + + # Record the hold after the debit and before the gate. A crash between the + # debit and this add leaks at most this one hold, reclaimed when the lease + # expires server-side; recording first would instead leave a record with no + # debit, which a later consume would refund into a double-spend. + resolved_config = deps.manager.resolve_config(credit_id) + record = ReservationRecord( + id=str(uuid.uuid4()), + lease_id=lease.lease_id, + company_id=resolved_company.id, + credit_type_id=credit_id, + event_subtype=event_subtype, + quantity_reserved=usage, + credits_reserved=credit_cost, + consumption_rate=consumption_rate, + expires_at=deps.clock() + resolved_config.reservation_ttl, + company=company, + user=user, + ) + try: + await deps.reservations.add(record) + except Exception as err: + log.error(f"Lease check: failed to persist reservation {record.id}: {err}") + # Undo the debit rather than strand it until lease expiry. consume + # claims whatever slice of the add landed and refunds it; a None says + # nothing landed, so refund the debit directly. Both are pinned to this + # lease. If the undo itself fails, accept the bounded leak: the slice + # comes back at lease expiry, which beats risking a double refund. + try: + if await deps.reservations.consume(record.id, 0) is None: + await deps.lease_store.refund(resolved_company.id, credit_id, credit_cost, lease.lease_id) + except Exception as undo_err: + log.warning( + f"Lease check: could not undo the local debit for {record.id} ({undo_err}); " + "the slice is reclaimed at lease expiry" + ) + return await failure("lease_store_error") + + # Gate against the lease's local view rather than the server's balance. The + # substituted figure is the PRE-reservation balance (what try_reserve + # returned plus what it debited, exact as of the debit), and credit_cost + # tells the engine what this call costs, so it evaluates the same + # arithmetic try_reserve just enforced, plus every non-credit rule. + pre_reservation = post_reserve_balance + credit_cost + substituted = _substitute_credit_balance(resolved_company, credit_id, pre_reservation) + try: + result = datastream.evaluate_flag( + flag, substituted, resolved_user, _credit_cost_options(credit_id, credit_cost) + ) + except Exception as err: + log.error(f"Lease check: rules evaluation failed for flag {flag_key}: {err}") + # The engine itself is down, so there is no fail-open re-evaluation to + # run: resolve the mode statically. + await _cancel_reservation(deps.reservations, record, log) + return await _emit_flag_check( + deps, + company, + user, + _static_failure_result(mode, flag_key, f"wasm_error: {err}", flag), + company_id=resolved_company.id, + user_id=resolved_user.id if resolved_user is not None else None, + ) + + # Engine-evaluated exits report the engine's resolved ids, mirroring the + # plain DataStream path's flag_check event. + engine_company_id = result.company_id or resolved_company.id + engine_user_id = result.user_id or (resolved_user.id if resolved_user is not None else None) + + if not result.value: + await _cancel_reservation(deps.reservations, record, log) + return await _emit_flag_check( + deps, + company, + user, + _check_result( + allowed=False, + value=False, + reason=result.reason or "denied_by_engine", + flag_key=result.flag_key or flag_key, + entitlement=_entitlement(result), + flag_id=result.flag_id, + ), + company_id=engine_company_id, + user_id=engine_user_id, + rule_id=result.rule_id, + ) + + # Allowed against the substituted balance, so the hold stands. Top the + # lease up in the background now that it has been drawn down. + deps.manager.extend_in_background(resolved_company.id, credit_id) + return await _emit_flag_check( + deps, + company, + user, + _check_result( + allowed=True, + value=True, + reason=result.reason or "lease_reserved", + flag_key=result.flag_key or flag_key, + reservation=_public_reservation(record), + entitlement=_entitlement(result), + flag_id=result.flag_id, + ), + company_id=engine_company_id, + user_id=engine_user_id, + rule_id=result.rule_id, + ) + + +async def _emit_flag_check( + deps: CreditCheckDeps, + req_company: Optional[Dict[str, str]], + req_user: Optional[Dict[str, str]], + result: "CheckResult", + *, + company_id: Optional[str] = None, + user_id: Optional[str] = None, + rule_id: Optional[str] = None, +) -> "CheckResult": + """Report a lease-path resolution and pass the result straight through. + + Analytics must never change a verdict the caller is already acting on, so a + failure here is logged and swallowed. + """ + try: + await deps.enqueue_flag_check( + EventBodyFlagCheck( + flag_key=result.flag_key, + value=result.value, + reason=result.reason, + error=result.error, + flag_id=result.flag_id, + company_id=company_id, + user_id=user_id, + rule_id=rule_id, + req_company=req_company, + req_user=req_user, + ) + ) + except Exception as err: + deps.logger.debug(f"Lease check: failed to report the flag_check event: {err}") + return result + + +async def _load_flag(datastream: CheckDataStream, flag_key: str, log: logging.Logger) -> Optional[RulesengineFlag]: + try: + return await datastream.get_flag(flag_key) + except Exception as err: + log.warning(f"Lease check: failed to load flag {flag_key}: {err}") + return None + + +async def _load_company( + datastream: CheckDataStream, keys: Dict[str, str], log: logging.Logger +) -> Optional[RulesengineCompany]: + try: + return await datastream.get_company(keys) + except Exception as err: + log.debug(f"Lease check: company fetch failed for keys {keys} ({err}), using a plain check") + return None + + +async def _load_user( + datastream: CheckDataStream, keys: Dict[str, str], log: logging.Logger +) -> Optional[RulesengineUser]: + try: + return await datastream.get_user(keys) + except Exception as err: + log.debug(f"Lease check: user fetch failed for keys {keys} ({err}), using a plain check") + return None + + +async def _handle_lease_failure( + *, + datastream: CheckDataStream, + log: logging.Logger, + mode: str, + flag_key: str, + reason: str, + flag: RulesengineFlag, + company: RulesengineCompany, + user: Optional[RulesengineUser], + credit_id: str, + options: "CheckOptions", +) -> "CheckResult": + """Resolve a check that could not gate: acquire failed, store unreachable, + or the lease is exhausted. + + fail-closed denies. fail-open means assume the credits are there, not skip + the evaluation: the rules still run with the balance substituted to an + effectively unlimited value, so plan targeting, overrides, and every + non-credit condition still apply, and a company that is not entitled stays + denied with the lease backend down. Only an error in that evaluation drops + to a blanket allow. + """ + if mode == "fail-closed": + return _static_failure_result(mode, flag_key, reason, flag) + + try: + substituted = _substitute_credit_balance(company, credit_id, FAIL_OPEN_BALANCE) + result = datastream.evaluate_flag(flag, substituted, user, _preflight_options(options)) + except Exception as err: + log.warning(f"Lease check: the fail-open evaluation failed ({err}); allowing") + return _static_failure_result(mode, flag_key, reason, flag) + return _check_result( + allowed=result.value, + value=result.value, + reason=f"{result.reason or 'evaluated'} ({reason}_fail_open)", + flag_key=result.flag_key or flag_key, + entitlement=_entitlement(result), + flag_id=result.flag_id or flag.id, + error=reason, + ) + + +def _static_failure_result( + mode: str, flag_key: str, reason: str, flag: Optional[RulesengineFlag] +) -> "CheckResult": + """Resolve a mode with no evaluation behind it: deny for fail-closed, + blanket allow for fail-open. + + Used when the engine is the thing that failed, and when the fail-open + evaluation itself errors. + """ + allowed = mode != "fail-closed" + return _check_result( + allowed=allowed, + value=allowed, + reason=f"{reason}_fail_open" if allowed else reason, + flag_key=flag_key, + flag_id=flag.id if flag is not None else None, + error=reason, + ) + + +async def _cancel_reservation( + reservations: ReservationStore, record: ReservationRecord, log: logging.Logger +) -> None: + """Claim the hold and refund all of it. Best effort: a failure leaves the + hold for the sweeper or for lease expiry.""" + try: + await reservations.consume(record.id, 0) + except Exception as err: + log.warning( + f"Lease check: failed to cancel reservation {record.id} ({err}); " + "its hold is reclaimed by the sweeper or at lease expiry" + ) + + +def _substitute_credit_balance( + company: RulesengineCompany, credit_id: str, balance: float +) -> RulesengineCompany: + balances = dict(company.credit_balances or {}) + balances[credit_id] = balance + return company.model_copy(update={"credit_balances": balances}) + + +def _credit_cost_options(credit_id: str, credit_cost: float) -> "CheckFlagOptions": + from ..client import CheckFlagOptions + + return CheckFlagOptions(credit_cost={credit_id: credit_cost}) + + +def _preflight_options(options: "CheckOptions") -> Optional["CheckFlagOptions"]: + from ..client import _check_options_to_flag_options + + return _check_options_to_flag_options(options) + + +def _entitlement(result: RulesengineCheckFlagResult) -> Optional[Any]: + if result.entitlement is None: + return None + from ..types import FeatureEntitlement + + return FeatureEntitlement.model_validate(result.entitlement.model_dump()) + + +def _public_reservation(record: ReservationRecord) -> "Reservation": + """The caller's handle on a hold carved out of a lease.""" + import datetime as dt + + from ..client import Reservation + + return Reservation( + id=record.id, + lease_id=record.lease_id, + mode="client", + company_id=record.company_id, + credit_type_id=record.credit_type_id, + event_subtype=record.event_subtype, + quantity_reserved=record.quantity_reserved, + credits_reserved=record.credits_reserved, + consumption_rate=record.consumption_rate, + expires_at=dt.datetime.fromtimestamp(record.expires_at, tz=dt.timezone.utc), + company=record.company, + user=record.user, + ) + + +def _check_result(**fields: Any) -> "CheckResult": + """CheckResult is defined on the client, which imports this module, so the + import waits until call time to keep the cycle from closing at import.""" + from ..client import CheckResult + + return CheckResult(**fields) diff --git a/src/schematic/leases/lease_manager.py b/src/schematic/leases/lease_manager.py new file mode 100644 index 00000000..141b93be --- /dev/null +++ b/src/schematic/leases/lease_manager.py @@ -0,0 +1,452 @@ +"""Lease lifecycle against the server: acquire, extend, release, sweep. + +Every path here resolves rather than raises. The manager's callers route a +missing lease through their fail-open/fail-closed handling, and several calls +are made fire-and-forget, where a raised exception would surface as an +unretrieved task exception instead. +""" + +from __future__ import annotations + +import asyncio +import datetime as dt +import logging +import time +from dataclasses import dataclass +from typing import Any, Awaitable, Dict, Optional, Protocol, Set + +from .lease_store import LeaseStore, lease_key +from .reservation_store import ReservationStore +from .types import ( + DEFAULT_SWEEP_INTERVAL, + Clock, + LeaseConfig, + LeaseState, + ResolvedLeaseConfig, + resolve_lease_config, +) + +logger = logging.getLogger(__name__) + + +@dataclass +class LeaseGrant: + """What the server says a lease is, after an acquire or an extend.""" + + lease_id: str + company_id: str + credit_type_id: str + # The server-authoritative TOTAL, not the increment an extend asked for. + granted_amount: float + expires_at: float + + +class LeaseWireClient(Protocol): + """The three lease calls the manager makes. + + Narrow on purpose: it keeps the manager independent of the generated + client's request and response models, and lets tests script the server. + """ + + async def acquire( + self, + company_id: str, + credit_type_id: str, + requested_amount: float, + expires_at: float, + timeout: Optional[float] = None, + ) -> LeaseGrant: ... + + async def extend( + self, + lease_id: str, + additional_amount: float, + expires_at: float, + timeout: Optional[float] = None, + ) -> LeaseGrant: ... + + async def release(self, lease_id: str) -> None: ... + + +class CreditsWireClient: + """Adapter over the generated async credits client (``AsyncCreditsClient``).""" + + def __init__(self, credits_client: Any, *, request_options: Optional[Any] = None) -> None: + self._credits = credits_client + self._request_options = request_options + + def _options(self, timeout: Optional[float]) -> Optional[Any]: + """The caller's per-check timeout wins over the client-wide options, + which is what a caller asking for one on this check means.""" + if timeout is None: + return self._request_options + return {"timeout": timeout} + + async def acquire( + self, + company_id: str, + credit_type_id: str, + requested_amount: float, + expires_at: float, + timeout: Optional[float] = None, + ) -> LeaseGrant: + response = await self._credits.acquire_credit_lease( + company_id=company_id, + credit_type_id=credit_type_id, + requested_amount=requested_amount, + expires_at=_to_datetime(expires_at), + request_options=self._options(timeout), + ) + return _grant_from_response(response) + + async def extend( + self, + lease_id: str, + additional_amount: float, + expires_at: float, + timeout: Optional[float] = None, + ) -> LeaseGrant: + response = await self._credits.extend_credit_lease( + lease_id, + additional_amount=additional_amount, + expires_at=_to_datetime(expires_at), + request_options=self._options(timeout), + ) + return _grant_from_response(response) + + async def release(self, lease_id: str) -> None: + await self._credits.release_credit_lease(lease_id, request_options=self._request_options) + + +class LeaseManager: + """Owns lease rows for one client: acquire on first use or after expiry, + extend when the local view dips below the water mark, release on close. + + Acquire and extend each get their own best-effort single-flight map keyed + by slot. Best-effort because callers racing ahead of the registration can + still issue duplicate wire calls, which is safe: the server is idempotent + for an active slot, ``replace`` keeps the first live lease, and ``extend`` + reconciles to a total. + """ + + def __init__( + self, + wire_client: LeaseWireClient, + lease_store: LeaseStore, + *, + reservation_store: Optional[ReservationStore] = None, + config: Optional[LeaseConfig] = None, + clock: Clock = time.time, + ) -> None: + self._wire = wire_client + self._lease_store = lease_store + self._reservation_store = reservation_store + self._config = config or LeaseConfig() + self._clock = clock + # Kept separate so an in-flight extend can never satisfy an acquire, + # or the other way round. + self._inflight_acquire: Dict[str, "asyncio.Future[Optional[LeaseState]]"] = {} + self._inflight_extend: Dict[str, "asyncio.Future[Optional[LeaseState]]"] = {} + self._background: Set["asyncio.Task[None]"] = set() + self._sweep_task: Optional["asyncio.Task[None]"] = None + self._stopped = False + + def resolve_config(self, credit_type_id: str) -> ResolvedLeaseConfig: + return resolve_lease_config(self._config, None, credit_type_id) + + @property + def sweep_interval(self) -> float: + return self._config.sweep_interval or DEFAULT_SWEEP_INTERVAL + + async def acquire_if_needed( + self, company_id: str, credit_type_id: str, timeout: Optional[float] = None + ) -> Optional[LeaseState]: + """The slot's live lease, acquiring one over the wire if none is live. + + ``timeout`` governs the wire call this caller starts. A caller that + joins an in-flight acquire rides the first caller's timeout, since + there is one shared call to time out. + """ + try: + existing = await self._lease_store.get(company_id, credit_type_id) + except Exception as err: + logger.error("Failed to read lease store for %s/%s: %s", company_id, credit_type_id, err) + return None + if existing is not None and existing.expires_at > self._clock(): + return existing + # An expired (or absent) slot is left for `replace` to overwrite: it + # guards on expiry and writes atomically. Dropping the stale row first + # would be a separate, non-atomic op that can interleave between a + # sibling pod's read and its replace, clobbering a lease that pod just + # installed. Reading a stale entry in the gap is harmless, since every + # path that acts on a lease re-guards on expiry. + + key = lease_key(company_id, credit_type_id) + inflight = self._inflight_acquire.get(key) + if inflight is not None: + return await asyncio.shield(inflight) + return await self._single_flight( + self._inflight_acquire, key, self._acquire(company_id, credit_type_id, timeout) + ) + + async def _acquire( + self, company_id: str, credit_type_id: str, timeout: Optional[float] = None + ) -> Optional[LeaseState]: + resolved = self.resolve_config(credit_type_id) + try: + grant = await self._wire.acquire( + company_id, + credit_type_id, + resolved.lease_size, + self._clock() + resolved.lease_duration, + timeout, + ) + wrote = await self._lease_store.replace( + lease_id=grant.lease_id, + company_id=grant.company_id or company_id, + credit_type_id=grant.credit_type_id or credit_type_id, + granted_amount=grant.granted_amount, + expires_at=grant.expires_at, + ) + if wrote: + return await self._lease_store.get(company_id, credit_type_id) + + # A sibling holds the slot with a live lease, or the slot's expired + # row was reconciled in place. The server is idempotent for an + # active slot, so a racing acquire is normally handed back the SAME + # lease the sibling installed, and releasing it would pull the + # shared lease out from under every pod. Only a *different* lease + # is a redundant hold nobody will draw on, so only that one is + # released. An empty slot (expired in the gap) releases nothing + # either: this lease is likely what the next acquire is handed. + current = await self._lease_store.get(company_id, credit_type_id) + if current is not None and current.lease_id != grant.lease_id: + logger.debug( + "Lost acquire race for %s/%s; releasing redundant lease %s", + company_id, + credit_type_id, + grant.lease_id, + ) + self._spawn(self._release(grant.lease_id)) + return current + except Exception as err: + logger.error("Failed to acquire credit lease for %s/%s: %s", company_id, credit_type_id, err) + return None + + async def maybe_extend( + self, + company_id: str, + credit_type_id: str, + required_credits: Optional[float] = None, + timeout: Optional[float] = None, + ) -> Optional[LeaseState]: + """Extend the slot's lease when the local view warrants it. + + Triggered by either the low-water-mark ratio (steady-state refresh) or + a ``required_credits`` hint above the local remaining (a check just + failed a reserve of that size). + """ + try: + entry = await self._lease_store.get(company_id, credit_type_id) + except Exception as err: + logger.warning("Failed to read lease store for %s/%s: %s", company_id, credit_type_id, err) + return None + if entry is None: + return None + # Never extend an expired lease: the server treats it as released and + # has already refunded its remainder, so the only correct move is a + # fresh acquire on the next check. + if entry.expires_at <= self._clock(): + return None + resolved = self.resolve_config(credit_type_id) + ratio = entry.local_remaining_credits / max(entry.granted_amount, 1) + below_watermark = ratio <= resolved.low_water_mark + below_required = required_credits is not None and entry.local_remaining_credits < required_credits + if not below_watermark and not below_required: + return entry + + key = lease_key(company_id, credit_type_id) + inflight = self._inflight_extend.get(key) + if inflight is not None: + return await asyncio.shield(inflight) + return await self._single_flight( + self._inflight_extend, key, self._extend(entry, resolved, required_credits, timeout) + ) + + async def _extend( + self, + entry: LeaseState, + resolved: ResolvedLeaseConfig, + required_credits: Optional[float], + timeout: Optional[float] = None, + ) -> Optional[LeaseState]: + # Size the extend to cover the request that triggered it: a single + # check needing more than remaining plus one tranche would otherwise + # fail its post-extend retry forever, however much balance the server + # has. The steady-state path keeps asking for the configured tranche. + shortfall = (required_credits - entry.local_remaining_credits) if required_credits is not None else 0.0 + try: + grant = await self._wire.extend( + entry.lease_id, + max(resolved.lease_size, shortfall), + self._clock() + resolved.lease_duration, + timeout, + ) + # Reconcile to the server's authoritative TOTAL, with the store + # computing the delta against its own current total: per-process + # single-flight does not cover sibling pods. Pinned to the lease + # the server extended, so an expiry mid-call cannot mint the delta + # onto a successor. + await self._lease_store.extend( + entry.company_id, + entry.credit_type_id, + grant.granted_amount, + grant.expires_at, + entry.lease_id, + ) + return await self._lease_store.get(entry.company_id, entry.credit_type_id) + except Exception as err: + logger.warning("Failed to extend credit lease %s: %s", entry.lease_id, err) + return None + + def extend_in_background(self, company_id: str, credit_type_id: str) -> None: + """Kick off a water-mark extend without waiting for it. + + A check that just drew the lease down should not pay for the top-up, so + the extend runs as tracked background work and never raises into the + caller. + """ + + async def run() -> None: + await self.maybe_extend(company_id, credit_type_id) + + self._spawn(run()) + + async def release_all_local_leases(self) -> None: + """Release every live lease this process exclusively holds. + + Only a per-process store answers ``list_leases``; a shared backend + returns ``None`` and is skipped, since sibling pods still draw on those + leases. Expired leases are skipped too: the server already swept them. + Best-effort, with failures falling back to server-side expiry. + """ + try: + entries = self._lease_store.list_leases() + except Exception as err: + logger.warning("Failed to enumerate leases on close: %s", err) + return + if not entries: + return + now = self._clock() + for entry in entries: + if entry.expires_at <= now: + continue + try: + await self._wire.release(entry.lease_id) + await self._lease_store.drop(entry.company_id, entry.credit_type_id) + logger.debug("Released credit lease %s on close", entry.lease_id) + except Exception as err: + logger.warning( + "Failed to release credit lease %s on close (it will expire server-side): %s", + entry.lease_id, + err, + ) + + def start_sweep(self) -> None: + """Run the expired-reservation sweep on an interval. Safe to call twice.""" + if self._reservation_store is None or self._sweep_task is not None or self._stopped: + return + try: + loop = asyncio.get_running_loop() + except RuntimeError: + logger.debug("No running event loop; the reservation sweep stays off") + return + self._sweep_task = loop.create_task(self._sweep_loop()) + + async def _sweep_loop(self) -> None: + store = self._reservation_store + assert store is not None + while True: + await asyncio.sleep(self.sweep_interval) + try: + await store.sweep_expired() + except asyncio.CancelledError: + raise + except Exception as err: + # Keep the loop alive: a sweep failure is transient (a Redis + # blip), and the next tick retries. + logger.debug("Reservation sweep failed: %s", err) + + def stop(self) -> None: + """Cancel the sweep loop. Pending releases are left to finish.""" + self._stopped = True + if self._sweep_task is not None: + self._sweep_task.cancel() + self._sweep_task = None + + async def _single_flight( + self, + registry: Dict[str, "asyncio.Future[Optional[LeaseState]]"], + key: str, + coro: Awaitable[Optional[LeaseState]], + ) -> Optional[LeaseState]: + task = asyncio.ensure_future(coro) + registry[key] = task + try: + return await asyncio.shield(task) + finally: + if registry.get(key) is task: + del registry[key] + + async def _release(self, lease_id: str) -> None: + try: + await self._wire.release(lease_id) + except Exception as err: + logger.warning("Failed to release redundant credit lease %s: %s", lease_id, err) + + def _spawn(self, coro: Awaitable[None]) -> None: + """Run a fire-and-forget step, holding a reference so it is not collected.""" + try: + task = asyncio.get_running_loop().create_task(_never_raises(coro)) + except RuntimeError: + logger.debug("No running event loop; skipping background lease work") + return + self._background.add(task) + task.add_done_callback(self._background.discard) + + async def _drain_background(self) -> None: + """Wait out pending fire-and-forget work. For tests and close paths.""" + while self._background: + await asyncio.gather(*list(self._background), return_exceptions=True) + + +async def _never_raises(coro: Awaitable[None]) -> None: + try: + await coro + except asyncio.CancelledError: + raise + except Exception as err: + logger.debug("Background lease work failed: %s", err) + + +def _to_datetime(epoch_seconds: float) -> dt.datetime: + return dt.datetime.fromtimestamp(epoch_seconds, tz=dt.timezone.utc) + + +def _epoch_seconds(value: dt.datetime) -> float: + # A naive timestamp from the API is UTC; reading it as local time would + # shift every expiry by the pod's offset. + if value.tzinfo is None: + return value.replace(tzinfo=dt.timezone.utc).timestamp() + return value.timestamp() + + +def _grant_from_response(response: Any) -> LeaseGrant: + data = response.data + return LeaseGrant( + lease_id=data.id, + company_id=data.company_id, + credit_type_id=data.credit_type_id, + granted_amount=float(data.granted_amount), + expires_at=_epoch_seconds(data.expires_at), + ) diff --git a/src/schematic/leases/lease_store.py b/src/schematic/leases/lease_store.py new file mode 100644 index 00000000..0742ff77 --- /dev/null +++ b/src/schematic/leases/lease_store.py @@ -0,0 +1,270 @@ +"""Lease slot storage: the contract, plus the per-process in-memory backend. + +At most one lease occupies a ``(company_id, credit_type_id)`` slot. Every +mutation is atomic per slot; ``RedisLeaseStore`` gets that from single-key Lua, +this one from a per-slot ``asyncio.Lock``. +""" + +from __future__ import annotations + +import abc +import asyncio +import math +import time +from contextlib import asynccontextmanager +from typing import AsyncIterator, Dict, List, Optional, Tuple + +from .types import Clock, LeaseState + + +def lease_key(company_id: str, credit_type_id: str) -> str: + return f"{company_id}:{credit_type_id}" + + +class LeaseStore(abc.ABC): + """Backing store for lease slots, shared by the in-memory and Redis backends.""" + + @abc.abstractmethod + async def get(self, company_id: str, credit_type_id: str) -> Optional[LeaseState]: + """Snapshot of the slot, expired or not. Callers re-guard on expiry.""" + + @abc.abstractmethod + async def replace( + self, + *, + lease_id: str, + company_id: str, + credit_type_id: str, + granted_amount: float, + expires_at: float, + ) -> bool: + """Install a fresh lease at its full grant, if the slot is free to take. + + A *live* lease holds the slot even when it carries a different id (a + sibling pod won the acquire race): its already-debited balance wins and + this reports ``False``. An *expired* row carrying the SAME id is not + rewritten either, since that would reset the balance and erase debits + whose reservations are still open; it is reconciled like an extend + (granted to the incoming total, expiry forward only, balance untouched) + and also reports ``False``. Returns ``True`` only when a fresh row was + written, which is what tells the manager whether the lease it just + acquired is redundant. + """ + + @abc.abstractmethod + async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[float]: + """Atomically check and debit, returning the post-debit balance. + + ``None`` when there is no lease, it has expired, the balance is short, + or ``credits`` is not a finite non-negative number. Returning the + balance (rather than a bool) lets the caller derive the pre-debit + figure as ``returned + credits`` without a racy follow-up read. + """ + + @abc.abstractmethod + async def refund( + self, + company_id: str, + credit_type_id: str, + credits: float, + pin_lease_id: Optional[str] = None, + ) -> None: + """Return credits to the slot's balance, clamped at the granted amount. + + With ``pin_lease_id``, the refund applies only while the slot still + holds that lease: a hold carved out of an expired lease must never + inflate its successor, whose grant the server already issued whole. + """ + + @abc.abstractmethod + async def extend( + self, + company_id: str, + credit_type_id: str, + granted_total: float, + new_expires_at: Optional[float] = None, + pin_lease_id: Optional[str] = None, + ) -> None: + """Reconcile the slot to the server-authoritative total. + + The delta is computed inside the store against the currently stored + total, never from a caller-held pre-wire-call read: two pods extending + concurrently from the same stale read would each apply a delta and mint + phantom credits. A total a sibling already applied is a no-op, so + applies converge in any order. Expiry only ever moves forward. + """ + + @abc.abstractmethod + async def drop(self, company_id: str, credit_type_id: str) -> None: + """Remove the slot entry, after a remote release.""" + + def list_leases(self) -> Optional[List[LeaseState]]: + """Every lease this store holds, or ``None`` when it cannot enumerate. + + Only a per-process store answers: its leases are exclusively this + process's, so releasing them on close is safe. A shared backend must + never enumerate and release, since sibling pods still draw on those + leases. + """ + return None + + +class _SlotLocks: + """Per-slot mutual exclusion, refcounted so idle slots do not accumulate.""" + + def __init__(self) -> None: + self._locks: Dict[str, Tuple[asyncio.Lock, int]] = {} + + @asynccontextmanager + async def hold(self, key: str) -> AsyncIterator[None]: + lock, waiters = self._locks.get(key, (asyncio.Lock(), 0)) + self._locks[key] = (lock, waiters + 1) + try: + async with lock: + yield + finally: + held, count = self._locks[key] + if count <= 1: + del self._locks[key] + else: + self._locks[key] = (held, count - 1) + + +class InMemoryLeaseStore(LeaseStore): + """Per-process lease slots. Single-pod gating only. + + Swap in ``RedisLeaseStore`` to gate across pods; both implement the same + contract. + """ + + def __init__(self, *, clock: Clock = time.time) -> None: + self._clock = clock + self._leases: Dict[str, LeaseState] = {} + self._locks = _SlotLocks() + + async def get(self, company_id: str, credit_type_id: str) -> Optional[LeaseState]: + entry = self._leases.get(lease_key(company_id, credit_type_id)) + return _copy(entry) if entry else None + + async def replace( + self, + *, + lease_id: str, + company_id: str, + credit_type_id: str, + granted_amount: float, + expires_at: float, + ) -> bool: + key = lease_key(company_id, credit_type_id) + async with self._locks.hold(key): + existing = self._leases.get(key) + if existing is not None and existing.expires_at > self._clock(): + return False + if existing is not None and existing.lease_id == lease_id: + # The same lease coming back over its own expired row: a stale + # acquire response for a lease the idempotent server also + # handed a racing sibling, which may since have extended it. + add = granted_amount - existing.granted_amount + if add > 0: + existing.granted_amount = granted_amount + existing.local_remaining_credits += add + if expires_at > existing.expires_at: + existing.expires_at = expires_at + return False + self._leases[key] = LeaseState( + lease_id=lease_id, + company_id=company_id, + credit_type_id=credit_type_id, + granted_amount=granted_amount, + local_remaining_credits=granted_amount, + expires_at=expires_at, + ) + return True + + async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[float]: + # NaN passes every comparison below, and a NaN balance would approve + # every later reserve, so it never reaches the arithmetic. + if not is_finite_non_negative(credits): + return None + key = lease_key(company_id, credit_type_id) + async with self._locks.hold(key): + entry = self._leases.get(key) + if entry is None: + return None + if entry.expires_at <= self._clock(): + return None + if entry.local_remaining_credits < credits: + return None + entry.local_remaining_credits -= credits + return entry.local_remaining_credits + + async def refund( + self, + company_id: str, + credit_type_id: str, + credits: float, + pin_lease_id: Optional[str] = None, + ) -> None: + if not is_finite_non_negative(credits) or credits <= 0: + return + key = lease_key(company_id, credit_type_id) + async with self._locks.hold(key): + entry = self._leases.get(key) + if entry is None: + return + if pin_lease_id is not None and entry.lease_id != pin_lease_id: + return + entry.local_remaining_credits = min( + entry.local_remaining_credits + credits, + entry.granted_amount, + ) + + async def extend( + self, + company_id: str, + credit_type_id: str, + granted_total: float, + new_expires_at: Optional[float] = None, + pin_lease_id: Optional[str] = None, + ) -> None: + key = lease_key(company_id, credit_type_id) + async with self._locks.hold(key): + entry = self._leases.get(key) + if entry is None: + return + if pin_lease_id is not None and entry.lease_id != pin_lease_id: + return + add = granted_total - entry.granted_amount + if add > 0: + entry.granted_amount = granted_total + entry.local_remaining_credits += add + if new_expires_at is not None and new_expires_at > entry.expires_at: + entry.expires_at = new_expires_at + + async def drop(self, company_id: str, credit_type_id: str) -> None: + key = lease_key(company_id, credit_type_id) + async with self._locks.hold(key): + self._leases.pop(key, None) + + def list_leases(self) -> Optional[List[LeaseState]]: + return [_copy(entry) for entry in self._leases.values()] + + +def _copy(entry: LeaseState) -> LeaseState: + return LeaseState( + lease_id=entry.lease_id, + company_id=entry.company_id, + credit_type_id=entry.credit_type_id, + granted_amount=entry.granted_amount, + local_remaining_credits=entry.local_remaining_credits, + expires_at=entry.expires_at, + ) + + +def is_finite_non_negative(value: float) -> bool: + """A credit amount fit for arithmetic: finite, and not negative.""" + try: + number = float(value) + except (TypeError, ValueError): + return False + return math.isfinite(number) and number >= 0 diff --git a/src/schematic/leases/redis_lease_store.py b/src/schematic/leases/redis_lease_store.py new file mode 100644 index 00000000..0accf652 --- /dev/null +++ b/src/schematic/leases/redis_lease_store.py @@ -0,0 +1,349 @@ +"""Redis-backed lease slots: one hash per slot, mutated by single-key Lua. + +The key layout, hash field names (camelCase), millisecond instants, Lua +scripts, and TTL grace windows match the Node SDK exactly, so Node and Python +pods can share one Redis and agree on every slot. +""" + +from __future__ import annotations + +import time +from typing import Any, Dict, List, Optional + +from .lease_store import LeaseStore, is_finite_non_negative, lease_key +from .types import DEFAULT_LEASE_DURATION, Clock, LeaseState + +DEFAULT_KEY_PREFIX = "schematic:" +LEASE_KEY_NAMESPACE = "credit-lease:" +# How long after the declared expiry the row survives before Redis evicts it. +# Gives the sweeper a window to refund expired reservations before the lease +# state underneath them disappears. +LEASE_TTL_GRACE_MS = 60_000 + +# Every Lua script below touches exactly ONE key (the lease hash), keeping +# them safe under Redis Cluster (multi-key scripts spanning slots raise +# CROSSSLOT). Only the lease hash needs atomic mutation; cross-key +# bookkeeping uses ordinary single-key commands. +# +# Expiry is decided against the *Redis server's* clock (`redis.call('TIME')`), +# not the calling pod's: with many pods sharing one lease, local clock skew +# would let pods disagree on whether the lease is live. The `LEASE_NOW_MS` +# snippet converts TIME to integer milliseconds (matching the stored +# `expiresAt`); `redis.replicate_commands()` first, so the non-deterministic +# TIME read is allowed alongside writes on Redis 5/6. +LEASE_NOW_MS = """ +redis.replicate_commands() +local t = redis.call('TIME') +local now = (tonumber(t[1]) * 1000) + math.floor(tonumber(t[2]) / 1000) +""" + +# Atomic `replace`. Writes the lease hash only when the slot is empty or the +# existing lease has expired. Returns 1 on write, 0 if a *live* lease already +# occupies the slot, even one with a different leaseId, e.g. installed by a +# sibling instance that raced this acquire. An expired row with the SAME +# leaseId is reconciled like an extend instead of rewritten, which would reset +# the balance and erase debits whose reservations are still open. +REPLACE_SCRIPT = ( + LEASE_NOW_MS + + """ +local existing_id = redis.call('HGET', KEYS[1], 'leaseId') +local existing_expiry = tonumber(redis.call('HGET', KEYS[1], 'expiresAt') or '0') +local new_id = ARGV[1] +local new_granted = ARGV[2] +local new_expiry = tonumber(ARGV[3]) +local grace = tonumber(ARGV[4]) + +if existing_id and existing_expiry > now then + return 0 +end + +if existing_id == new_id then + local granted = tonumber(redis.call('HGET', KEYS[1], 'grantedAmount') or '0') + local add = tonumber(new_granted) - granted + if add > 0 then + local remaining = tonumber(redis.call('HGET', KEYS[1], 'localRemainingCredits') or '0') + redis.call('HSET', KEYS[1], + 'grantedAmount', new_granted, + 'localRemainingCredits', tostring(remaining + add)) + end + if new_expiry > existing_expiry then + redis.call('HSET', KEYS[1], 'expiresAt', ARGV[3]) + redis.call('PEXPIREAT', KEYS[1], new_expiry + grace) + end + return 0 +end + +redis.call('DEL', KEYS[1]) +redis.call('HSET', KEYS[1], + 'leaseId', new_id, + 'companyId', ARGV[5], + 'creditTypeId', ARGV[6], + 'grantedAmount', new_granted, + 'localRemainingCredits', new_granted, + 'expiresAt', ARGV[3]) +redis.call('PEXPIREAT', KEYS[1], new_expiry + grace) +return 1 +""" +) + +# Atomic check-and-decrement on `localRemainingCredits`. Returns the post-debit +# balance as a string (a Lua number reply truncates to integer, which would +# corrupt fractional credit costs); nil if there is no lease, the lease has +# expired, or there is insufficient remaining. The expiry guard compares +# against the Redis server clock, so a reserve against an expired-but-not-yet- +# evicted row during the TTL grace window is rejected. +TRY_RESERVE_SCRIPT = ( + LEASE_NOW_MS + + """ +local raw = redis.call('HGET', KEYS[1], 'localRemainingCredits') +if not raw then return false end +local expiry = tonumber(redis.call('HGET', KEYS[1], 'expiresAt') or '0') +if expiry <= now then return false end +local remaining = tonumber(raw) +local requested = tonumber(ARGV[1]) +if remaining < requested then return false end +local new_remaining = remaining - requested +redis.call('HSET', KEYS[1], 'localRemainingCredits', tostring(new_remaining)) +return tostring(new_remaining) +""" +) + +# Refund credits, clamped at `grantedAmount`. ARGV[2], when non-empty, pins the +# refund to a specific leaseId: if the slot now holds a different lease, the +# refund is dropped: the expired lease's unspent remainder was already +# returned to the company balance server-side, so crediting the successor would +# mint phantom credits. +REFUND_SCRIPT = """ +local raw_remaining = redis.call('HGET', KEYS[1], 'localRemainingCredits') +if not raw_remaining then return 0 end +local required_lease = ARGV[2] +if required_lease and required_lease ~= '' then + local current_lease = redis.call('HGET', KEYS[1], 'leaseId') + if current_lease ~= required_lease then return 0 end +end +local remaining = tonumber(raw_remaining) +local granted = tonumber(redis.call('HGET', KEYS[1], 'grantedAmount') or '0') +local refund = tonumber(ARGV[1]) +local new_balance = remaining + refund +if new_balance > granted then new_balance = granted end +redis.call('HSET', KEYS[1], 'localRemainingCredits', tostring(new_balance)) +return 1 +""" + +# Reconcile the lease to the server-authoritative grantedAmount total +# (ARGV[1]), crediting the difference to localRemainingCredits. The delta is +# computed HERE, atomically against the hash's current total, never by the +# caller from a pre-wire-call read: two pods extending the same shared lease +# concurrently would each apply a delta against the same stale read and mint +# phantom credits. Reconciling to the absolute total converges regardless of +# arrival order. Expiry only ever moves forward. ARGV[4], when non-empty, pins +# the extend to a leaseId, mirroring REFUND_SCRIPT. +EXTEND_SCRIPT = """ +local raw_granted = redis.call('HGET', KEYS[1], 'grantedAmount') +if not raw_granted then return 0 end +local required_lease = ARGV[4] +if required_lease and required_lease ~= '' then + local current_lease = redis.call('HGET', KEYS[1], 'leaseId') + if current_lease ~= required_lease then return 0 end +end +local granted = tonumber(raw_granted) +local target = tonumber(ARGV[1]) +local add = target - granted +if add > 0 then + local remaining = tonumber(redis.call('HGET', KEYS[1], 'localRemainingCredits') or '0') + redis.call('HSET', KEYS[1], + 'grantedAmount', tostring(target), + 'localRemainingCredits', tostring(remaining + add)) +end +local new_expiry = tonumber(ARGV[2]) +local grace = tonumber(ARGV[3]) +local current_expiry = tonumber(redis.call('HGET', KEYS[1], 'expiresAt') or '0') +if new_expiry > current_expiry then + redis.call('HSET', KEYS[1], 'expiresAt', ARGV[2]) + redis.call('PEXPIREAT', KEYS[1], new_expiry + grace) +end +return 1 +""" + + +class LuaScript: + """One Lua script, run by EVALSHA with a lazy SCRIPT LOAD and an EVAL fallback. + + A Redis that has dropped its script cache (restart, SCRIPT FLUSH) answers + NOSCRIPT; the fallback re-sends the body and re-loads it on the next call. + """ + + def __init__(self, body: str) -> None: + self.body = body + self._sha: Optional[str] = None + + async def run(self, client: Any, keys: List[str], args: List[str]) -> Any: + if self._sha is None: + self._sha = to_str(await client.script_load(self.body)) + try: + return await client.evalsha(self._sha, len(keys), *keys, *args) + except Exception as err: + if not _is_noscript(err): + raise + self._sha = None + return await client.eval(self.body, len(keys), *keys, *args) + + +class RedisLeaseStore(LeaseStore): + """Lease slots in Redis: one hash per slot, atomic via single-key Lua. + + ``client`` is a connected ``redis.asyncio.Redis``. Balances are stored as + strings so fractional credit amounts survive the round trip. + """ + + def __init__( + self, + client: Any, + *, + key_prefix: str = DEFAULT_KEY_PREFIX, + default_lease_duration: float = DEFAULT_LEASE_DURATION, + clock: Clock = time.time, + ) -> None: + self._client = client + self._key_prefix = key_prefix + # Only reached when a direct caller extends without an expiry; the + # lease manager always passes one. + self._default_lease_duration = default_lease_duration + self._clock = clock + self._replace = LuaScript(REPLACE_SCRIPT) + self._try_reserve = LuaScript(TRY_RESERVE_SCRIPT) + self._refund = LuaScript(REFUND_SCRIPT) + self._extend = LuaScript(EXTEND_SCRIPT) + + def hash_key(self, company_id: str, credit_type_id: str) -> str: + """Public so the reservation store can target the same lease hash.""" + return f"{self._key_prefix}{LEASE_KEY_NAMESPACE}{lease_key(company_id, credit_type_id)}" + + async def get(self, company_id: str, credit_type_id: str) -> Optional[LeaseState]: + raw = decode_hash(await self._client.hgetall(self.hash_key(company_id, credit_type_id))) + if not raw.get("leaseId"): + return None + return LeaseState( + lease_id=raw["leaseId"], + company_id=raw.get("companyId", company_id), + credit_type_id=raw.get("creditTypeId", credit_type_id), + granted_amount=float(raw.get("grantedAmount", "0")), + local_remaining_credits=float(raw.get("localRemainingCredits", "0")), + expires_at=float(raw.get("expiresAt", "0")) / 1000.0, + ) + + async def replace( + self, + *, + lease_id: str, + company_id: str, + credit_type_id: str, + granted_amount: float, + expires_at: float, + ) -> bool: + result = await self._replace.run( + self._client, + [self.hash_key(company_id, credit_type_id)], + # No client clock here: the script reads `now` from the Redis + # server via TIME, so every pod agrees on expiry. + [ + lease_id, + format_amount(granted_amount), + to_epoch_ms(expires_at), + str(LEASE_TTL_GRACE_MS), + company_id, + credit_type_id, + ], + ) + return _to_number(result) == 1 + + async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[float]: + # Reject non-finite/negative debits before they reach the script: the + # string form of NaN parses back to a Lua nan, slips through the `<` + # comparison, and would poison the SHARED balance for every pod. + if not is_finite_non_negative(credits): + return None + result = await self._try_reserve.run( + self._client, + [self.hash_key(company_id, credit_type_id)], + [format_amount(credits)], + ) + if result is None or result is False: + return None + return float(to_str(result)) + + async def refund( + self, + company_id: str, + credit_type_id: str, + credits: float, + pin_lease_id: Optional[str] = None, + ) -> None: + if not is_finite_non_negative(credits) or credits <= 0: + return + await self._refund.run( + self._client, + [self.hash_key(company_id, credit_type_id)], + # An empty string disables the lease pin (Lua has no nil ARGV). + [format_amount(credits), pin_lease_id or ""], + ) + + async def extend( + self, + company_id: str, + credit_type_id: str, + granted_total: float, + new_expires_at: Optional[float] = None, + pin_lease_id: Optional[str] = None, + ) -> None: + expiry = new_expires_at if new_expires_at is not None else self._clock() + self._default_lease_duration + await self._extend.run( + self._client, + [self.hash_key(company_id, credit_type_id)], + # granted_total is the server-authoritative TOTAL; the script + # computes the delta against the stored total. + [format_amount(granted_total), to_epoch_ms(expiry), str(LEASE_TTL_GRACE_MS), pin_lease_id or ""], + ) + + async def drop(self, company_id: str, credit_type_id: str) -> None: + # A plain single-key delete: no secondary index to keep in sync. + await self._client.delete(self.hash_key(company_id, credit_type_id)) + + +def _is_noscript(err: Exception) -> bool: + """Did Redis answer that it no longer holds this script? + + Matched by exception name as well as message so the check does not depend + on importing redis, nor on a client's exact wording. + """ + return type(err).__name__ == "NoScriptError" or "NOSCRIPT" in str(err).upper() + + +def to_str(value: Any) -> str: + return value.decode("utf-8") if isinstance(value, bytes) else str(value) + + +def decode_hash(raw: Any) -> Dict[str, str]: + if not raw: + return {} + return {to_str(key): to_str(value) for key, value in raw.items()} + + +def _to_number(value: Any) -> float: + if value is None or value is False: + return 0.0 + if value is True: + return 1.0 + return float(to_str(value)) + + +def to_epoch_ms(epoch_seconds: float) -> str: + """Instants cross the wire as integer milliseconds, as the Node SDK writes them.""" + return str(int(round(epoch_seconds * 1000))) + + +def format_amount(value: float) -> str: + """Format a credit amount the way JavaScript would, so shared rows read alike.""" + if float(value).is_integer(): + return str(int(value)) + return repr(float(value)) diff --git a/src/schematic/leases/redis_reservation_store.py b/src/schematic/leases/redis_reservation_store.py new file mode 100644 index 00000000..be3c3e84 --- /dev/null +++ b/src/schematic/leases/redis_reservation_store.py @@ -0,0 +1,278 @@ +"""Redis-backed reservation table: one hash per hold, two secondary indexes. + +Key layout, hash fields (camelCase), millisecond instants, the claim script, +and the TTL grace window match the Node SDK, so Node and Python pods sharing +one Redis read each other's holds. + +Every mutation is a single-key operation (or single-key Lua), so the store is +correct on standalone and clustered Redis alike: the unspent-slice refund is +delegated to the lease store rather than reaching across to the lease hash +inside a multi-key script. +""" + +from __future__ import annotations + +import json +import time +from typing import Any, Dict, List, Optional, Tuple + +from .lease_store import LeaseStore +from .redis_lease_store import DEFAULT_KEY_PREFIX, LuaScript, decode_hash, format_amount, to_epoch_ms, to_str +from .reservation_store import ReservationStore, clamp_consumption +from .types import Clock, ReservationRecord + +RES_KEY_NAMESPACE = "credit-reservation:" +# Sorted set scoring open reservations by expiry so the sweeper can pop expired +# entries in O(log n). Members encode the full (company, credit, id) tuple. +RES_INDEX_KEY = "credit-reservations:byExpiry" +# Per-(company, credit) index of open holds, one hash of id -> creditsReserved, +# so reserved_credits reads a tenant's holds with one HGETALL. The hash is also +# the source of truth for that sum: a field exists exactly while its +# reservation is open and unrefunded. +RES_BYCREDIT_NAMESPACE = "credit-reservations:byCredit:" +# Buffer past expiry before Redis evicts the row, so the sweeper has a window +# to refund. +RES_TTL_GRACE_MS = 30_000 + +# Page size for the sweeper's ZRANGEBYSCORE. Without a limit, a backlog of +# expired holds (after a Redis outage or a long pod pause) would come back as +# one giant reply on every pod's next tick; paging bounds the reply while the +# per-member ZREM keeps offset 0 advancing through the backlog. +SWEEP_BATCH_SIZE = 256 +# Upper bound on pages per tick: keeps one sweep's work bounded and guards +# against an endless loop if ZREM persistently fails. Anything left over is +# picked up next tick. +MAX_SWEEP_BATCHES = 16 + +# Atomic claim: read the reservation hash and delete it in one step, returning +# its fields (or nil if it was already gone). Touches a single key. The atomic +# read-then-delete is what makes consume exactly-once: of two racing callers (a +# normal settle and a sweeper, say) only one gets the fields back and proceeds +# to refund. The refund to the lease hash is a separate single-key op; a crash +# in the gap leaves the unspent slice held on the lease until the lease itself +# expires, never double-refunded. +CLAIM_SCRIPT = """ +local raw = redis.call('HGETALL', KEYS[1]) +if #raw == 0 then return nil end +redis.call('DEL', KEYS[1]) +return raw +""" + + +class RedisReservationStore(ReservationStore): + """Reservation table in Redis, refunding through a lease store it is given. + + ``client`` is a connected ``redis.asyncio.Redis``. + """ + + def __init__( + self, + client: Any, + lease_store: LeaseStore, + *, + key_prefix: str = DEFAULT_KEY_PREFIX, + clock: Clock = time.time, + ) -> None: + self._client = client + self._lease_store = lease_store + self._key_prefix = key_prefix + self._clock = clock + self._claim = LuaScript(CLAIM_SCRIPT) + + def _hash_key(self, reservation_id: str) -> str: + return f"{self._key_prefix}{RES_KEY_NAMESPACE}{reservation_id}" + + def _index_key(self) -> str: + return f"{self._key_prefix}{RES_INDEX_KEY}" + + def _by_credit_key(self, company_id: str, credit_type_id: str) -> str: + return f"{self._key_prefix}{RES_BYCREDIT_NAMESPACE}{company_id}:{credit_type_id}" + + async def add(self, reservation: ReservationRecord) -> None: + expires_ms = int(to_epoch_ms(reservation.expires_at)) + hash_key = self._hash_key(reservation.id) + # The hash goes out first so the reservation exists before anything + # references it. These are independent single-key ops rather than one + # multi-key script: a partial failure at worst leaves an un-indexed + # reservation that the TTL reaps, never a double-spend. + await self._client.hset( + hash_key, + mapping={ + "id": reservation.id, + "leaseId": reservation.lease_id, + "companyId": reservation.company_id, + "creditTypeId": reservation.credit_type_id, + "eventSubtype": reservation.event_subtype, + "quantityReserved": format_amount(reservation.quantity_reserved), + "creditsReserved": format_amount(reservation.credits_reserved), + "consumptionRate": format_amount(reservation.consumption_rate), + "expiresAt": str(expires_ms), + "evalCtx": _encode_eval_ctx(reservation), + }, + ) + await self._client.pexpireat(hash_key, expires_ms + RES_TTL_GRACE_MS) + member = _encode_member(reservation.company_id, reservation.credit_type_id, reservation.id) + await self._client.zadd(self._index_key(), {member: expires_ms}) + await self._client.hset( + self._by_credit_key(reservation.company_id, reservation.credit_type_id), + reservation.id, + format_amount(reservation.credits_reserved), + ) + + async def get(self, reservation_id: str) -> Optional[ReservationRecord]: + raw = decode_hash(await self._client.hgetall(self._hash_key(reservation_id))) + if not raw.get("id"): + return None + return _decode_reservation(raw) + + async def consume(self, reservation_id: str, credits_consumed: float) -> Optional[float]: + claimed = await self._claim.run(self._client, [self._hash_key(reservation_id)], []) + raw = _decode_flat(claimed) + if not raw or not raw.get("id"): + return None + + company_id = raw["companyId"] + credit_type_id = raw["creditTypeId"] + reserved = float(raw.get("creditsReserved", "0")) + + # Index cleanup, single-key ops. The per-tenant hash loses the slice + # BEFORE the refund below, so the lease (local remaining plus this + # hash) never transiently double-counts it. + member = _encode_member(company_id, credit_type_id, reservation_id) + await _ignore_errors(self._client.zrem(self._index_key(), member)) + await _ignore_errors(self._client.hdel(self._by_credit_key(company_id, credit_type_id), reservation_id)) + + consumed = clamp_consumption(credits_consumed, reserved) + refund = reserved - consumed + if refund > 0: + # The lease store owns the lease hash, which keeps this cross-key + # write out of a single Lua script. Pinned to the reservation's + # lease so a hold carved out of an expired lease cannot inflate a + # successor's balance. + await self._lease_store.refund(company_id, credit_type_id, refund, raw.get("leaseId")) + return consumed + + async def reserved_credits(self, company_id: str, credit_type_id: str) -> float: + raw = await _ignore_errors(self._client.hgetall(self._by_credit_key(company_id, credit_type_id))) + total = 0.0 + for value in decode_hash(raw).values(): + try: + total += float(value) + except ValueError: + continue + return total + + async def sweep_expired(self, now: Optional[float] = None) -> int: + cutoff = int(to_epoch_ms(self._clock() if now is None else now)) + swept = 0 + # Page through expired members rather than fetching them all at once. + # Each processed member is removed below, so re-reading at offset 0 + # advances through the backlog. + for _ in range(MAX_SWEEP_BATCHES): + expired = await self._client.zrangebyscore( + self._index_key(), 0, cutoff, start=0, num=SWEEP_BATCH_SIZE + ) + if not expired: + break + for member in expired: + member_str = to_str(member) + decoded = _decode_member(member_str) + if decoded is None: + # Nothing but `add` writes members, so this is + # belt-and-braces: drop it rather than let it wedge the + # sweeper. + await _ignore_errors(self._client.zrem(self._index_key(), member_str)) + continue + company_id, credit_type_id, reservation_id = decoded + refunded = await self.consume(reservation_id, 0) + # Always drop the member just read. On the success path + # `consume` already removed it, so this is idempotent; it also + # covers the hash-evicted path below. + await _ignore_errors(self._client.zrem(self._index_key(), member_str)) + if refunded is not None: + swept += 1 + continue + # No reservation hash: either a racing settle consumed it (and + # reconciled the byCredit field, making this a no-op) or the + # hash TTL-evicted before the sweeper reached it, orphaning the + # field. Reconcile so reserved_credits stops summing an evicted + # hold. Deliberately no refund: without the hash, exactly-once + # cannot be arbitrated across racing sweepers, so the slice + # waits for the lease to expire server-side. + await _ignore_errors( + self._client.hdel(self._by_credit_key(company_id, credit_type_id), reservation_id) + ) + if len(expired) < SWEEP_BATCH_SIZE: + break + return swept + + async def count(self) -> int: + result = await _ignore_errors(self._client.zcard(self._index_key())) + return int(result or 0) + + +def _encode_member(company_id: str, credit_type_id: str, reservation_id: str) -> str: + """Expiry-index members carry the whole tuple. + + The sweeper needs company and credit to clean the per-tenant hash even + after the reservation hash has TTL-evicted, at which point the claim + returns nil and cannot report them; otherwise the orphaned field would + inflate reserved_credits forever. The delimiter is absent from Schematic + ids and from the reservation id. + """ + return f"{company_id}|{credit_type_id}|{reservation_id}" + + +def _decode_member(member: str) -> Optional[Tuple[str, str, str]]: + parts = member.split("|") + if len(parts) != 3: + return None + return parts[0], parts[1], parts[2] + + +def _encode_eval_ctx(reservation: ReservationRecord) -> str: + ctx: Dict[str, Any] = {} + if reservation.company is not None: + ctx["company"] = reservation.company + if reservation.user is not None: + ctx["user"] = reservation.user + # Compact, like the Node SDK writes it, so a shared row reads identically. + return json.dumps(ctx, separators=(",", ":")) + + +def _decode_flat(raw: Any) -> Dict[str, str]: + """Decode the flat [field, value, ...] reply the claim script returns.""" + if not raw or not isinstance(raw, (list, tuple)): + return {} + items: List[str] = [to_str(item) for item in raw] + return {items[i]: items[i + 1] for i in range(0, len(items) - 1, 2)} + + +def _decode_reservation(raw: Dict[str, str]) -> ReservationRecord: + ctx: Dict[str, Any] = {} + if raw.get("evalCtx"): + try: + ctx = json.loads(raw["evalCtx"]) + except ValueError: + ctx = {} + return ReservationRecord( + id=raw["id"], + lease_id=raw.get("leaseId", ""), + company_id=raw.get("companyId", ""), + credit_type_id=raw.get("creditTypeId", ""), + event_subtype=raw.get("eventSubtype", ""), + quantity_reserved=float(raw.get("quantityReserved", "0")), + credits_reserved=float(raw.get("creditsReserved", "0")), + consumption_rate=float(raw.get("consumptionRate", "0")), + expires_at=float(raw.get("expiresAt", "0")) / 1000.0, + company=ctx.get("company"), + user=ctx.get("user"), + ) + + +async def _ignore_errors(awaitable: Any) -> Any: + """Index bookkeeping is best-effort: a failed cleanup must not abort a settle.""" + try: + return await awaitable + except Exception: + return None diff --git a/src/schematic/leases/reservation_store.py b/src/schematic/leases/reservation_store.py new file mode 100644 index 00000000..121631d1 --- /dev/null +++ b/src/schematic/leases/reservation_store.py @@ -0,0 +1,125 @@ +"""Reservation table: the contract, plus the per-process in-memory backend. + +A reservation is a hold carved out of a lease. ``add`` does not debit: the +debit already landed in ``LeaseStore.try_reserve``, and that ordering is what +bounds a crash to a leaked hold rather than a double-spend. +""" + +from __future__ import annotations + +import abc +import math +import time +from typing import Dict, List, Optional + +from .lease_store import LeaseStore +from .types import Clock, ReservationRecord + + +class ReservationStore(abc.ABC): + """Backing store for open reservations, shared by both backends.""" + + @abc.abstractmethod + async def add(self, reservation: ReservationRecord) -> None: + """Register a reservation. Idempotent on id.""" + + @abc.abstractmethod + async def get(self, reservation_id: str) -> Optional[ReservationRecord]: + """Look up a reservation, or ``None`` once it has been claimed or swept.""" + + @abc.abstractmethod + async def consume(self, reservation_id: str, credits_consumed: float) -> Optional[float]: + """Claim a reservation exactly once and refund its unspent slice. + + The claim is atomic and comes first: a racing settle or sweep finds + nothing to claim, gets ``None``, and refunds nothing. On a successful + claim, ``credits_consumed`` is clamped to ``[0, credits_reserved]``, the + remainder is refunded to the lease (pinned to the reservation's lease), + and the clamped figure is returned. A crash between the claim and the + refund loses the refund, never double-refunds. + """ + + @abc.abstractmethod + async def reserved_credits(self, company_id: str, credit_type_id: str) -> float: + """Sum of ``credits_reserved`` across the slot's open reservations. + + A hold counts exactly while it is in the table, so + ``local_remaining_credits + reserved_credits`` stays exact between + operations. + """ + + @abc.abstractmethod + async def sweep_expired(self, now: Optional[float] = None) -> int: + """Remove every reservation past its TTL, refunding each full hold. + + Refunds are pinned to the originating lease, so a hold carved from a + lease that has since expired is dropped rather than credited to its + successor. Returns the number swept. + """ + + @abc.abstractmethod + async def count(self) -> int: + """Open reservations across every slot.""" + + +class InMemoryReservationStore(ReservationStore): + """Per-process reservation table refunding into a per-process lease store.""" + + def __init__(self, lease_store: LeaseStore, *, clock: Clock = time.time) -> None: + self._lease_store = lease_store + self._clock = clock + self._reservations: Dict[str, ReservationRecord] = {} + + async def add(self, reservation: ReservationRecord) -> None: + self._reservations[reservation.id] = reservation + + async def get(self, reservation_id: str) -> Optional[ReservationRecord]: + return self._reservations.get(reservation_id) + + async def consume(self, reservation_id: str, credits_consumed: float) -> Optional[float]: + # The claim: a dict pop with no await in it, so of two racing callers + # exactly one comes away with the record. + reservation = self._reservations.pop(reservation_id, None) + if reservation is None: + return None + consumed = clamp_consumption(credits_consumed, reservation.credits_reserved) + refund = reservation.credits_reserved - consumed + if refund > 0: + await self._lease_store.refund( + reservation.company_id, + reservation.credit_type_id, + refund, + reservation.lease_id, + ) + return consumed + + async def reserved_credits(self, company_id: str, credit_type_id: str) -> float: + return sum( + reservation.credits_reserved + for reservation in self._reservations.values() + if reservation.company_id == company_id and reservation.credit_type_id == credit_type_id + ) + + async def sweep_expired(self, now: Optional[float] = None) -> int: + cutoff = self._clock() if now is None else now + expired: List[str] = [ + reservation_id + for reservation_id, reservation in self._reservations.items() + if reservation.expires_at <= cutoff + ] + swept = 0 + for reservation_id in expired: + # Route through consume so the sweep claims exactly once too. + if await self.consume(reservation_id, 0) is not None: + swept += 1 + return swept + + async def count(self) -> int: + return len(self._reservations) + + +def clamp_consumption(credits_consumed: float, credits_reserved: float) -> float: + """Local bookkeeping never debits a lease past the hold it took.""" + if math.isnan(credits_consumed) or credits_consumed < 0: + return 0.0 + return min(credits_consumed, credits_reserved) diff --git a/src/schematic/leases/track.py b/src/schematic/leases/track.py new file mode 100644 index 00000000..76c68ab8 --- /dev/null +++ b/src/schematic/leases/track.py @@ -0,0 +1,82 @@ +"""Settling a reservation: consume the hold, then build the billing event. + +The event is built from the caller-held handle rather than the store, so the +usage is still billed when the hold has already been swept. The server is the +source of truth for real consumption; the local bookkeeping only keeps the +lease's view honest. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +from ..types import EventBodyTrack +from .reservation_store import ReservationStore + +if TYPE_CHECKING: + from ..client import Reservation, TrackWithReservationOptions + + +@dataclass +class ReservationConsumeResult: + """What a settle did locally, and what it owes the server.""" + + track: EventBodyTrack + # True when the hold was still open and this call debited the consumed + # slice and refunded the rest. False when it had already been swept at its + # TTL, already settled, or the store was unreachable: the lease balance was + # not touched here, so it reads high until the lease rolls over, and the + # track event is a recovery emit. + settled_locally: bool + + +async def consume_reservation_and_build_event( + reservations: ReservationStore, + reservation: "Reservation", + actual_quantity: float, + options: Optional["TrackWithReservationOptions"] = None, +) -> ReservationConsumeResult: + """Settle a hold against its lease and build the track event that bills it.""" + credits = actual_quantity * reservation.consumption_rate + consumed = await reservations.consume(reservation.id, credits) + return ReservationConsumeResult( + track=build_reservation_track_event(reservation, settled_quantity(actual_quantity), options), + settled_locally=consumed is not None, + ) + + +def build_reservation_track_event( + reservation: "Reservation", + actual_quantity: int, + options: Optional["TrackWithReservationOptions"] = None, +) -> EventBodyTrack: + """Build the track event that settles a reservation, from the handle alone. + + Kept free of store access so the client can still bill the usage when the + local settle fails against an unreachable store. + """ + return EventBodyTrack( + company=reservation.company, + event=reservation.event_subtype, + # A client-mode hold routes the server-side consumption through the + # lease's sub-ledger, instead of decrementing a grant the acquire + # already pre-debited. In server mode the hold lives on the server and + # settles by id; never send both, since the server prefers the lease id + # and there is no lease behind it. + lease_id=None if reservation.mode == "server" else reservation.lease_id, + quantity=actual_quantity, + reservation_id=reservation.id if reservation.mode == "server" else None, + traits=options.traits if options is not None else None, + user=reservation.user, + ) + + +def settled_quantity(actual_quantity: float) -> int: + """Cast a settled usage onto the integer a track event records. + + The hold can be sized from a fractional usage but the event's quantity is + an integer, so a partial unit settles as a whole one rather than as none. + """ + return int(actual_quantity) if float(actual_quantity).is_integer() else math.ceil(actual_quantity) diff --git a/src/schematic/leases/types.py b/src/schematic/leases/types.py new file mode 100644 index 00000000..4a9d8a89 --- /dev/null +++ b/src/schematic/leases/types.py @@ -0,0 +1,162 @@ +"""Shared types, defaults, and config resolution for client-mode credit leases. + +Durations are seconds (floats), like every other duration on this SDK, and +instants are epoch seconds (floats) rather than ``datetime`` objects: the +stores compare them against a clock the conformance runner can drive, and +Redis stores them as numbers anyway. The wire adapter converts at the +boundary. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, Mapping, Optional + +# Reads the current time as epoch seconds. Injected into every store and the +# lease manager so tests and the conformance runner can drive a virtual clock. +Clock = Callable[[], float] + +# Lease lifetime requested at acquire and extend (expires_at = now + duration). +DEFAULT_LEASE_DURATION = 300.0 +# Reservation lifetime, and the sweep deadline. Size it above the longest +# expected gap between check() and track_with_reservation(): a settle arriving +# after the TTL still bills the server but no longer re-debits the lease. +DEFAULT_RESERVATION_TTL = 60.0 +# Credits requested per acquire, and the minimum extend tranche. +DEFAULT_LEASE_SIZE = 10_000.0 +# Remaining/granted ratio at or below which a background extend is kicked off. +DEFAULT_LOW_WATER_MARK = 0.25 +# Expired-reservation sweep cadence. +DEFAULT_SWEEP_INTERVAL = 1.0 +# How long a prewarm waits for a freshly identified company to surface in the +# datastream cache before giving up. +DEFAULT_PREWARM_RESOLVE_TIMEOUT = 5.0 + + +@dataclass +class LeaseState: + """The local view of the one lease a ``(company, credit type)`` slot holds.""" + + lease_id: str + company_id: str + credit_type_id: str + # Server-authoritative total granted to this lease. Grows on extend. + granted_amount: float + # Granted minus outstanding holds and consumption. Starts at the full + # grant when the lease is installed. + local_remaining_credits: float + # Epoch seconds. Past it the lease is dead: the server has refunded the + # remainder to the company balance, so the local balance is stale. + expires_at: float + + +@dataclass +class ReservationRecord: + """One credit hold carved out of a lease by a check.""" + + id: str + # The lease the hold was carved from. Pins refunds so a hold from an + # expired lease can never inflate its successor. + lease_id: str + company_id: str + credit_type_id: str + # Event subtype the settling track event is billed as. + event_subtype: str + # Caller-declared usage, in event units. + quantity_reserved: float + # quantity_reserved * consumption_rate. + credits_reserved: float + consumption_rate: float + # Epoch seconds. The sweeper refunds the full hold past this instant. + expires_at: float + # Evaluation context the hold was issued for, threaded onto the track + # event so the server attributes usage to the same company and user. + company: Optional[Dict[str, str]] = None + user: Optional[Dict[str, str]] = None + + +@dataclass(frozen=True) +class ResolvedLeaseConfig: + """Config for a single credit type, after overrides and defaults.""" + + lease_duration: float = DEFAULT_LEASE_DURATION + reservation_ttl: float = DEFAULT_RESERVATION_TTL + lease_size: float = DEFAULT_LEASE_SIZE + low_water_mark: float = DEFAULT_LOW_WATER_MARK + + +@dataclass +class LeaseConfigOverride: + """Per-credit-type overrides of the four resolvable knobs.""" + + lease_duration: Optional[float] = None + reservation_ttl: Optional[float] = None + lease_size: Optional[float] = None + low_water_mark: Optional[float] = None + + +@dataclass +class LeaseConfig: + """Client-wide lease knobs, in seconds, plus per-credit-type overrides. + + The user-facing configuration dataclass lives on the client; this is the + plain-keyword form the lease machinery resolves against. + """ + + lease_duration: Optional[float] = None + reservation_ttl: Optional[float] = None + lease_size: Optional[float] = None + low_water_mark: Optional[float] = None + sweep_interval: Optional[float] = None + overrides: Mapping[str, LeaseConfigOverride] = field(default_factory=dict) + + +def resolve_lease_config( + config: Optional[LeaseConfig] = None, + overrides: Optional[Mapping[str, LeaseConfigOverride]] = None, + credit_type_id: Optional[str] = None, +) -> ResolvedLeaseConfig: + """Resolve the knobs for one credit type: override, then client config, then default. + + ``overrides`` wins over ``config.overrides`` so a caller can resolve against + a one-off override map without rebuilding the config. + """ + override: Optional[LeaseConfigOverride] = None + if credit_type_id is not None: + table = overrides if overrides is not None else (config.overrides if config else None) + if table: + override = table.get(credit_type_id) + + def pick(name: str, default: float) -> float: + if override is not None: + value = getattr(override, name) + if value is not None: + return float(value) + if config is not None: + value = getattr(config, name) + if value is not None: + return float(value) + return default + + return ResolvedLeaseConfig( + lease_duration=pick("lease_duration", DEFAULT_LEASE_DURATION), + # Uncapped: a client-mode TTL only tells the local sweeper when to + # refund an unsettled hold, and never reaches the server. + reservation_ttl=pick("reservation_ttl", DEFAULT_RESERVATION_TTL), + lease_size=pick("lease_size", DEFAULT_LEASE_SIZE), + low_water_mark=pick("low_water_mark", DEFAULT_LOW_WATER_MARK), + ) + + +def is_valid_quantity(value: Any) -> bool: + """Whether a caller-supplied quantity can size a credit hold. + + A bool is an int in Python, and NaN and infinity are floats that slip + through every numeric comparison, so a hold would be sized from any of them + with nothing rejecting it. A NaN debit is the worst of the three: it + poisons a possibly shared lease balance into approving every later reserve. + """ + if isinstance(value, bool) or not isinstance(value, (int, float)): + return False + return math.isfinite(value) and value >= 0 diff --git a/tests/conformance/test_vectors.py b/tests/conformance/test_vectors.py new file mode 100644 index 00000000..3c3693e1 --- /dev/null +++ b/tests/conformance/test_vectors.py @@ -0,0 +1,459 @@ +"""Runs the language-agnostic conformance vectors against this SDK. + +The vectors and the semantics they pin live in ``conformance/`` at the repo +root, copied verbatim from schematic-node (the reference implementation). This +runner is the only language-specific piece; every port reimplements it and must +pass the same vectors, on every store backend it ships. + +The flow-level ops (``check`` / ``track``) drive the real orchestration +against a scripted rules engine and a scripted DataStream, so what they pin is +what the SDK does around the engine, not the engine itself. +""" + +from __future__ import annotations + +import json +import logging +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional +from unittest import mock + +import pytest +from lease_support import ( + CrashingRefundLeaseStore, + ScriptedDataStream, + ScriptedEngine, + ScriptedWireClient, + VirtualClock, + make_fake_redis, +) + +from schematic.client import CheckOptions, Reservation +from schematic.leases import ( + CreditCheckDeps, + InMemoryLeaseStore, + InMemoryReservationStore, + LeaseConfig, + LeaseManager, + LeaseState, + LeaseStore, + RedisLeaseStore, + RedisReservationStore, + ReservationRecord, + ReservationStore, + check_with_lease, + consume_reservation_and_build_event, +) + +VECTORS_DIR = Path(__file__).resolve().parents[2] / "conformance" / "vectors" +BACKENDS = ("in_memory", "redis") +# What the vectors write for the balance the fail-open evaluation substitutes: +# the largest integer a JSON number carries exactly. +MAX_SAFE_INTEGER = 2**53 - 1 + + +def _load_cases() -> List[Any]: + cases: List[Any] = [] + for path in sorted(VECTORS_DIR.glob("*.json")): + document = json.loads(path.read_text()) + for vector in document["vectors"]: + for backend in BACKENDS: + allowed = vector.get("backends") + if allowed and backend not in allowed: + continue + cases.append( + pytest.param( + backend, + vector, + id=f"{backend}-{document['category']}-{vector['name']}", + ) + ) + return cases + + +class Harness: + """One vector's stores, manager, clock, and reservation handles.""" + + def __init__(self, backend: str, config: Dict[str, Any]) -> None: + self.clock = VirtualClock() + self.handles: Dict[str, Reservation] = {} + self.wire = ScriptedWireClient() + self.leases: LeaseStore + self.reservations: ReservationStore + if backend == "in_memory": + self.leases = InMemoryLeaseStore(clock=self.clock) + self.crash = CrashingRefundLeaseStore(self.leases) + self.reservations = InMemoryReservationStore(self.crash, clock=self.clock) + else: + client = make_fake_redis() + self.leases = RedisLeaseStore(client, clock=self.clock) + self.crash = CrashingRefundLeaseStore(self.leases) + self.reservations = RedisReservationStore(client, self.crash, clock=self.clock) + self.manager = LeaseManager( + self.wire, + self.leases, + reservation_store=self.reservations, + config=LeaseConfig( + lease_duration=_seconds(config.get("lease_duration_ms")), + reservation_ttl=_seconds(config.get("reservation_ttl_ms")), + lease_size=config.get("lease_size"), + low_water_mark=config.get("low_water_mark"), + ), + clock=self.clock, + ) + + def at_ms(self, offset_ms: float) -> float: + return self.clock.at_ms(offset_ms) + + def reservation_id(self, op: Dict[str, Any]) -> str: + if "handle" in op: + reservation = self.handles.get(op["handle"]) + if reservation is None: + raise AssertionError(f"unknown reservation handle: {op['handle']}") + return reservation.id + if "id" not in op: + raise AssertionError(f"op {op['op']} needs an id or handle") + return str(op["id"]) + + async def drain(self) -> None: + """Let the manager's fire-and-forget work (a redundant release) finish.""" + await self.manager._drain_background() + + +@pytest.mark.parametrize("backend,vector", _load_cases()) +async def test_vector(backend: str, vector: Dict[str, Any]) -> None: + harness = Harness(backend, (vector.get("given") or {}).get("config") or {}) + # The Redis backend decides expiry against the store's own clock (TIME), + # so the virtual clock has to be the process clock too, not just the one + # the stores read. + with mock.patch("time.time", harness.clock): + for lease in (vector.get("given") or {}).get("leases") or []: + written = await harness.leases.replace( + lease_id=lease["lease_id"], + company_id=lease["company_id"], + credit_type_id=lease["credit_type_id"], + granted_amount=lease["granted_amount"], + expires_at=harness.at_ms(lease["expires_at_ms"]), + ) + assert written is True + for op in vector["operations"]: + handler = _HANDLERS.get(op["op"]) + if handler is None: + raise AssertionError(f"unknown conformance op: {op['op']}") + await handler(harness, op, op.get("expect") or {}) + harness.manager.stop() + + +async def _op_advance_clock(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + h.clock.advance_ms(op.get("ms") or 0) + + +async def _op_replace_lease(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + written = await h.leases.replace( + lease_id=op["lease_id"], + company_id=op["company_id"], + credit_type_id=op["credit_type_id"], + granted_amount=op["granted_amount"], + expires_at=h.at_ms(op["expires_at_ms"]), + ) + if "written" in expect: + assert written is expect["written"] + + +async def _op_drop_lease(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + await h.leases.drop(op["company_id"], op["credit_type_id"]) + + +async def _op_try_reserve(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + balance = await h.leases.try_reserve(op["company_id"], op["credit_type_id"], op["credits"]) + if "balance" in expect: + _assert_number(balance, expect["balance"]) + + +async def _op_refund_lease(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + await h.leases.refund(op["company_id"], op["credit_type_id"], op["credits"], op.get("pin_lease_id")) + + +async def _op_extend_lease(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + expires_at = h.at_ms(op["expires_at_ms"]) if "expires_at_ms" in op else None + await h.leases.extend( + op["company_id"], + op["credit_type_id"], + op["granted_total"], + expires_at, + op.get("pin_lease_id"), + ) + + +async def _op_get_lease(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + entry = await h.leases.get(op["company_id"], op["credit_type_id"]) + if "exists" in expect: + assert (entry is not None) is expect["exists"] + if "lease_id" in expect: + assert (entry.lease_id if entry else None) == expect["lease_id"] + if "granted_amount" in expect: + assert entry is not None and entry.granted_amount == expect["granted_amount"] + if "local_remaining_credits" in expect: + assert entry is not None and entry.local_remaining_credits == expect["local_remaining_credits"] + + +async def _op_add_reservation(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + await h.reservations.add( + ReservationRecord( + id=op["id"], + lease_id=op["lease_id"], + company_id=op["company_id"], + credit_type_id=op["credit_type_id"], + event_subtype=op["event_subtype"], + quantity_reserved=op["quantity_reserved"], + credits_reserved=op["credits_reserved"], + consumption_rate=op["consumption_rate"], + expires_at=h.at_ms(op["expires_at_ms"]), + company={"id": op["company_id"]}, + ) + ) + + +async def _op_consume_reservation(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + reservation_id = h.reservation_id(op) + if op.get("crash_before_refund"): + h.crash.arm() + with pytest.raises(RuntimeError, match="simulated crash before refund"): + await h.reservations.consume(reservation_id, op["credits"]) + assert expect.get("throws") is True + return + consumed = await h.reservations.consume(reservation_id, op["credits"]) + if "consumed" in expect: + _assert_number(consumed, expect["consumed"]) + + +async def _op_get_reservation(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + reservation = await h.reservations.get(h.reservation_id(op)) + if "exists" in expect: + assert (reservation is not None) is expect["exists"] + + +async def _op_reserved_credits(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + total = await h.reservations.reserved_credits(op["company_id"], op["credit_type_id"]) + assert total == expect["total"] + + +async def _op_reservation_count(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + assert await h.reservations.count() == expect["count"] + + +async def _op_sweep_expired(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + swept = await h.reservations.sweep_expired() + if "swept" in expect: + assert swept == expect["swept"] + + +async def _op_acquire_if_needed(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + if op.get("server"): + h.wire.acquire_responses.append(_server_script(h, op["server"])) + install = op.get("install_during_wire") + if install: + + async def install_lease() -> None: + await h.leases.replace( + lease_id=install["lease_id"], + company_id=install["company_id"], + credit_type_id=install["credit_type_id"], + granted_amount=install["granted_amount"], + expires_at=h.at_ms(install["expires_at_ms"]), + ) + + h.wire.during_acquire = install_lease + entry = await h.manager.acquire_if_needed(op["company_id"], op["credit_type_id"]) + await h.drain() + if "lease_id" in expect: + assert (entry.lease_id if entry else None) == expect["lease_id"] + if "wire_acquires" in expect: + assert len(h.wire.acquire_calls) == expect["wire_acquires"] + if "last_acquire_requested_amount" in expect: + assert h.wire.acquire_calls[-1]["requested_amount"] == expect["last_acquire_requested_amount"] + if "released_lease_ids" in expect: + assert h.wire.release_calls == expect["released_lease_ids"] + + +async def _op_maybe_extend(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + if op.get("server"): + h.wire.extend_responses.append(_server_script(h, op["server"])) + await h.manager.maybe_extend(op["company_id"], op["credit_type_id"], op.get("required_credits")) + await h.drain() + if "wire_extends" in expect: + assert len(h.wire.extend_calls) == expect["wire_extends"] + if "last_extend_additional_amount" in expect: + assert h.wire.extend_calls[-1]["additional_amount"] == expect["last_extend_additional_amount"] + if "last_extend_lease_id" in expect: + assert h.wire.extend_calls[-1]["lease_id"] == expect["last_extend_lease_id"] + + +async def _op_release_all_local_leases(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + await h.manager.release_all_local_leases() + if "released_lease_ids" in expect: + assert h.wire.release_calls == expect["released_lease_ids"] + if "remaining_slots" in expect: + remaining: Optional[List[LeaseState]] = h.leases.list_leases() + assert len(remaining or []) == expect["remaining_slots"] + + +async def _op_check(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + flag_key = op.get("flag_key") or "flag" + company = op.get("company") or {"id": "co_1"} + engine = ScriptedEngine(op.get("engine") or [], flag_key) + datastream = ScriptedDataStream(engine, flag_key, company) + for call, script in (op.get("server") or {}).items(): + queue = h.wire.acquire_responses if call == "acquire" else h.wire.extend_responses + queue.append(_server_script(h, script)) + + fallback_called = False + + async def fallback() -> Any: + nonlocal fallback_called + fallback_called = True + from schematic.client import CheckResult + + return CheckResult(allowed=True, value=True, reason="fallback", flag_key=flag_key) + + async def enqueue_flag_check(body: Any) -> None: + """The vectors pin the lease flow, not analytics, so drop the event.""" + + deps = CreditCheckDeps( + datastream=datastream, + lease_store=h.leases, + reservations=h.reservations, + manager=h.manager, + logger=logging.getLogger("conformance"), + enqueue_flag_check=enqueue_flag_check, + clock=h.clock, + ) + options = CheckOptions( + usage=op.get("usage"), + event_subtype=op.get("event_subtype"), + on_acquire_failure=op.get("on_acquire_failure") or "fail-closed", + ) + result = await check_with_lease(deps, flag_key, {"id": company["id"]}, None, options, fallback) + await h.drain() + + if "allowed" in expect: + assert result.allowed is expect["allowed"] + if "reason" in expect: + assert result.reason == expect["reason"] + if "err" in expect: + assert result.error == expect["err"] + if "has_reservation" in expect: + assert (result.reservation is not None) is expect["has_reservation"] + if "fallback_called" in expect: + assert fallback_called is expect["fallback_called"] + if expect.get("reservation"): + reservation = result.reservation + assert reservation is not None + for field, attribute in ( + ("lease_id", "lease_id"), + ("credit_type_id", "credit_type_id"), + ("event_subtype", "event_subtype"), + ("quantity_reserved", "quantity_reserved"), + ("credits_reserved", "credits_reserved"), + ("consumption_rate", "consumption_rate"), + ): + if field in expect["reservation"]: + assert getattr(reservation, attribute) == expect["reservation"][field] + if "engine_calls" in expect: + _assert_engine_calls(engine.calls, expect["engine_calls"], _credit_id(op)) + if "wire_extends" in expect: + assert len(h.wire.extend_calls) == expect["wire_extends"] + if "last_extend_additional_amount" in expect: + assert h.wire.extend_calls[-1]["additional_amount"] == expect["last_extend_additional_amount"] + if op.get("save_reservation_as") and result.reservation is not None: + h.handles[op["save_reservation_as"]] = result.reservation + + +async def _op_track(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + reservation = h.handles.get(op["handle"]) + if reservation is None: + raise AssertionError(f"unknown reservation handle: {op['handle']}") + outcome = await consume_reservation_and_build_event(h.reservations, reservation, op["actual_quantity"]) + if "settled_locally" in expect: + assert outcome.settled_locally is expect["settled_locally"] + track = expect.get("track") or {} + if "event" in track: + assert outcome.track.event == track["event"] + if "quantity" in track: + assert outcome.track.quantity == track["quantity"] + if "lease_id" in track: + assert outcome.track.lease_id == track["lease_id"] + + +def _credit_id(op: Dict[str, Any]) -> Optional[str]: + """The credit the vector's engine_calls expectations are keyed on.""" + for scripted in op.get("engine") or []: + credit_id = (scripted.get("entitlement") or {}).get("credit_id") + if credit_id: + return str(credit_id) + balances = (op.get("company") or {}).get("credit_balances") or {} + return next(iter(balances), None) + + +def _assert_engine_calls( + recorded: List[Dict[str, Any]], expected: List[Dict[str, Any]], credit_id: Optional[str], +) -> None: + assert len(recorded) == len(expected) + for got, want in zip(recorded, expected): + if "credit_balance" in want: + balance = want["credit_balance"] + assert credit_id is not None + assert got["credit_balances"].get(credit_id) == ( + MAX_SAFE_INTEGER if balance == "max_safe_integer" else balance + ) + if "credit_cost" in want: + assert (got["credit_cost"] or {}).get(credit_id) == want["credit_cost"] + if "event_usage" in want: + assert got["event_usage"] == want["event_usage"] + if "usage" in want: + assert got["usage"] == want["usage"] + + +_Handler = Callable[[Harness, Dict[str, Any], Dict[str, Any]], Any] + +_HANDLERS: Dict[str, _Handler] = { + "advance_clock": _op_advance_clock, + "replace_lease": _op_replace_lease, + "drop_lease": _op_drop_lease, + "try_reserve": _op_try_reserve, + "refund_lease": _op_refund_lease, + "extend_lease": _op_extend_lease, + "get_lease": _op_get_lease, + "add_reservation": _op_add_reservation, + "consume_reservation": _op_consume_reservation, + "get_reservation": _op_get_reservation, + "reserved_credits": _op_reserved_credits, + "reservation_count": _op_reservation_count, + "sweep_expired": _op_sweep_expired, + "acquire_if_needed": _op_acquire_if_needed, + "maybe_extend": _op_maybe_extend, + "release_all_local_leases": _op_release_all_local_leases, + "check": _op_check, + "track": _op_track, +} + + +def _server_script(h: Harness, server: Dict[str, Any]) -> Dict[str, Any]: + """Turn a vector's server script into one the wire stand-in can serve.""" + if server.get("error") is not None: + return {"error": server["error"]} + lease: Dict[str, Any] = dict(server["lease"]) + lease["expires_at"] = h.at_ms(lease["expires_at_ms"]) + return {"lease": lease} + + +def _assert_number(actual: Optional[float], expected: Optional[float]) -> None: + """``None`` is the refused result, so it never compares equal to a figure.""" + if expected is None: + assert actual is None + return + assert actual is not None and actual == expected + + +def _seconds(milliseconds: Optional[float]) -> Optional[float]: + return None if milliseconds is None else milliseconds / 1000.0 diff --git a/tests/custom/test_client.py b/tests/custom/test_client.py index 668ffb4d..9f76216d 100644 --- a/tests/custom/test_client.py +++ b/tests/custom/test_client.py @@ -1,23 +1,50 @@ +import asyncio +import datetime as dt import time import unittest +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest from httpx import AsyncClient, Client +from lease_support import ScriptedDataStream, ScriptedEngine, make_fake_redis -from schematic.cache import LocalCache +from schematic.cache import LocalCache, RedisCache from schematic.client import ( + INSUFFICIENT_CREDITS_REASON, + MAX_RESERVATION_TTL, REASON_FLAG_NOT_FOUND, REASON_OFFLINE, + RESERVATION_TTL_SKEW_ALLOWANCE, AsyncSchematic, AsyncSchematicConfig, CheckFlagOptions, + CheckOptions, + CreditLeaseConfig, + DataStreamConfig, + EventUsage, IdentifyOptions, + Reservation, Schematic, SchematicConfig, TrackOptions, + TrackWithReservationOptions, + _is_valid_quantity, +) +from schematic.core.api_error import ApiError as CoreApiError +from schematic.errors import PaymentRequiredError +from schematic.leases import LeaseConfigOverride +from schematic.types import ( + ApiError, + CheckAndReserveFlagResponseData, + CheckFlagResponseData, + EventBodyIdentifyCompany, + FeatureEntitlement, + FlagCheckReservationResponseData, + PreflightEventUsageRequestBody, + PreflightRequestBody, + RulesengineCheckFlagResult, ) -from schematic.types import CheckFlagResponseData, FeatureEntitlement class TestSchematic(unittest.TestCase): @@ -1530,5 +1557,1427 @@ async def test_check_flags_datastream_failure_falls_back_to_bulk_api(self): await client.event_buffer.stop() +TTL_SECONDS = 120.0 + +CREDIT_ENTITLEMENT = FeatureEntitlement( + feature_id="feat", + feature_key="inference", + value_type="credit", +) + + +def _held_reservation(**overrides) -> FlagCheckReservationResponseData: + fields = dict( + id="rsv_1", + company_id="co_1", + credit_type_id="bilcr_inference", + consumption_rate=10.0, + credits_reserved=500.0, + quantity_reserved=50.0, + event_subtype="inference_tokens", + expires_at=dt.datetime.now(dt.timezone.utc) + dt.timedelta(seconds=TTL_SECONDS), + ) + fields.update(overrides) + return FlagCheckReservationResponseData(**fields) # type: ignore[arg-type] + + +def _reserve_response(**overrides): + fields = dict( + flag="inference", + flag_id="flag_1", + value=True, + reason="matched", + company_id="co_1", + user_id="user_1", + rule_id="rule_1", + entitlement=CREDIT_ENTITLEMENT, + reservation=_held_reservation(), + ) + fields.update(overrides) + return MagicMock(data=CheckAndReserveFlagResponseData(**fields)) # type: ignore[arg-type] + + +class TestSchematicPreflight(unittest.TestCase): + """Preflight options on the sync REST check path.""" + + def setUp(self): + config = SchematicConfig( + event_buffer_period=1, + logger=MagicMock(), + httpx_client=MagicMock(spec=Client), + ) + self.schematic = Schematic("api_key", config) + self.data = CheckFlagResponseData(value=True, flag="inference", reason="matched") + self.schematic.features.check_flag = MagicMock(return_value=MagicMock(data=self.data)) + + def tearDown(self): + self.schematic.event_buffer.stop() + + def test_plain_check_sends_no_preflight_kwarg(self): + self.schematic.check_flag("inference", company={"id": "co_1"}) + self.assertNotIn("preflight", self.schematic.features.check_flag.call_args.kwargs) + + def test_usage_is_forwarded_as_preflight(self): + self.schematic.check_flag("inference", company={"id": "co_1"}, options=CheckFlagOptions(usage=5)) + preflight = self.schematic.features.check_flag.call_args.kwargs["preflight"] + self.assertEqual(preflight, PreflightRequestBody(usage=5)) + + def test_event_usage_and_credit_cost_are_forwarded_as_preflight(self): + self.schematic.check_flag( + "inference", + company={"id": "co_1"}, + options=CheckFlagOptions( + event_usage=EventUsage(event_subtype="inference_tokens", quantity=7), + credit_cost={"bilcr_inference": 12.5}, + ), + ) + preflight = self.schematic.features.check_flag.call_args.kwargs["preflight"] + self.assertEqual( + preflight, + PreflightRequestBody( + credit_cost={"bilcr_inference": 12.5}, + event_usage=PreflightEventUsageRequestBody(event_subtype="inference_tokens", quantity=7), + ), + ) + + def test_preflighted_check_neither_reads_nor_writes_the_cache(self): + company = {"id": "co_1"} + options = CheckFlagOptions(usage=5) + + # Two preflighted checks both hit the API: the answer is specific to + # the simulated usage, so it is never served from the cache. + self.schematic.check_flag("inference", company=company, options=options) + self.schematic.check_flag("inference", company=company, options=options) + self.assertEqual(self.schematic.features.check_flag.call_count, 2) + + # And nothing they returned was written to the cache: the first plain + # check still has to ask the API. + self.schematic.check_flag("inference", company=company) + self.assertEqual(self.schematic.features.check_flag.call_count, 3) + + def test_plain_check_still_caches(self): + company = {"id": "co_1"} + self.schematic.check_flag("inference", company=company) + self.schematic.check_flag("inference", company=company) + self.assertEqual(self.schematic.features.check_flag.call_count, 1) + + +class TestQuantityValidation(unittest.TestCase): + """What a usage, or a settled quantity, has to be to size a credit hold.""" + + def test_accepts_finite_non_negative_numbers(self): + for value in (0, 50, 100.0, 0.5): + with self.subTest(value=value): + self.assertTrue(_is_valid_quantity(value)) + + def test_rejects_bools_negatives_and_non_finite_floats(self): + for value in (True, -1, float("nan"), float("inf")): + with self.subTest(value=value): + self.assertFalse(_is_valid_quantity(value)) + + +class TestSchematicServerReservation(unittest.TestCase): + """check() and track_with_reservation() against the server hold path.""" + + def setUp(self): + self.schematic = self._client() + + def tearDown(self): + self.schematic.event_buffer.stop() + + def _client(self, **config_overrides) -> Schematic: + config_kwargs = dict( + event_buffer_period=1, + logger=MagicMock(), + httpx_client=MagicMock(spec=Client), + credit_leases=CreditLeaseConfig(mode="server", default_reservation_ttl=TTL_SECONDS), + ) + config_kwargs.update(config_overrides) + client = Schematic("api_key", SchematicConfig(**config_kwargs)) # type: ignore[arg-type] + client.features.check_and_reserve_flag = MagicMock(return_value=_reserve_response()) + client.features.check_flag = MagicMock( + return_value=MagicMock(data=CheckFlagResponseData(value=True, flag="inference", reason="plain check")) + ) + client.credits.release_credit_reservation = MagicMock() + client.flag_check_cache_providers = [] + return client + + def test_returns_a_reservation_handle_built_from_the_response(self): + before = dt.datetime.now(dt.timezone.utc) + with patch.object(self.schematic.event_buffer, "push") as mock_push: + result = self.schematic.check( + "inference", + company={"id": "co_1"}, + user={"id": "user_1"}, + options=CheckOptions(usage=50, event_subtype="inference_tokens"), + ) + after = dt.datetime.now(dt.timezone.utc) + + self.assertTrue(result.allowed) + self.assertTrue(result.value) + self.assertEqual(result.reason, "matched") + self.assertEqual(result.flag_key, "inference") + self.assertEqual(result.flag_id, "flag_1") + self.assertEqual(result.entitlement, CREDIT_ENTITLEMENT) + + assert result.reservation is not None + self.assertEqual(result.reservation.id, "rsv_1") + # No lease exists server side; the handle mirrors the id. + self.assertEqual(result.reservation.lease_id, "rsv_1") + self.assertEqual(result.reservation.mode, "server") + self.assertEqual(result.reservation.company_id, "co_1") + self.assertEqual(result.reservation.credit_type_id, "bilcr_inference") + self.assertEqual(result.reservation.event_subtype, "inference_tokens") + self.assertEqual(result.reservation.quantity_reserved, 50.0) + self.assertEqual(result.reservation.credits_reserved, 500.0) + self.assertEqual(result.reservation.consumption_rate, 10.0) + self.assertEqual(result.reservation.company, {"id": "co_1"}) + self.assertEqual(result.reservation.user, {"id": "user_1"}) + + kwargs = self.schematic.features.check_and_reserve_flag.call_args.kwargs + self.assertEqual(self.schematic.features.check_and_reserve_flag.call_args.args, ("inference",)) + self.assertEqual(kwargs["quantity"], 50) + self.assertEqual(kwargs["company"], {"id": "co_1"}) + self.assertEqual(kwargs["user"], {"id": "user_1"}) + self.assertEqual( + kwargs["preflight"], + PreflightRequestBody( + event_usage=PreflightEventUsageRequestBody(event_subtype="inference_tokens", quantity=50) + ), + ) + ttl = dt.timedelta(seconds=TTL_SECONDS) + self.assertGreaterEqual(kwargs["expires_at"], before + ttl) + self.assertLessEqual(kwargs["expires_at"], after + ttl) + self.assertEqual(kwargs["request_options"], {"max_retries": 0}) + + # The server logs the flag check for check-and-reserve itself. + mock_push.assert_not_called() + + def test_sends_the_generic_usage_preflight_without_an_event_subtype(self): + self.schematic.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50)) + kwargs = self.schematic.features.check_and_reserve_flag.call_args.kwargs + self.assertEqual(kwargs["preflight"], PreflightRequestBody(usage=50)) + + def test_forwards_the_per_check_timeout(self): + self.schematic.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50, timeout=2.5)) + kwargs = self.schematic.features.check_and_reserve_flag.call_args.kwargs + self.assertEqual(kwargs["request_options"], {"max_retries": 0, "timeout": 2.5}) + + def test_never_retries_check_and_reserve(self): + # The call has no idempotency key, so a retried 5xx that the server + # already committed would take a second hold. + self.schematic.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50)) + kwargs = self.schematic.features.check_and_reserve_flag.call_args.kwargs + self.assertEqual(kwargs["request_options"]["max_retries"], 0) + + def test_a_fractional_usage_sizes_the_hold_and_rounds_the_preflight_up(self): + self.schematic.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=0.5)) + kwargs = self.schematic.features.check_and_reserve_flag.call_args.kwargs + self.assertEqual(kwargs["quantity"], 0.5) + # The hold takes the fraction; the preflight's usage is an integer, and + # rounding it down would ask about less usage than is about to land. + self.assertEqual(kwargs["preflight"], PreflightRequestBody(usage=1)) + + def test_an_integral_float_usage_reaches_the_preflight_unchanged(self): + self.schematic.check( + "inference", + company={"id": "co_1"}, + options=CheckOptions(usage=100.0, event_subtype="inference_tokens"), + ) + kwargs = self.schematic.features.check_and_reserve_flag.call_args.kwargs + self.assertEqual(kwargs["quantity"], 100.0) + self.assertEqual( + kwargs["preflight"], + PreflightRequestBody( + event_usage=PreflightEventUsageRequestBody(event_subtype="inference_tokens", quantity=100) + ), + ) + + def test_a_reservation_ttl_above_the_cap_is_clamped(self): + client = self._client(credit_leases=CreditLeaseConfig(default_reservation_ttl=7200.0)) + try: + # Short of the cap by the skew allowance, so a client running + # slightly fast still asks for something the server accepts. + self.assertEqual(client._reservation_ttl, MAX_RESERVATION_TTL - RESERVATION_TTL_SKEW_ALLOWANCE) + warning = " ".join(str(call.args[0]) for call in client.logger.warning.call_args_list) + self.assertIn("one hour cap", warning) + self.assertIn( + f"server-mode holds will be clamped to {MAX_RESERVATION_TTL - RESERVATION_TTL_SKEW_ALLOWANCE}s", + warning, + ) + finally: + client.event_buffer.stop() + + def test_denies_without_a_reservation_when_credits_are_short(self): + self.schematic.features.check_and_reserve_flag.return_value = _reserve_response( + value=False, reason="Insufficient credits", reservation=None, + ) + result = self.schematic.check( + "inference", company={"id": "co_1"}, options=CheckOptions(usage=50, event_subtype="inference_tokens"), + ) + self.assertFalse(result.allowed) + self.assertFalse(result.value) + self.assertEqual(result.reason, "Insufficient credits") + self.assertIsNone(result.reservation) + self.schematic.credits.release_credit_reservation.assert_not_called() + + def test_allows_without_a_reservation_when_the_feature_is_not_credit_metered(self): + self.schematic.features.check_and_reserve_flag.return_value = _reserve_response( + reason="company entitlement", + reservation=None, + entitlement=FeatureEntitlement(feature_id="feat", feature_key="inference", value_type="boolean"), + ) + result = self.schematic.check( + "inference", company={"id": "co_1"}, options=CheckOptions(usage=50, event_subtype="inference_tokens"), + ) + self.assertTrue(result.allowed) + self.assertIsNone(result.reservation) + self.assertEqual(result.reason, "company entitlement") + + def test_payment_required_denies_even_with_fail_open(self): + self.schematic.features.check_and_reserve_flag.side_effect = PaymentRequiredError( + body=ApiError(error="credit balance exhausted") + ) + result = self.schematic.check( + "inference", + company={"id": "co_1"}, + options=CheckOptions(usage=50, on_acquire_failure="fail-open", default_value=True), + ) + self.assertFalse(result.allowed) + self.assertFalse(result.value) + self.assertEqual(result.reason, INSUFFICIENT_CREDITS_REASON) + self.assertEqual(result.error, "credit balance exhausted") + self.assertIsNone(result.reservation) + + def test_a_402_api_error_denies_even_with_fail_open(self): + # The generated features client has no 402 branch, so a real 402 from + # check-and-reserve arrives as the base ApiError. + self.schematic.features.check_and_reserve_flag.side_effect = CoreApiError( + status_code=402, body={"error": "credit balance exhausted"}, + ) + result = self.schematic.check( + "inference", + company={"id": "co_1"}, + options=CheckOptions(usage=50, on_acquire_failure="fail-open", default_value=True), + ) + self.assertFalse(result.allowed) + self.assertFalse(result.value) + self.assertEqual(result.reason, INSUFFICIENT_CREDITS_REASON) + self.assertEqual(result.error, "credit balance exhausted") + self.assertIsNone(result.reservation) + + def test_fails_closed_when_check_and_reserve_errors(self): + self.schematic.features.check_and_reserve_flag.side_effect = Exception("ECONNRESET") + result = self.schematic.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50)) + self.assertFalse(result.allowed) + self.assertFalse(result.value) + self.assertEqual(result.reason, "server_reservation_failed") + self.assertEqual(result.error, "server_reservation_failed") + self.assertIsNone(result.reservation) + + def test_fails_open_to_the_per_check_default_value(self): + self.schematic.features.check_and_reserve_flag.side_effect = Exception("ECONNRESET") + result = self.schematic.check( + "inference", + company={"id": "co_1"}, + options=CheckOptions(usage=50, on_acquire_failure="fail-open", default_value=True), + ) + self.assertTrue(result.allowed) + self.assertTrue(result.value) + self.assertEqual(result.reason, "server_reservation_failed_fail_open") + self.assertEqual(result.error, "server_reservation_failed") + + def test_fails_open_to_the_client_level_flag_default(self): + client = self._client(flag_defaults={"inference": True}) + try: + client.features.check_and_reserve_flag.side_effect = Exception("ECONNRESET") + result = client.check( + "inference", company={"id": "co_1"}, options=CheckOptions(usage=50, on_acquire_failure="fail-open"), + ) + self.assertTrue(result.allowed) + self.assertEqual(result.reason, "server_reservation_failed_fail_open") + finally: + client.event_buffer.stop() + + # The same client with no configured default stays denied. + self.schematic.features.check_and_reserve_flag.side_effect = Exception("ECONNRESET") + denied = self.schematic.check( + "inference", company={"id": "co_1"}, options=CheckOptions(usage=50, on_acquire_failure="fail-open"), + ) + self.assertFalse(denied.allowed) + + def test_zero_usage_falls_back_to_a_plain_check(self): + result = self.schematic.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=0)) + self.schematic.features.check_and_reserve_flag.assert_not_called() + self.schematic.features.check_flag.assert_called_once() + self.assertTrue(result.allowed) + self.assertEqual(result.reason, "plain check") + self.assertIsNone(result.reservation) + + def test_invalid_usage_resolves_through_the_failure_contract(self): + for usage in (-5, float("nan"), float("inf"), True): + with self.subTest(usage=usage): + denied = self.schematic.check( + "inference", company={"id": "co_1"}, options=CheckOptions(usage=usage), # type: ignore[arg-type] + ) + self.assertFalse(denied.allowed) + self.assertEqual(denied.reason, "invalid_usage") + self.assertEqual(denied.error, "invalid_usage") + + opened = self.schematic.check( + "inference", + company={"id": "co_1"}, + options=CheckOptions(usage=-5, on_acquire_failure="fail-open", default_value=True), + ) + self.assertTrue(opened.allowed) + self.assertEqual(opened.reason, "invalid_usage_fail_open") + + self.schematic.features.check_and_reserve_flag.assert_not_called() + self.schematic.features.check_flag.assert_not_called() + + def test_releases_a_hold_that_names_no_event_subtype(self): + self.schematic.features.check_and_reserve_flag.return_value = _reserve_response( + reservation=_held_reservation(id="rsv_orphan", event_subtype=None), + ) + result = self.schematic.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50)) + self.schematic.credits.release_credit_reservation.assert_called_once_with("rsv_orphan") + self.assertFalse(result.allowed) + self.assertEqual(result.reason, "missing_event_subtype") + self.assertIsNone(result.reservation) + + def test_a_failed_release_is_swallowed(self): + self.schematic.features.check_and_reserve_flag.return_value = _reserve_response( + reservation=_held_reservation(id="rsv_orphan", event_subtype=None), + ) + self.schematic.credits.release_credit_reservation.side_effect = Exception("boom") + result = self.schematic.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50)) + self.assertEqual(result.reason, "missing_event_subtype") + self.schematic.logger.warning.assert_called() + + def test_a_released_hold_keeps_the_server_verdict_when_failing_open(self): + self.schematic.features.check_and_reserve_flag.return_value = _reserve_response( + reservation=_held_reservation(id="rsv_orphan", event_subtype=None), + ) + result = self.schematic.check( + "inference", company={"id": "co_1"}, options=CheckOptions(usage=50, on_acquire_failure="fail-open"), + ) + self.schematic.credits.release_credit_reservation.assert_called_once_with("rsv_orphan") + # The server evaluated the flag and allowed it; only the settle is + # impossible, and fail-open assumes the credits are there. + self.assertTrue(result.allowed) + self.assertTrue(result.value) + self.assertEqual(result.reason, "matched") + self.assertEqual(result.flag_id, "flag_1") + self.assertEqual(result.entitlement, CREDIT_ENTITLEMENT) + self.assertEqual(result.error, "missing_event_subtype") + self.assertIsNone(result.reservation) + + def test_no_credit_lease_config_falls_back_to_a_plain_check(self): + client = self._client(credit_leases=None) + try: + result = client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50)) + client.features.check_and_reserve_flag.assert_not_called() + client.features.check_flag.assert_called_once() + self.assertIsNone(result.reservation) + # The preflight still rides along on the plain check. + self.assertEqual( + client.features.check_flag.call_args.kwargs["preflight"], PreflightRequestBody(usage=50), + ) + finally: + client.event_buffer.stop() + + def test_client_mode_falls_back_and_warns_at_construction(self): + client = self._client(credit_leases=CreditLeaseConfig(mode="client")) + try: + warning = " ".join(str(call.args[0]) for call in client.logger.warning.call_args_list) + self.assertIn("'client'", warning) + result = client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50)) + client.features.check_and_reserve_flag.assert_not_called() + client.features.check_flag.assert_called_once() + self.assertIsNone(result.reservation) + finally: + client.event_buffer.stop() + + def test_offline_check_returns_the_flag_default(self): + client = self._client(offline=True, flag_defaults={"inference": True}) + try: + result = client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50)) + client.features.check_and_reserve_flag.assert_not_called() + self.assertTrue(result.allowed) + self.assertEqual(result.reason, REASON_OFFLINE) + self.assertIsNone(result.reservation) + finally: + client.event_buffer.stop() + + def _reservation_handle(self) -> Reservation: + result = self.schematic.check( + "inference", + company={"id": "co_1"}, + user={"id": "user_1"}, + options=CheckOptions(usage=50, event_subtype="inference_tokens"), + ) + assert result.reservation is not None + return result.reservation + + def test_track_with_reservation_settles_by_reservation_id(self): + reservation = self._reservation_handle() + with patch.object(self.schematic.event_buffer, "push") as mock_push: + self.schematic.track_with_reservation( + reservation, 20, TrackWithReservationOptions(traits={"model": "opus"}), + ) + + pushed = mock_push.call_args.args[0] + self.assertEqual(pushed.event_type, "track") + self.assertEqual(pushed.body.event, "inference_tokens") + self.assertEqual(pushed.body.quantity, 20) + self.assertEqual(pushed.body.reservation_id, "rsv_1") + # The server prefers lease_id when both are set, and there is no lease. + self.assertIsNone(pushed.body.lease_id) + self.assertEqual(pushed.body.company, {"id": "co_1"}) + self.assertEqual(pushed.body.user, {"id": "user_1"}) + self.assertEqual(pushed.body.traits, {"model": "opus"}) + self.assertEqual(pushed.idempotency_key, "lease-reservation:rsv_1") + self.schematic.credits.release_credit_reservation.assert_not_called() + + def test_track_with_reservation_skips_an_invalid_quantity(self): + reservation = self._reservation_handle() + with patch.object(self.schematic.event_buffer, "push") as mock_push: + for quantity in (-1, float("nan"), float("inf"), True): + self.schematic.track_with_reservation(reservation, quantity) # type: ignore[arg-type] + mock_push.assert_not_called() + + def test_track_with_reservation_settles_a_fractional_quantity_as_a_whole_unit(self): + reservation = self._reservation_handle() + with patch.object(self.schematic.event_buffer, "push") as mock_push: + self.schematic.track_with_reservation(reservation, 0.5) + # A track event's quantity is an integer, so a partial unit bills as one. + self.assertEqual(mock_push.call_args.args[0].body.quantity, 1) + + def test_track_with_reservation_without_a_hold_says_to_use_track(self): + with patch.object(self.schematic.event_buffer, "push") as mock_push: + self.schematic.track_with_reservation(None, 5) + mock_push.assert_not_called() + self.assertIn("track()", str(self.schematic.logger.error.call_args.args[0])) + + def test_track_with_reservation_is_a_no_op_when_offline(self): + reservation = self._reservation_handle() + self.schematic.offline = True + with patch.object(self.schematic.event_buffer, "push") as mock_push: + self.schematic.track_with_reservation(reservation, 20) + mock_push.assert_not_called() + + +def _async_server_client(**config_overrides) -> AsyncSchematic: + config_kwargs = dict( + event_buffer_period=1, + logger=MagicMock(), + httpx_client=MagicMock(spec=AsyncClient), + credit_leases=CreditLeaseConfig(mode="server", default_reservation_ttl=TTL_SECONDS), + ) + config_kwargs.update(config_overrides) + client = AsyncSchematic("test_key", AsyncSchematicConfig(**config_kwargs)) # type: ignore[arg-type] + client.features.check_and_reserve_flag = AsyncMock(return_value=_reserve_response()) + client.features.check_flag = AsyncMock( + return_value=MagicMock(data=CheckFlagResponseData(value=True, flag="inference", reason="plain check")) + ) + client.credits.release_credit_reservation = AsyncMock() + client.flag_check_cache_providers = [] + return client + + +@pytest.mark.asyncio +class TestAsyncSchematicPreflight: + """Preflight options on the async check paths.""" + + @pytest.fixture(autouse=True) + async def setup_and_teardown(self): + config = AsyncSchematicConfig( + logger=MagicMock(), + httpx_client=MagicMock(spec=AsyncClient), + event_buffer_period=1, + ) + self.client = AsyncSchematic("test_key", config) + self.client.features.check_flag = AsyncMock( + return_value=MagicMock(data=CheckFlagResponseData(value=True, flag="inference", reason="matched")) + ) + yield + await self.client.event_buffer.stop() + + async def test_plain_check_sends_no_preflight_kwarg(self): + await self.client.check_flag("inference", company={"id": "co_1"}) + assert "preflight" not in self.client.features.check_flag.call_args.kwargs + + async def test_usage_is_forwarded_as_preflight(self): + await self.client.check_flag("inference", company={"id": "co_1"}, options=CheckFlagOptions(usage=5)) + assert self.client.features.check_flag.call_args.kwargs["preflight"] == PreflightRequestBody(usage=5) + + async def test_preflighted_check_neither_reads_nor_writes_the_cache(self): + company = {"id": "co_1"} + options = CheckFlagOptions(usage=5) + + await self.client.check_flag("inference", company=company, options=options) + await self.client.check_flag("inference", company=company, options=options) + assert self.client.features.check_flag.call_count == 2 + + await self.client.check_flag("inference", company=company) + assert self.client.features.check_flag.call_count == 3 + + async def test_plain_check_still_caches(self): + company = {"id": "co_1"} + await self.client.check_flag("inference", company=company) + await self.client.check_flag("inference", company=company) + assert self.client.features.check_flag.call_count == 1 + + async def test_datastream_check_receives_the_options(self): + ds_result = RulesengineCheckFlagResult(value=True, flag_key="inference", reason="matched") + mock_ds = MagicMock() + mock_ds.check_flag = AsyncMock(return_value=ds_result) + self.client._datastream_client = mock_ds + + options = CheckFlagOptions(usage=5) + await self.client.check_flag("inference", company={"id": "co_1"}, options=options) + + assert mock_ds.check_flag.call_args.kwargs["options"] is options + self.client.features.check_flag.assert_not_called() + + async def test_check_threads_its_preflight_through_the_datastream_fallback(self): + ds_result = RulesengineCheckFlagResult(value=True, flag_key="inference", reason="matched") + mock_ds = MagicMock() + mock_ds.check_flag = AsyncMock(return_value=ds_result) + self.client._datastream_client = mock_ds + + # No credit leases configured, so check() is a plain check that still + # carries the caller's preflight. + result = await self.client.check( + "inference", company={"id": "co_1"}, options=CheckOptions(usage=5, event_subtype="inference_tokens"), + ) + + assert result.allowed is True + assert result.reservation is None + threaded = mock_ds.check_flag.call_args.kwargs["options"] + assert threaded.event_usage == EventUsage(event_subtype="inference_tokens", quantity=5) + + +@pytest.mark.asyncio +class TestAsyncSchematicServerReservation: + """check() and track_with_reservation() on the async client.""" + + @pytest.fixture(autouse=True) + async def setup_and_teardown(self): + self.client = _async_server_client() + yield + await self.client.event_buffer.stop() + + async def test_returns_a_reservation_handle_built_from_the_response(self): + before = dt.datetime.now(dt.timezone.utc) + with patch.object(self.client.event_buffer, "push", new=AsyncMock()) as mock_push: + result = await self.client.check( + "inference", + company={"id": "co_1"}, + user={"id": "user_1"}, + options=CheckOptions(usage=50, event_subtype="inference_tokens"), + ) + after = dt.datetime.now(dt.timezone.utc) + + assert result.allowed is True + assert result.value is True + assert result.reason == "matched" + assert result.flag_key == "inference" + assert result.flag_id == "flag_1" + assert result.entitlement == CREDIT_ENTITLEMENT + + assert result.reservation is not None + assert result.reservation.id == "rsv_1" + assert result.reservation.lease_id == "rsv_1" + assert result.reservation.mode == "server" + assert result.reservation.company_id == "co_1" + assert result.reservation.credit_type_id == "bilcr_inference" + assert result.reservation.event_subtype == "inference_tokens" + assert result.reservation.quantity_reserved == 50.0 + assert result.reservation.credits_reserved == 500.0 + assert result.reservation.consumption_rate == 10.0 + assert result.reservation.company == {"id": "co_1"} + assert result.reservation.user == {"id": "user_1"} + + kwargs = self.client.features.check_and_reserve_flag.call_args.kwargs + assert kwargs["quantity"] == 50 + assert kwargs["company"] == {"id": "co_1"} + assert kwargs["user"] == {"id": "user_1"} + assert kwargs["preflight"] == PreflightRequestBody( + event_usage=PreflightEventUsageRequestBody(event_subtype="inference_tokens", quantity=50) + ) + ttl = dt.timedelta(seconds=TTL_SECONDS) + assert before + ttl <= kwargs["expires_at"] <= after + ttl + + # The server logs the flag check for check-and-reserve itself. + mock_push.assert_not_called() + + async def test_sends_the_generic_usage_preflight_without_an_event_subtype(self): + await self.client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50)) + kwargs = self.client.features.check_and_reserve_flag.call_args.kwargs + assert kwargs["preflight"] == PreflightRequestBody(usage=50) + + async def test_forwards_the_per_check_timeout(self): + await self.client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50, timeout=2.5)) + kwargs = self.client.features.check_and_reserve_flag.call_args.kwargs + assert kwargs["request_options"] == {"max_retries": 0, "timeout": 2.5} + + async def test_never_retries_check_and_reserve(self): + # The call has no idempotency key, so a retried 5xx that the server + # already committed would take a second hold. + await self.client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50)) + kwargs = self.client.features.check_and_reserve_flag.call_args.kwargs + assert kwargs["request_options"] == {"max_retries": 0} + + async def test_a_fractional_usage_sizes_the_hold_and_rounds_the_preflight_up(self): + await self.client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=0.5)) + kwargs = self.client.features.check_and_reserve_flag.call_args.kwargs + assert kwargs["quantity"] == 0.5 + assert kwargs["preflight"] == PreflightRequestBody(usage=1) + + async def test_a_reservation_ttl_above_the_cap_is_clamped(self): + client = _async_server_client(credit_leases=CreditLeaseConfig(default_reservation_ttl=7200.0)) + try: + assert client._reservation_ttl == MAX_RESERVATION_TTL - RESERVATION_TTL_SKEW_ALLOWANCE + warning = " ".join(str(call.args[0]) for call in client.logger.warning.call_args_list) + assert "one hour cap" in warning + assert ( + f"server-mode holds will be clamped to {MAX_RESERVATION_TTL - RESERVATION_TTL_SKEW_ALLOWANCE}s" + in warning + ) + finally: + await client.event_buffer.stop() + + async def test_denies_without_a_reservation_when_credits_are_short(self): + self.client.features.check_and_reserve_flag.return_value = _reserve_response( + value=False, reason="Insufficient credits", reservation=None, + ) + result = await self.client.check( + "inference", company={"id": "co_1"}, options=CheckOptions(usage=50, event_subtype="inference_tokens"), + ) + assert result.allowed is False + assert result.reason == "Insufficient credits" + assert result.reservation is None + self.client.credits.release_credit_reservation.assert_not_called() + + async def test_payment_required_denies_even_with_fail_open(self): + self.client.features.check_and_reserve_flag.side_effect = PaymentRequiredError( + body=ApiError(error="credit balance exhausted") + ) + result = await self.client.check( + "inference", + company={"id": "co_1"}, + options=CheckOptions(usage=50, on_acquire_failure="fail-open", default_value=True), + ) + assert result.allowed is False + assert result.reason == INSUFFICIENT_CREDITS_REASON + assert result.error == "credit balance exhausted" + + async def test_a_402_api_error_denies_even_with_fail_open(self): + # The generated features client has no 402 branch, so a real 402 from + # check-and-reserve arrives as the base ApiError. + self.client.features.check_and_reserve_flag.side_effect = CoreApiError( + status_code=402, body={"error": "credit balance exhausted"}, + ) + result = await self.client.check( + "inference", + company={"id": "co_1"}, + options=CheckOptions(usage=50, on_acquire_failure="fail-open", default_value=True), + ) + assert result.allowed is False + assert result.value is False + assert result.reason == INSUFFICIENT_CREDITS_REASON + assert result.error == "credit balance exhausted" + assert result.reservation is None + + async def test_fails_closed_when_check_and_reserve_errors(self): + self.client.features.check_and_reserve_flag.side_effect = Exception("ECONNRESET") + result = await self.client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50)) + assert result.allowed is False + assert result.reason == "server_reservation_failed" + assert result.error == "server_reservation_failed" + + async def test_fails_open_to_the_per_check_default_value(self): + self.client.features.check_and_reserve_flag.side_effect = Exception("ECONNRESET") + result = await self.client.check( + "inference", + company={"id": "co_1"}, + options=CheckOptions(usage=50, on_acquire_failure="fail-open", default_value=True), + ) + assert result.allowed is True + assert result.reason == "server_reservation_failed_fail_open" + assert result.error == "server_reservation_failed" + + async def test_fails_open_to_the_client_level_flag_default(self): + client = _async_server_client(flag_defaults={"inference": True}) + try: + client.features.check_and_reserve_flag.side_effect = Exception("ECONNRESET") + result = await client.check( + "inference", company={"id": "co_1"}, options=CheckOptions(usage=50, on_acquire_failure="fail-open"), + ) + assert result.allowed is True + assert result.reason == "server_reservation_failed_fail_open" + finally: + await client.event_buffer.stop() + + self.client.features.check_and_reserve_flag.side_effect = Exception("ECONNRESET") + denied = await self.client.check( + "inference", company={"id": "co_1"}, options=CheckOptions(usage=50, on_acquire_failure="fail-open"), + ) + assert denied.allowed is False + + async def test_zero_usage_falls_back_to_a_plain_check(self): + result = await self.client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=0)) + self.client.features.check_and_reserve_flag.assert_not_called() + self.client.features.check_flag.assert_called_once() + assert result.allowed is True + assert result.reason == "plain check" + assert result.reservation is None + + async def test_invalid_usage_resolves_through_the_failure_contract(self): + for usage in (-5, float("nan"), float("inf"), True): + denied = await self.client.check( + "inference", company={"id": "co_1"}, options=CheckOptions(usage=usage), # type: ignore[arg-type] + ) + assert denied.allowed is False + assert denied.reason == "invalid_usage" + assert denied.error == "invalid_usage" + + opened = await self.client.check( + "inference", + company={"id": "co_1"}, + options=CheckOptions(usage=-5, on_acquire_failure="fail-open", default_value=True), + ) + assert opened.allowed is True + assert opened.reason == "invalid_usage_fail_open" + + self.client.features.check_and_reserve_flag.assert_not_called() + self.client.features.check_flag.assert_not_called() + + async def test_releases_a_hold_that_names_no_event_subtype(self): + self.client.features.check_and_reserve_flag.return_value = _reserve_response( + reservation=_held_reservation(id="rsv_orphan", event_subtype=None), + ) + result = await self.client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50)) + self.client.credits.release_credit_reservation.assert_awaited_once_with("rsv_orphan") + assert result.allowed is False + assert result.reason == "missing_event_subtype" + assert result.reservation is None + + async def test_a_failed_release_is_swallowed(self): + self.client.features.check_and_reserve_flag.return_value = _reserve_response( + reservation=_held_reservation(id="rsv_orphan", event_subtype=None), + ) + self.client.credits.release_credit_reservation.side_effect = Exception("boom") + result = await self.client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50)) + assert result.reason == "missing_event_subtype" + self.client.logger.warning.assert_called() + + async def test_a_released_hold_keeps_the_server_verdict_when_failing_open(self): + self.client.features.check_and_reserve_flag.return_value = _reserve_response( + reservation=_held_reservation(id="rsv_orphan", event_subtype=None), + ) + result = await self.client.check( + "inference", company={"id": "co_1"}, options=CheckOptions(usage=50, on_acquire_failure="fail-open"), + ) + self.client.credits.release_credit_reservation.assert_awaited_once_with("rsv_orphan") + # The server evaluated the flag and allowed it; only the settle is + # impossible, and fail-open assumes the credits are there. + assert result.allowed is True + assert result.value is True + assert result.reason == "matched" + assert result.flag_id == "flag_1" + assert result.entitlement == CREDIT_ENTITLEMENT + assert result.error == "missing_event_subtype" + assert result.reservation is None + + async def test_no_credit_lease_config_falls_back_to_a_plain_check(self): + client = _async_server_client(credit_leases=None) + try: + result = await client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50)) + client.features.check_and_reserve_flag.assert_not_called() + client.features.check_flag.assert_called_once() + assert result.reservation is None + assert client.features.check_flag.call_args.kwargs["preflight"] == PreflightRequestBody(usage=50) + finally: + await client.event_buffer.stop() + + async def test_client_mode_falls_back_and_warns_at_construction(self): + client = _async_server_client(credit_leases=CreditLeaseConfig(mode="client")) + try: + warning = " ".join(str(call.args[0]) for call in client.logger.warning.call_args_list) + assert "'client'" in warning + result = await client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50)) + client.features.check_and_reserve_flag.assert_not_called() + client.features.check_flag.assert_called_once() + assert result.reservation is None + finally: + await client.event_buffer.stop() + + async def test_offline_check_returns_the_flag_default(self): + client = _async_server_client(offline=True, flag_defaults={"inference": True}) + try: + result = await client.check("inference", company={"id": "co_1"}, options=CheckOptions(usage=50)) + client.features.check_and_reserve_flag.assert_not_called() + assert result.allowed is True + assert result.reason == REASON_OFFLINE + assert result.reservation is None + finally: + await client.event_buffer.stop() + + async def _reservation_handle(self) -> Reservation: + result = await self.client.check( + "inference", + company={"id": "co_1"}, + user={"id": "user_1"}, + options=CheckOptions(usage=50, event_subtype="inference_tokens"), + ) + assert result.reservation is not None + return result.reservation + + async def test_track_with_reservation_settles_by_reservation_id(self): + reservation = await self._reservation_handle() + with patch.object(self.client.event_buffer, "push", new=AsyncMock()) as mock_push: + await self.client.track_with_reservation( + reservation, 20, TrackWithReservationOptions(traits={"model": "opus"}), + ) + + pushed = mock_push.call_args.args[0] + assert pushed.event_type == "track" + assert pushed.body.event == "inference_tokens" + assert pushed.body.quantity == 20 + assert pushed.body.reservation_id == "rsv_1" + assert pushed.body.lease_id is None + assert pushed.body.company == {"id": "co_1"} + assert pushed.body.user == {"id": "user_1"} + assert pushed.body.traits == {"model": "opus"} + assert pushed.idempotency_key == "lease-reservation:rsv_1" + + async def test_track_with_reservation_updates_datastream_company_metrics(self): + reservation = await self._reservation_handle() + mock_ds = MagicMock() + mock_ds.is_connected = MagicMock(return_value=True) + mock_ds.update_company_metrics = AsyncMock() + self.client._datastream_client = mock_ds + + with patch.object(self.client.event_buffer, "push", new=AsyncMock()): + await self.client.track_with_reservation(reservation, 20) + + mock_ds.update_company_metrics.assert_awaited_once_with({"id": "co_1"}, "inference_tokens", 20) + + async def test_track_with_reservation_skips_an_invalid_quantity(self): + reservation = await self._reservation_handle() + with patch.object(self.client.event_buffer, "push", new=AsyncMock()) as mock_push: + for quantity in (-1, float("nan"), float("inf"), True): + await self.client.track_with_reservation(reservation, quantity) # type: ignore[arg-type] + mock_push.assert_not_called() + + async def test_track_with_reservation_settles_a_fractional_quantity_as_a_whole_unit(self): + reservation = await self._reservation_handle() + with patch.object(self.client.event_buffer, "push", new=AsyncMock()) as mock_push: + await self.client.track_with_reservation(reservation, 0.5) + # A track event's quantity is an integer, so a partial unit bills as one. + assert mock_push.call_args.args[0].body.quantity == 1 + + async def test_track_with_reservation_without_a_hold_says_to_use_track(self): + with patch.object(self.client.event_buffer, "push", new=AsyncMock()) as mock_push: + await self.client.track_with_reservation(None, 5) + mock_push.assert_not_called() + assert "track()" in str(self.client.logger.error.call_args.args[0]) + + async def test_track_with_reservation_is_a_no_op_when_offline(self): + reservation = await self._reservation_handle() + self.client.offline = True + with patch.object(self.client.event_buffer, "push", new=AsyncMock()) as mock_push: + await self.client.track_with_reservation(reservation, 20) + mock_push.assert_not_called() + + +LEASE_PROBE = { + "value": True, + "reason": "probe", + "entitlement": { + "value_type": "credit", + "credit_id": "bilcr_inference", + "consumption_rate": 10, + "event_subtype": "inference_tokens", + }, +} +LEASE_GATE = {"value": True, "reason": "matched"} + + +def _lease_datastream(results: list, **overrides) -> ScriptedDataStream: + """A DataStream stub carrying a scripted engine, wired for the client.""" + datastream = ScriptedDataStream( + ScriptedEngine(results, "inference"), + "inference", + {"id": "co_1", "credit_balances": {"bilcr_inference": 5000}}, + **overrides, + ) + datastream.is_connected = MagicMock(return_value=True) # type: ignore[attr-defined] + datastream.close = AsyncMock() # type: ignore[attr-defined] + datastream.update_company_metrics = AsyncMock() # type: ignore[attr-defined] + return datastream + + +def _lease_grant(lease_id: str = "lse_1", granted_amount: float = 1000.0) -> MagicMock: + return MagicMock( + data=MagicMock( + id=lease_id, + company_id="co_1", + credit_type_id="bilcr_inference", + granted_amount=granted_amount, + expires_at=dt.datetime.now(dt.timezone.utc) + dt.timedelta(seconds=300), + ) + ) + + +def _async_lease_client(**config_overrides) -> AsyncSchematic: + config_kwargs = dict( + event_buffer_period=1, + logger=MagicMock(), + httpx_client=MagicMock(spec=AsyncClient), + use_datastream=True, + credit_leases=CreditLeaseConfig(default_lease_size=1000.0, sweep_interval=60.0), + ) + config_kwargs.update(config_overrides) + client = AsyncSchematic("test_key", AsyncSchematicConfig(**config_kwargs)) # type: ignore[arg-type] + client.features.check_and_reserve_flag = AsyncMock(return_value=_reserve_response()) + client.features.check_flag = AsyncMock( + return_value=MagicMock(data=CheckFlagResponseData(value=True, flag="inference", reason="plain check")) + ) + client.credits.acquire_credit_lease = AsyncMock(return_value=_lease_grant()) + client.credits.extend_credit_lease = AsyncMock(return_value=_lease_grant()) + client.credits.release_credit_lease = AsyncMock() + client.flag_check_cache_providers = [] + return client + + +@pytest.mark.asyncio +class TestAsyncSchematicClientLeases: + """Routing, settling, prewarming, and shutdown with client-mode leases.""" + + async def _drain(self, client: AsyncSchematic) -> None: + if client._lease_manager is not None: + await client._lease_manager._drain_background() + await client.event_buffer.stop() + + async def _check(self, client: AsyncSchematic, **option_overrides) -> Any: + options = CheckOptions(usage=50, event_subtype="inference_tokens") + for name, value in option_overrides.items(): + setattr(options, name, value) + return await client.check("inference", company={"id": "co_1"}, options=options) + + async def test_auto_with_datastream_gates_on_a_local_lease(self): + client = _async_lease_client() + client._datastream_client = _lease_datastream([LEASE_PROBE, LEASE_GATE]) + try: + result = await self._check(client) + assert result.allowed is True + assert result.reservation is not None + assert result.reservation.mode == "client" + assert result.reservation.lease_id == "lse_1" + assert result.reservation.credits_reserved == 500 + client.credits.acquire_credit_lease.assert_awaited_once() + client.features.check_and_reserve_flag.assert_not_called() + finally: + await self._drain(client) + + async def test_auto_falls_back_to_server_mode_when_datastream_fails_to_start(self): + client = _async_lease_client() + client._datastream_client.start = AsyncMock(side_effect=RuntimeError("no socket")) # type: ignore[union-attr] + try: + await client.initialize() + assert client._datastream_client is None + result = await self._check(client) + # The plumbing is built but unusable without DataStream, so the + # check gates over the API rather than going ungated. + client.features.check_and_reserve_flag.assert_awaited_once() + client.credits.acquire_credit_lease.assert_not_awaited() + assert result.reservation is not None + assert result.reservation.mode == "server" + finally: + await self._drain(client) + + async def test_client_mode_without_datastream_checks_plainly_and_warns(self): + client = _async_lease_client(use_datastream=False, credit_leases=CreditLeaseConfig(mode="client")) + try: + warnings = " ".join(str(call.args[0]) for call in client.logger.warning.call_args_list) + assert "DataStream is not enabled" in warnings + result = await self._check(client) + assert result.reservation is None + client.credits.acquire_credit_lease.assert_not_awaited() + client.features.check_and_reserve_flag.assert_not_called() + client.features.check_flag.assert_awaited_once() + finally: + await self._drain(client) + + async def test_server_mode_warns_about_the_client_only_options(self): + client = _async_lease_client( + use_datastream=False, + credit_leases=CreditLeaseConfig( + mode="server", + default_lease_size=500.0, + overrides={"bilcr_inference": LeaseConfigOverride(lease_size=10.0)}, + ), + ) + try: + warnings = " ".join(str(call.args[0]) for call in client.logger.warning.call_args_list) + assert "default_lease_size" in warnings + assert "overrides" in warnings + assert client._lease_manager is None + assert client._lease_store is None + finally: + await self._drain(client) + + async def test_auto_without_datastream_warns_about_the_client_only_options(self): + client = _async_lease_client( + use_datastream=False, + credit_leases=CreditLeaseConfig(default_lease_size=500.0, sweep_interval=60.0), + ) + try: + warnings = " ".join(str(call.args[0]) for call in client.logger.warning.call_args_list) + assert "resolves to server mode" in warnings + assert "default_lease_size" in warnings + finally: + await self._drain(client) + + async def test_auto_with_datastream_keeps_the_client_only_options(self): + client = _async_lease_client() + try: + warnings = " ".join(str(call.args[0]) for call in client.logger.warning.call_args_list) + assert "resolves to server mode" not in warnings + finally: + await self._drain(client) + + async def test_no_shared_backend_warns_that_gating_is_per_process(self): + client = _async_lease_client() + try: + warnings = " ".join(str(call.args[0]) for call in client.logger.warning.call_args_list) + assert "without a shared Redis backend" in warnings + finally: + await self._drain(client) + + async def test_track_with_reservation_settles_the_local_hold(self): + client = _async_lease_client() + client._datastream_client = _lease_datastream([LEASE_PROBE, LEASE_GATE]) + try: + result = await self._check(client) + assert result.reservation is not None + with patch.object(client.event_buffer, "push", new=AsyncMock()) as mock_push: + await client.track_with_reservation(result.reservation, 20) + + pushed = mock_push.call_args.args[0] + assert pushed.body.event == "inference_tokens" + assert pushed.body.quantity == 20 + assert pushed.body.lease_id == "lse_1" + assert pushed.body.reservation_id is None + assert pushed.idempotency_key == f"lease-reservation:{result.reservation.id}" + # 1000 granted, 500 held, 200 of it actually consumed. + entry = await client._lease_store.get("co_1", "bilcr_inference") + assert entry is not None and entry.local_remaining_credits == 800 + finally: + await self._drain(client) + + async def test_track_with_reservation_emits_even_when_the_settle_raises(self): + client = _async_lease_client() + client._datastream_client = _lease_datastream([LEASE_PROBE, LEASE_GATE]) + try: + result = await self._check(client) + assert result.reservation is not None + client._reservations.consume = AsyncMock(side_effect=RuntimeError("redis down")) + with patch.object(client.event_buffer, "push", new=AsyncMock()) as mock_push: + await client.track_with_reservation(result.reservation, 7) + + # The server is the source of truth for consumption, so the usage + # is billed whatever the local bookkeeping did. + pushed = mock_push.call_args.args[0] + assert pushed.body.quantity == 7 + assert pushed.body.lease_id == "lse_1" + assert pushed.idempotency_key == f"lease-reservation:{result.reservation.id}" + finally: + await self._drain(client) + + async def test_an_unconfigured_client_still_bills_a_client_mode_handle(self): + client = _async_lease_client(credit_leases=None) + reservation = Reservation( + id="res_orphan", + lease_id="lse_x", + mode="client", + company_id="co_1", + credit_type_id="bilcr_inference", + event_subtype="inference_tokens", + quantity_reserved=10, + credits_reserved=100, + consumption_rate=10, + expires_at=dt.datetime.now(dt.timezone.utc), + company={"id": "co_1"}, + ) + try: + with patch.object(client.event_buffer, "push", new=AsyncMock()) as mock_push: + await client.track_with_reservation(reservation, 7) + pushed = mock_push.call_args.args[0] + assert pushed.body.lease_id == "lse_x" + assert pushed.idempotency_key == "lease-reservation:res_orphan" + finally: + await self._drain(client) + + async def test_prewarm_acquires_a_lease_per_credit_type(self): + client = _async_lease_client() + client._datastream_client = _lease_datastream([]) + try: + await client.prewarm({"id": "co_1"}, ["bilcr_inference"]) + client.credits.acquire_credit_lease.assert_awaited_once() + assert client.credits.acquire_credit_lease.call_args.kwargs["credit_type_id"] == "bilcr_inference" + finally: + await self._drain(client) + + async def test_prewarm_resolves_a_company_that_carries_only_secondary_keys(self): + client = _async_lease_client() + client._datastream_client = _lease_datastream([]) + try: + await client.prewarm({"external_id": "ext-co-1"}, ["bilcr_inference"]) + client.credits.acquire_credit_lease.assert_awaited_once() + assert client.credits.acquire_credit_lease.call_args.kwargs["company_id"] == "co_1" + finally: + await self._drain(client) + + async def test_prewarm_gives_up_when_the_company_never_surfaces(self): + client = _async_lease_client( + credit_leases=CreditLeaseConfig(default_lease_size=1000.0, prewarm_resolve_timeout=0.05) + ) + client._datastream_client = _lease_datastream( + [], company_error=RuntimeError("DataStream client is not connected") + ) + try: + await client.prewarm({"external_id": "ext-co-missing"}, ["bilcr_inference"]) + client.credits.acquire_credit_lease.assert_not_awaited() + finally: + await self._drain(client) + + async def test_prewarm_with_no_wait_still_warms_a_cached_company(self): + client = _async_lease_client( + credit_leases=CreditLeaseConfig(default_lease_size=1000.0, prewarm_resolve_timeout=0) + ) + # The fetch would fail, so only the cache can answer here. + client._datastream_client = _lease_datastream( + [], company_cached=True, company_error=RuntimeError("DataStream client is not connected") + ) + try: + await client.prewarm({"external_id": "ext-co-1"}, ["bilcr_inference"]) + client.credits.acquire_credit_lease.assert_awaited_once() + assert client.credits.acquire_credit_lease.call_args.kwargs["company_id"] == "co_1" + finally: + await self._drain(client) + + async def test_prewarm_with_no_wait_gives_up_on_an_uncached_company(self): + client = _async_lease_client( + credit_leases=CreditLeaseConfig(default_lease_size=1000.0, prewarm_resolve_timeout=0) + ) + client._datastream_client = _lease_datastream([]) + try: + await client.prewarm({"external_id": "ext-co-1"}, ["bilcr_inference"]) + client.credits.acquire_credit_lease.assert_not_awaited() + finally: + await self._drain(client) + + async def test_prewarm_is_a_no_op_in_server_mode(self): + client = _async_server_client() + try: + await client.prewarm({"id": "co_1"}, ["bilcr_inference"]) + debug = " ".join(str(call.args[0]) for call in client.logger.debug.call_args_list) + assert "no-op in server mode" in debug + finally: + await client.event_buffer.stop() + + async def test_identify_kicks_off_a_prewarm(self): + client = _async_lease_client() + client.prewarm = AsyncMock() # type: ignore[method-assign] + try: + with patch.object(client.event_buffer, "push", new=AsyncMock()): + await client.identify( + {"id": "user_1"}, + company=EventBodyIdentifyCompany(keys={"id": "co_1"}), + options=IdentifyOptions(prewarm=["bilcr_inference"]), + ) + await asyncio.sleep(0) + await asyncio.sleep(0) + client.prewarm.assert_awaited_once_with({"id": "co_1"}, ["bilcr_inference"]) + finally: + await self._drain(client) + + async def test_identify_without_prewarm_warms_nothing(self): + client = _async_lease_client() + client.prewarm = AsyncMock() # type: ignore[method-assign] + try: + with patch.object(client.event_buffer, "push", new=AsyncMock()): + await client.identify({"id": "user_1"}, company=EventBodyIdentifyCompany(keys={"id": "co_1"})) + await asyncio.sleep(0) + client.prewarm.assert_not_awaited() + finally: + await self._drain(client) + + async def test_identify_flushes_the_buffer_before_prewarming(self): + client = _async_lease_client() + client.prewarm = AsyncMock() # type: ignore[method-assign] + try: + with patch.object(client.event_buffer, "push", new=AsyncMock()): + with patch.object(client.event_buffer, "flush", new=AsyncMock()) as mock_flush: + await client.identify( + {"id": "user_1"}, + company=EventBodyIdentifyCompany(keys={"id": "co_1"}), + options=IdentifyOptions(prewarm=["bilcr_inference"]), + ) + # The prewarm polls for the company this identify creates, + # so the identify has to be on the wire before it starts. + mock_flush.assert_awaited_once() + client.prewarm.assert_not_awaited() + await asyncio.sleep(0) + await asyncio.sleep(0) + client.prewarm.assert_awaited_once_with({"id": "co_1"}, ["bilcr_inference"]) + finally: + await self._drain(client) + + async def test_identify_prewarms_even_when_the_flush_fails(self): + client = _async_lease_client() + client.prewarm = AsyncMock() # type: ignore[method-assign] + try: + with patch.object(client.event_buffer, "push", new=AsyncMock()): + with patch.object( + client.event_buffer, "flush", new=AsyncMock(side_effect=RuntimeError("api down")) + ): + await client.identify( + {"id": "user_1"}, + company=EventBodyIdentifyCompany(keys={"id": "co_1"}), + options=IdentifyOptions(prewarm=["bilcr_inference"]), + ) + await asyncio.sleep(0) + await asyncio.sleep(0) + client.prewarm.assert_awaited_once_with({"id": "co_1"}, ["bilcr_inference"]) + finally: + await self._drain(client) + + async def test_a_lease_gated_check_reports_one_flag_check_event(self): + client = _async_lease_client() + client._datastream_client = _lease_datastream([LEASE_PROBE, LEASE_GATE]) + try: + with patch.object(client.event_buffer, "push", new=AsyncMock()) as mock_push: + result = await self._check(client) + assert result.allowed is True + events = [call.args[0] for call in mock_push.call_args_list] + assert [event.event_type for event in events] == ["flag_check"] + body = events[0].body + assert body.flag_key == "inference" + assert body.value is True + assert body.reason == "matched" + assert body.company_id == "co_1" + assert body.req_company == {"id": "co_1"} + finally: + await self._drain(client) + + async def test_client_mode_keeps_a_reservation_ttl_past_the_server_cap(self): + client = _async_lease_client( + credit_leases=CreditLeaseConfig( + mode="client", default_lease_size=1000.0, default_reservation_ttl=7200.0 + ) + ) + try: + # The TTL only drives the local sweeper here, so the server's cap + # does not apply and nothing is clamped or warned about. + assert client._reservation_ttl == 7200.0 + assert client._lease_manager.resolve_config("bilcr_inference").reservation_ttl == 7200.0 + warning = " ".join(str(call.args[0]) for call in client.logger.warning.call_args_list) + assert "one hour cap" not in warning + finally: + await self._drain(client) + + async def test_shutdown_stops_the_sweep_and_releases_a_per_process_lease(self): + client = _async_lease_client() + client._datastream_client = _lease_datastream([LEASE_PROBE, LEASE_GATE]) + await self._check(client) + client._lease_manager.start_sweep() + assert client._lease_manager._sweep_task is not None + + await client.shutdown() + + assert client._lease_manager._sweep_task is None + client.credits.release_credit_lease.assert_awaited_once_with("lse_1", request_options=None) + + async def test_shutdown_leaves_a_shared_lease_for_the_pods_still_drawing_on_it(self): + redis_client = make_fake_redis() + client = _async_lease_client( + credit_leases=CreditLeaseConfig(default_lease_size=1000.0, redis_client=redis_client) + ) + assert client._lease_backend_shared is True + await client._lease_store.replace( + lease_id="lse_shared", + company_id="co_1", + credit_type_id="bilcr_inference", + granted_amount=1000, + expires_at=time.time() + 300, + ) + + await client.shutdown() + + client.credits.release_credit_lease.assert_not_awaited() + survivor = await client._lease_store.get("co_1", "bilcr_inference") + assert survivor is not None and survivor.lease_id == "lse_shared" + + async def test_a_redis_backed_datastream_cache_backs_the_leases_too(self): + redis_client = make_fake_redis() + client = _async_lease_client( + datastream=DataStreamConfig(company_cache=RedisCache(redis_client, prefix="acme")), + ) + try: + assert client._lease_backend_shared is True + assert type(client._lease_store).__name__ == "RedisLeaseStore" + assert type(client._reservations).__name__ == "RedisReservationStore" + finally: + await self._drain(client) + + +class TestSchematicClientModeWarning(unittest.TestCase): + """The sync client cannot run client mode, and says so.""" + + def test_client_mode_points_at_the_async_client(self): + logger = MagicMock() + client = Schematic( + "api_key", + SchematicConfig( + event_buffer_period=1, + logger=logger, + httpx_client=MagicMock(spec=Client), + credit_leases=CreditLeaseConfig(mode="client"), + ), + ) + try: + warnings = " ".join(str(call.args[0]) for call in logger.warning.call_args_list) + self.assertIn("AsyncSchematic", warnings) + self.assertIsNone(client._effective_lease_mode()) + finally: + client.event_buffer.stop() + + + def test_auto_warns_about_the_client_only_options_too(self): + # Every 'auto' on the sync client resolves to server mode, so the + # client-only knobs are just as ignored as under an explicit 'server'. + logger = MagicMock() + client = Schematic( + "api_key", + SchematicConfig( + event_buffer_period=1, + logger=logger, + httpx_client=MagicMock(spec=Client), + credit_leases=CreditLeaseConfig(default_lease_size=500.0), + ), + ) + try: + warnings = " ".join(str(call.args[0]) for call in logger.warning.call_args_list) + self.assertIn("resolves to server mode", warnings) + self.assertIn("default_lease_size", warnings) + finally: + client.event_buffer.stop() + + if __name__ == "__main__": unittest.main() diff --git a/tests/datastream/test_datastream_client.py b/tests/datastream/test_datastream_client.py index be00292b..5764b633 100644 --- a/tests/datastream/test_datastream_client.py +++ b/tests/datastream/test_datastream_client.py @@ -489,6 +489,66 @@ async def test_flag_evaluation_with_cached_user(self, logger: logging.Logger) -> assert result.flag_key == "usr-flag" +class TestDataStreamClientCheckFlagOptions: + """check_flag threads the caller's preflight options to the engine.""" + + def _client(self, logger: logging.Logger) -> DataStreamClient: + cache = MockCacheProvider() + return DataStreamClient(DataStreamClientOptions( + api_key="test-key", + logger=logger, + replicator_mode=True, + company_cache=cache, + company_lookup_cache=cache, + user_cache=cache, + user_lookup_cache=cache, + flag_cache=cache, + )) + + async def _cache_flag(self, client: DataStreamClient) -> None: + await client._handle_message(DataStreamResp( + data={ + "key": "inference", "id": "f1", "default_value": True, "rules": [], + "account_id": "acc_1", "environment_id": "env_1", + }, + entity_type=EntityType.FLAG.value, + message_type=MessageType.FULL.value, + )) + + async def test_options_reach_the_rules_engine(self, logger: logging.Logger) -> None: + from schematic.client import CheckFlagOptions + + client = self._client(logger) + await self._cache_flag(client) + engine = MagicMock() + engine.is_initialized = MagicMock(return_value=True) + engine.get_version_key = MagicMock(return_value="1") + engine.check_flag = MagicMock( + return_value=RulesengineCheckFlagResult(value=True, flag_key="inference", reason="matched") + ) + client._rules_engine = engine + + options = CheckFlagOptions(usage=5) + await client.check_flag(CheckFlagRequestBody(), "inference", options) + + assert engine.check_flag.call_args.args[3] is options + + async def test_a_plain_check_passes_no_options(self, logger: logging.Logger) -> None: + client = self._client(logger) + await self._cache_flag(client) + engine = MagicMock() + engine.is_initialized = MagicMock(return_value=True) + engine.get_version_key = MagicMock(return_value="1") + engine.check_flag = MagicMock( + return_value=RulesengineCheckFlagResult(value=True, flag_key="inference", reason="matched") + ) + client._rules_engine = engine + + await client.check_flag(CheckFlagRequestBody(), "inference") + + assert engine.check_flag.call_args.args[3] is None + + class TestDataStreamClientPartialMerge: """Spec test #4: Partial entity message merges into cache.""" diff --git a/tests/datastream/test_rules_engine.py b/tests/datastream/test_rules_engine.py index 0020e55a..d20019f5 100644 --- a/tests/datastream/test_rules_engine.py +++ b/tests/datastream/test_rules_engine.py @@ -6,6 +6,8 @@ from schematic.types import ( RulesengineCheckFlagResult, RulesengineCompany, + RulesengineCompanyMetric, + RulesengineCondition, RulesengineFlag, RulesengineRule, ) @@ -239,6 +241,180 @@ async def test_billing_metric_override_populates_reset_at(self) -> None: assert result.feature_usage_reset_at is not None +class TestRulesEngineOptionsEnvelope: + """The preflight options block the engine evaluates against. + + The engine's serde struct is snake_case with defaulted fields, so unset + options are dropped and the key is left off entirely when the caller + preflighted nothing, keeping envelopes for plain checks unchanged. + """ + + @pytest.fixture + async def engine(self) -> RulesEngineClient: + e = RulesEngineClient() + await e.initialize() + return e + + def _capture(self, engine: RulesEngineClient) -> list[str]: + captured: list[str] = [] + original = engine._call_wasm + + def spy(input_json: str) -> str: + captured.append(input_json) + return original(input_json) + + engine._call_wasm = spy # type: ignore[method-assign] + return captured + + async def test_no_options_leaves_the_key_off(self, engine: RulesEngineClient) -> None: + import json + + captured = self._capture(engine) + engine.check_flag(_make_flag(default_value=True)) + assert "options" not in json.loads(captured[0]) + + async def test_options_with_no_preflight_leave_the_key_off(self, engine: RulesEngineClient) -> None: + import json + + from schematic.client import CheckFlagOptions + + captured = self._capture(engine) + engine.check_flag(_make_flag(default_value=True), None, None, CheckFlagOptions(default_value=True)) + assert "options" not in json.loads(captured[0]) + + async def test_usage_is_carried_in_the_options_block(self, engine: RulesEngineClient) -> None: + import json + + from schematic.client import CheckFlagOptions + + captured = self._capture(engine) + result = engine.check_flag(_make_flag(default_value=True), None, None, CheckFlagOptions(usage=5)) + assert json.loads(captured[0])["options"] == {"usage": 5} + assert result.value is True + + async def test_event_usage_and_credit_cost_are_carried_snake_cased(self, engine: RulesEngineClient) -> None: + import json + + from schematic.client import CheckFlagOptions, EventUsage + + captured = self._capture(engine) + engine.check_flag( + _make_flag(default_value=True), + None, + None, + CheckFlagOptions( + event_usage=EventUsage(event_subtype="inference_tokens", quantity=7), + credit_cost={"bilcr_inference": 12.5}, + ), + ) + assert json.loads(captured[0])["options"] == { + "credit_cost": {"bilcr_inference": 12.5}, + "event_usage": {"event_subtype": "inference_tokens", "quantity": 7}, + } + + +class TestRulesEnginePreflightVerdict: + """A preflight moves the verdict, not just the envelope. + + The company sits at 95 against a `usage < 100` condition: entitled as it + stands, and denied once the usage the caller is about to record counts + against the same condition. + """ + + @pytest.fixture + async def engine(self) -> RulesEngineClient: + e = RulesEngineClient() + await e.initialize() + return e + + def _metered_company(self) -> RulesengineCompany: + company_id = "co_metered" + company_condition = RulesengineCondition( + id="cond_company", + account_id="acc_1", + environment_id="env_1", + condition_type="company", + operator="eq", + resource_ids=[company_id], + trait_value="", + ) + metric_condition = RulesengineCondition( + id="cond_metric", + account_id="acc_1", + environment_id="env_1", + condition_type="metric", + operator="lt", + resource_ids=[], + event_subtype="api-calls", + metric_value=100, + metric_period="current_month", + metric_period_month_reset="billing_cycle", + trait_value="100", + ) + override_rule = RulesengineRule( + id="rule_override", + flag_id="flag1", + account_id="acc_1", + environment_id="env_1", + name="Company Override", + rule_type="company_override", + value=True, + priority=0, + conditions=[company_condition, metric_condition], + condition_groups=[], + ) + metric = RulesengineCompanyMetric( + account_id="acc_1", + environment_id="env_1", + company_id=company_id, + event_subtype="api-calls", + period="current_month", + month_reset="billing_cycle", + value=95, + created_at="2023-01-01T00:00:00Z", + ) + return RulesengineCompany( + id=company_id, + account_id="acc_1", + environment_id="env_1", + keys={"id": company_id}, + traits=[], + metrics=[metric], + rules=[override_rule], + entitlements=[], + billing_product_ids=[], + credit_balances={}, + plan_ids=[], + plan_version_ids=[], + ) + + def _flag(self) -> RulesengineFlag: + return _make_flag(id="flag1", key="api-access", default_value=False) + + async def test_allows_without_a_preflight(self, engine: RulesEngineClient) -> None: + result = engine.check_flag(self._flag(), self._metered_company()) + assert result.value is True + + async def test_usage_that_crosses_the_limit_denies(self, engine: RulesEngineClient) -> None: + from schematic.client import CheckFlagOptions + + result = engine.check_flag(self._flag(), self._metered_company(), None, CheckFlagOptions(usage=10)) + assert result.value is False + + async def test_event_usage_for_another_subtype_leaves_the_verdict_alone( + self, engine: RulesEngineClient, + ) -> None: + from schematic.client import CheckFlagOptions, EventUsage + + result = engine.check_flag( + self._flag(), + self._metered_company(), + None, + CheckFlagOptions(event_usage=EventUsage(event_subtype="other-events", quantity=10)), + ) + assert result.value is True + + class TestRulesEngineFileNotFound: async def test_missing_wasm_raises(self) -> None: engine = RulesEngineClient(wasm_path="/nonexistent/rulesengine.wasm") diff --git a/tests/lease_support.py b/tests/lease_support.py new file mode 100644 index 00000000..cb5a49fe --- /dev/null +++ b/tests/lease_support.py @@ -0,0 +1,366 @@ +"""Harness shared by the lease unit tests and the conformance runner. + +Everything here is test-only: the virtual clock, the fakeredis client the +verbatim Lua scripts can run against, the crash seam the bounded-leak tests +need, and scriptable stand-ins for the lease wire API, the rules engine, and +DataStream. +""" + +from __future__ import annotations + +import datetime as dt +from typing import Any, Awaitable, Dict, List, Optional, cast + +import fakeredis.aioredis + +from schematic.leases import LeaseGrant, LeaseState, ReservationRecord +from schematic.leases.lease_store import LeaseStore +from schematic.types import ( + RulesengineCheckFlagResult, + RulesengineCompany, + RulesengineFeatureEntitlement, + RulesengineFlag, +) + +# The fixed virtual instant every vector and test starts from. +T0 = dt.datetime(2026, 1, 1, tzinfo=dt.timezone.utc).timestamp() + + +class VirtualClock: + """A clock only ``advance`` moves, so no test depends on wall time.""" + + def __init__(self, now: float = T0) -> None: + self._now = now + + def __call__(self) -> float: + return self._now + + def advance_ms(self, milliseconds: float) -> None: + self._now += milliseconds / 1000.0 + + def at_ms(self, offset_ms: float) -> float: + """An absolute position on the virtual timeline, as the vectors express it.""" + return T0 + offset_ms / 1000.0 + + +class LuaCompatFakeRedis(fakeredis.aioredis.FakeRedis): + """fakeredis with the one Lua builtin its runtime lacks papered over. + + ``redis.replicate_commands()`` is what lets a real Redis 5/6 read TIME in a + writing script; fakeredis's Lua runtime does not define it. Stripping the + call here (rather than dropping it from the scripts) keeps the shipped Lua + byte-identical to the Node SDK's, which is what lets both fleets share one + Redis. The scripts are sent through SCRIPT LOAD as well as EVAL, so both + are rewritten and the server-assigned digest stays consistent. + """ + + @staticmethod + def _strip(script: str) -> str: + return script.replace("redis.replicate_commands()\n", "") + + async def script_load(self, script: str) -> Any: # type: ignore[override] + return await cast(Awaitable[Any], super().script_load(self._strip(script))) + + async def eval(self, script: str, numkeys: int, *keys_and_args: Any) -> Any: # type: ignore[override] + return await cast(Awaitable[Any], super().eval(self._strip(script), numkeys, *keys_and_args)) + + +def make_fake_redis() -> LuaCompatFakeRedis: + return LuaCompatFakeRedis(decode_responses=True) + + +class CrashingRefundLeaseStore(LeaseStore): + """Lease store whose ``refund`` raises once while armed. + + The reservation store refunds through the store it is handed, so wrapping + that one reproduces a process death between the claim and the refund while + the test body keeps reading the real store. + """ + + def __init__(self, target: LeaseStore) -> None: + self._target = target + self._armed = False + + def arm(self) -> None: + self._armed = True + + async def get(self, company_id: str, credit_type_id: str) -> Optional[LeaseState]: + return await self._target.get(company_id, credit_type_id) + + async def replace( + self, + *, + lease_id: str, + company_id: str, + credit_type_id: str, + granted_amount: float, + expires_at: float, + ) -> bool: + return await self._target.replace( + lease_id=lease_id, + company_id=company_id, + credit_type_id=credit_type_id, + granted_amount=granted_amount, + expires_at=expires_at, + ) + + async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[float]: + return await self._target.try_reserve(company_id, credit_type_id, credits) + + async def refund( + self, + company_id: str, + credit_type_id: str, + credits: float, + pin_lease_id: Optional[str] = None, + ) -> None: + if self._armed: + self._armed = False + raise RuntimeError("simulated crash before refund") + await self._target.refund(company_id, credit_type_id, credits, pin_lease_id) + + async def extend( + self, + company_id: str, + credit_type_id: str, + granted_total: float, + new_expires_at: Optional[float] = None, + pin_lease_id: Optional[str] = None, + ) -> None: + await self._target.extend(company_id, credit_type_id, granted_total, new_expires_at, pin_lease_id) + + async def drop(self, company_id: str, credit_type_id: str) -> None: + await self._target.drop(company_id, credit_type_id) + + def list_leases(self) -> Optional[List[LeaseState]]: + return self._target.list_leases() + + +class ScriptedWireClient: + """Stands in for the lease API: queued responses in, recorded calls out.""" + + def __init__(self) -> None: + self.acquire_responses: List[Dict[str, Any]] = [] + self.extend_responses: List[Dict[str, Any]] = [] + self.acquire_calls: List[Dict[str, Any]] = [] + self.extend_calls: List[Dict[str, Any]] = [] + self.release_calls: List[str] = [] + # Runs while an acquire is in flight, for emulating a sibling pod + # winning the race. + self.during_acquire: Optional[Any] = None + + async def acquire( + self, + company_id: str, + credit_type_id: str, + requested_amount: float, + expires_at: float, + timeout: Optional[float] = None, + ) -> LeaseGrant: + self.acquire_calls.append( + { + "company_id": company_id, + "credit_type_id": credit_type_id, + "requested_amount": requested_amount, + "expires_at": expires_at, + "timeout": timeout, + } + ) + during = self.during_acquire + if during is not None: + self.during_acquire = None + await during() + scripted = self.acquire_responses.pop(0) if self.acquire_responses else None + lease = _scripted_lease(scripted, "unscripted acquire wire call") + return LeaseGrant( + lease_id=lease.get("lease_id", "lse_unnamed"), + company_id=company_id, + credit_type_id=credit_type_id, + granted_amount=float(lease.get("granted_amount", 0)), + expires_at=lease["expires_at"], + ) + + async def extend( + self, + lease_id: str, + additional_amount: float, + expires_at: float, + timeout: Optional[float] = None, + ) -> LeaseGrant: + self.extend_calls.append( + { + "lease_id": lease_id, + "additional_amount": additional_amount, + "expires_at": expires_at, + "timeout": timeout, + } + ) + scripted = self.extend_responses.pop(0) if self.extend_responses else None + lease = _scripted_lease(scripted, "unscripted extend wire call") + return LeaseGrant( + lease_id=lease_id, + company_id=lease.get("company_id", "co_wire"), + credit_type_id=lease.get("credit_type_id", "ct_wire"), + granted_amount=float(lease.get("granted_total", lease.get("granted_amount", 0))), + expires_at=lease["expires_at"], + ) + + async def release(self, lease_id: str) -> None: + self.release_calls.append(lease_id) + + +def _scripted_lease(scripted: Optional[Dict[str, Any]], missing: str) -> Dict[str, Any]: + if scripted is None: + raise RuntimeError(missing) + if scripted.get("error") is not None or not scripted.get("lease"): + raise RuntimeError(str(scripted.get("error") or missing)) + lease: Dict[str, Any] = dict(scripted["lease"]) + return lease + + +def make_reservation(**overrides: Any) -> ReservationRecord: + fields: Dict[str, Any] = { + "id": "res_1", + "lease_id": "lse_1", + "company_id": "co_1", + "credit_type_id": "ct_1", + "event_subtype": "inference_tokens", + "quantity_reserved": 10, + "credits_reserved": 100, + "consumption_rate": 10, + "expires_at": T0 + 60, + "company": {"id": "co_1"}, + } + fields.update(overrides) + return ReservationRecord(**fields) + + +class ScriptedEngine: + """Stands in for the WASM rules engine: queued verdicts in, calls out. + + The vectors treat the engine as an oracle. What they pin is the + orchestration around it, so every call records the credit balance the SDK + substituted and the preflight it threaded. + """ + + def __init__(self, results: List[Dict[str, Any]], flag_key: str = "flag") -> None: + self._results = list(results) + self._flag_key = flag_key + self.calls: List[Dict[str, Any]] = [] + + def __call__( + self, + flag: Any, + company: Any, + user: Any, + options: Any = None, + ) -> RulesengineCheckFlagResult: + event_usage = getattr(options, "event_usage", None) + self.calls.append( + { + "credit_balances": dict(getattr(company, "credit_balances", None) or {}), + "credit_cost": getattr(options, "credit_cost", None), + "event_usage": ( + {"event_subtype": event_usage.event_subtype, "quantity": event_usage.quantity} + if event_usage is not None + else None + ), + "usage": getattr(options, "usage", None), + } + ) + if not self._results: + raise RuntimeError(f"unscripted engine call for flag {self._flag_key}") + scripted = self._results.pop(0) + entitlement = scripted.get("entitlement") + return RulesengineCheckFlagResult( + value=scripted["value"], + reason=scripted.get("reason") or "", + flag_key=self._flag_key, + flag_id="flag_1", + entitlement=_scripted_entitlement(entitlement, self._flag_key) if entitlement else None, + ) + + +class ScriptedDataStream: + """The slice of ``DataStreamClient`` a lease-bearing check touches. + + The keyword arguments stage the misses the check flow has to survive: a + flag that is not cached, a company or user the socket cannot resolve. + """ + + def __init__( + self, + engine: Any, + flag_key: str, + company: Dict[str, Any], + *, + user: Optional[Any] = None, + missing_flag: bool = False, + company_error: Optional[Exception] = None, + user_error: Optional[Exception] = None, + company_cached: bool = False, + ) -> None: + self._engine = engine + self._flag_key = flag_key + self._company = make_company(company["id"], company.get("credit_balances") or {}) + self._user = user + self._missing_flag = missing_flag + self._company_error = company_error + self._user_error = user_error + self._company_cached = company_cached + + async def get_flag(self, flag_key: str) -> Optional[RulesengineFlag]: + if self._missing_flag: + return None + return RulesengineFlag( + id="flag_1", + key=self._flag_key, + account_id="acc_1", + environment_id="env_1", + default_value=False, + rules=[], + ) + + async def get_company(self, keys: Dict[str, str]) -> RulesengineCompany: + if self._company_error is not None: + raise self._company_error + return self._company + + async def get_cached_company(self, keys: Dict[str, str]) -> Optional[RulesengineCompany]: + return self._company if self._company_cached else None + + async def get_user(self, keys: Dict[str, str]) -> Any: + if self._user_error is not None: + raise self._user_error + return self._user + + def evaluate_flag(self, flag: Any, company: Any, user: Any, options: Any = None) -> RulesengineCheckFlagResult: + return self._engine(flag, company, user, options) + + +def make_company(company_id: str, credit_balances: Dict[str, float]) -> RulesengineCompany: + return RulesengineCompany( + id=company_id, + account_id="acc_1", + environment_id="env_1", + keys={"id": company_id}, + traits=[], + metrics=[], + rules=[], + entitlements=[], + billing_product_ids=[], + credit_balances=dict(credit_balances), + plan_ids=[], + plan_version_ids=[], + ) + + +def _scripted_entitlement(spec: Dict[str, Any], flag_key: str) -> RulesengineFeatureEntitlement: + return RulesengineFeatureEntitlement( + feature_id=spec.get("feature_id") or "feat_1", + feature_key=spec.get("feature_key") or flag_key, + value_type=spec["value_type"], + credit_id=spec.get("credit_id"), + consumption_rate=spec.get("consumption_rate"), + event_subtype=spec.get("event_subtype"), + ) diff --git a/tests/leases/__init__.py b/tests/leases/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/leases/conftest.py b/tests/leases/conftest.py new file mode 100644 index 00000000..09213037 --- /dev/null +++ b/tests/leases/conftest.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +from typing import Iterator +from unittest import mock + +import pytest +from lease_support import VirtualClock, make_fake_redis + + +@pytest.fixture +def clock() -> VirtualClock: + return VirtualClock() + + +@pytest.fixture +def frozen_clock(clock: VirtualClock) -> Iterator[VirtualClock]: + """A virtual clock that is also the process clock. + + The Redis stores decide expiry against the store's own clock (Redis TIME), + which fakeredis reads from ``time.time``, so a test that moves the virtual + clock has to move that one too. + """ + with mock.patch("time.time", clock): + yield clock + + +@pytest.fixture +def redis_client(frozen_clock: VirtualClock) -> object: + return make_fake_redis() diff --git a/tests/leases/test_check_and_track.py b/tests/leases/test_check_and_track.py new file mode 100644 index 00000000..e7a7b555 --- /dev/null +++ b/tests/leases/test_check_and_track.py @@ -0,0 +1,675 @@ +"""The client-mode check and settle flow against scripted stores and engine. + +Ports schematic-node's check-and-track suite. The engine is scripted rather +than real (``test_wasm_credit_gate`` drives the real one), so what these pin is +the orchestration: which balance the engine is handed, when a hold is taken, +and what happens to it when something downstream says no. +""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +import pytest +from lease_support import ScriptedDataStream, ScriptedWireClient, VirtualClock + +from schematic.client import CheckOptions, CheckResult, Reservation +from schematic.leases import ( + CreditCheckDeps, + InMemoryLeaseStore, + InMemoryReservationStore, + LeaseConfig, + LeaseManager, + LeaseStore, + ReservationRecord, + check_with_lease, + consume_reservation_and_build_event, +) +from schematic.leases.check import FAIL_OPEN_BALANCE +from schematic.types import ( + EventBodyFlagCheck, + RulesengineCheckFlagResult, + RulesengineFeatureEntitlement, + RulesengineUser, +) + +FLAG_KEY = "inference" +CREDIT_ID = "bilcr_inference" +EVENT_SUBTYPE = "inference_tokens" +LEASE_SIZE = 1000.0 +LEASE_DURATION = 300.0 +COMPANY = {"id": "co_1"} + +CREDIT_ENTITLEMENT = RulesengineFeatureEntitlement( + feature_id="feat", + feature_key=FLAG_KEY, + value_type="credit", + credit_id=CREDIT_ID, + consumption_rate=10, + event_subtype=EVENT_SUBTYPE, +) + + +def _verdict( + value: bool, reason: str, entitlement: Optional[RulesengineFeatureEntitlement] = None, +) -> RulesengineCheckFlagResult: + return RulesengineCheckFlagResult( + value=value, reason=reason, flag_key=FLAG_KEY, flag_id="flag_1", entitlement=entitlement + ) + + +class FlowEngine: + """The engine the lease path asks twice: the probe, then the gate. + + The gate is the call carrying ``credit_cost``, which is also what puts the + fail-open re-evaluation on the probe branch, as it is in the reference + implementation. + """ + + def __init__( + self, + *, + probe: Optional[RulesengineCheckFlagResult] = None, + gate: Optional[RulesengineCheckFlagResult] = None, + probe_error: Optional[Exception] = None, + gate_error: Optional[Exception] = None, + ) -> None: + self.probe = probe if probe is not None else _verdict(True, "probe", CREDIT_ENTITLEMENT) + self.gate = gate if gate is not None else _verdict(True, "matched", CREDIT_ENTITLEMENT) + self.probe_error = probe_error + self.gate_error = gate_error + self.calls: List[Dict[str, Any]] = [] + + def __call__(self, flag: Any, company: Any, user: Any, options: Any = None) -> RulesengineCheckFlagResult: + gating = bool(options is not None and getattr(options, "credit_cost", None)) + self.calls.append({"company": company, "user": user, "options": options, "gating": gating}) + if gating: + if self.gate_error is not None: + raise self.gate_error + return self.gate + if self.probe_error is not None: + raise self.probe_error + return self.probe + + @property + def gate_call(self) -> Optional[Dict[str, Any]]: + return next((call for call in self.calls if call["gating"]), None) + + def balance(self, index: int) -> float: + return float(self.calls[index]["company"].credit_balances[CREDIT_ID]) + + +class _ProbeThenExplodes(FlowEngine): + """Answers the probe, then fails every evaluation after it.""" + + def __call__(self, flag: Any, company: Any, user: Any, options: Any = None) -> RulesengineCheckFlagResult: + result = super().__call__(flag, company, user, options) + if len(self.calls) > 1: + raise RuntimeError("wasm exploded") + return result + + +class Fallback: + """The plain flag check the lease path defers to.""" + + def __init__(self) -> None: + self.called = False + + async def __call__(self) -> CheckResult: + self.called = True + return CheckResult(allowed=True, value=True, reason="fallback", flag_key=FLAG_KEY) + + +class UnreachableLeaseStore(InMemoryLeaseStore): + """A store that can be read but never debited, as an unreachable Redis is.""" + + async def try_reserve(self, company_id: str, credit_type_id: str, credits: float) -> Optional[float]: + raise RuntimeError("redis down") + + +class RecordedFlagChecks: + """Stands in for the client's event buffer so the tests can see what the + lease path reported.""" + + def __init__(self) -> None: + self.events: List[EventBodyFlagCheck] = [] + self.explode = False + + async def __call__(self, body: EventBodyFlagCheck) -> None: + if self.explode: + raise RuntimeError("event buffer down") + self.events.append(body) + + +@dataclass +class Flow: + deps: CreditCheckDeps + engine: FlowEngine + wire: ScriptedWireClient + leases: LeaseStore + reservations: InMemoryReservationStore + manager: LeaseManager + flag_checks: RecordedFlagChecks = field(default_factory=RecordedFlagChecks) + fallback: Fallback = field(default_factory=Fallback) + + async def check(self, **option_overrides: Any) -> CheckResult: + options = CheckOptions(usage=50, event_subtype=EVENT_SUBTYPE) + for name, value in option_overrides.items(): + setattr(options, name, value) + result = await check_with_lease(self.deps, FLAG_KEY, COMPANY, None, options, self.fallback) + await self.manager._drain_background() + return result + + async def remaining(self) -> Optional[float]: + entry = await self.leases.get(COMPANY["id"], CREDIT_ID) + return entry.local_remaining_credits if entry is not None else None + + +def make_flow( + clock: VirtualClock, + *, + engine: Optional[FlowEngine] = None, + lease_store: Optional[LeaseStore] = None, + acquire: str = "ok", + credit_balances: Optional[Dict[str, float]] = None, + **datastream_kwargs: Any, +) -> Flow: + engine = engine or FlowEngine() + leases = lease_store if lease_store is not None else InMemoryLeaseStore(clock=clock) + reservations = InMemoryReservationStore(leases, clock=clock) + wire = ScriptedWireClient() + if acquire == "ok": + wire.acquire_responses.append( + {"lease": {"lease_id": "lse_1", "granted_amount": LEASE_SIZE, "expires_at": clock() + LEASE_DURATION}} + ) + elif acquire == "error": + wire.acquire_responses.append({"error": "lease 503"}) + manager = LeaseManager( + wire, + leases, + reservation_store=reservations, + config=LeaseConfig( + lease_duration=LEASE_DURATION, reservation_ttl=60.0, lease_size=LEASE_SIZE, low_water_mark=0.25 + ), + clock=clock, + ) + datastream = ScriptedDataStream( + engine, + FLAG_KEY, + {"id": COMPANY["id"], "credit_balances": credit_balances if credit_balances is not None else {CREDIT_ID: 5000}}, + **datastream_kwargs, + ) + flag_checks = RecordedFlagChecks() + return Flow( + deps=CreditCheckDeps( + datastream=datastream, + lease_store=leases, + reservations=reservations, + manager=manager, + logger=logging.getLogger("lease-flow-test"), + enqueue_flag_check=flag_checks, + clock=clock, + ), + engine=engine, + wire=wire, + leases=leases, + reservations=reservations, + manager=manager, + flag_checks=flag_checks, + ) + + +class TestCheckWithLease: + async def test_issues_a_reservation_when_the_engine_allows(self, clock: VirtualClock) -> None: + flow = make_flow(clock) + result = await flow.check() + + assert result.allowed is True + assert result.reservation is not None + assert result.reservation.mode == "client" + assert result.reservation.lease_id == "lse_1" + assert result.reservation.quantity_reserved == 50 + assert result.reservation.credits_reserved == 500 + assert result.reservation.consumption_rate == 10 + assert len(flow.wire.acquire_calls) == 1 + # The probe sees the company's real balance; the gate sees the + # pre-reservation lease balance and the cost this call just debited. + assert flow.engine.balance(0) == 5000 + assert flow.engine.balance(1) == LEASE_SIZE + assert flow.engine.gate_call is not None + assert flow.engine.gate_call["options"].credit_cost == {CREDIT_ID: 500} + assert await flow.remaining() == 500 + assert await flow.reservations.count() == 1 + + async def test_leases_the_credit_the_matched_entitlement_names(self, clock: VirtualClock) -> None: + # Entitlement-first resolution: the credit and the rate come off the + # probe's entitlement, whatever the flag's conditions look like. + ai_entitlement = RulesengineFeatureEntitlement( + feature_id="feat", + feature_key=FLAG_KEY, + value_type="credit", + credit_id="bilcr_ai", + consumption_rate=5, + event_subtype=EVENT_SUBTYPE, + ) + engine = FlowEngine( + probe=_verdict(True, "probe", ai_entitlement), gate=_verdict(True, "matched", ai_entitlement) + ) + flow = make_flow(clock, engine=engine) + flow.wire.acquire_responses[0]["lease"]["lease_id"] = "lse_ai" + result = await flow.check() + + assert result.reservation is not None + assert result.reservation.credit_type_id == "bilcr_ai" + assert result.reservation.consumption_rate == 5 + assert result.reservation.credits_reserved == 250 + assert flow.wire.acquire_calls[0]["credit_type_id"] == "bilcr_ai" + assert len(flow.engine.calls) == 2 + + async def test_skips_the_lease_when_the_entitlement_is_not_credit_metered(self, clock: VirtualClock) -> None: + # A boolean grant (an override, say) draws no credit, so the lease path + # never acquires: the plain check decides, with no reserve to cancel. + entitlement = RulesengineFeatureEntitlement(feature_id="feat", feature_key=FLAG_KEY, value_type="boolean") + flow = make_flow(clock, engine=FlowEngine(probe=_verdict(True, "override", entitlement))) + result = await flow.check() + + assert flow.fallback.called is True + assert result.allowed is True + assert result.reservation is None + assert flow.wire.acquire_calls == [] + assert flow.engine.gate_call is None + + async def test_denies_and_refunds_when_the_gate_denies(self, clock: VirtualClock) -> None: + flow = make_flow(clock, engine=FlowEngine(gate=_verdict(False, "denied_by_targeting"))) + result = await flow.check() + + assert result.allowed is False + assert result.reason == "denied_by_targeting" + assert result.reservation is None + assert await flow.remaining() == LEASE_SIZE + assert await flow.reservations.count() == 0 + + async def test_fails_closed_when_the_acquire_fails(self, clock: VirtualClock) -> None: + flow = make_flow(clock, acquire="error") + result = await flow.check(on_acquire_failure="fail-closed") + + assert result.allowed is False + assert result.reason == "lease_acquire_failed" + assert result.error == "lease_acquire_failed" + assert result.reservation is None + assert flow.engine.gate_call is None + + async def test_defaults_to_fail_closed(self, clock: VirtualClock) -> None: + flow = make_flow(clock, acquire="error") + result = await flow.check() + + assert result.allowed is False + assert result.reservation is None + + async def test_fail_open_re_evaluates_with_the_balance_assumed_sufficient(self, clock: VirtualClock) -> None: + flow = make_flow(clock, acquire="error") + result = await flow.check(on_acquire_failure="fail-open") + + assert result.allowed is True + assert result.error == "lease_acquire_failed" + assert result.reservation is None + # Fail-open runs the rules, it does not skip them: only the credit + # balance is assumed sufficient, and the caller's usage still rides in. + assert flow.engine.balance(1) == FAIL_OPEN_BALANCE + preflight = flow.engine.calls[1]["options"].event_usage + assert (preflight.event_subtype, preflight.quantity) == (EVENT_SUBTYPE, 50) + + async def test_fail_open_still_denies_a_company_the_rules_do_not_entitle(self, clock: VirtualClock) -> None: + flow = make_flow( + clock, + engine=FlowEngine(probe=_verdict(False, "no matching rule", CREDIT_ENTITLEMENT)), + acquire="error", + ) + result = await flow.check(on_acquire_failure="fail-open") + + assert result.allowed is False + assert result.reason == "no matching rule (lease_acquire_failed_fail_open)" + assert result.error == "lease_acquire_failed" + + async def test_fail_open_allows_outright_when_the_re_evaluation_errors(self, clock: VirtualClock) -> None: + # The probe resolves the credit, the acquire fails, and then the + # fail-open re-evaluation throws, leaving only the blanket allow. + flow = make_flow(clock, engine=_ProbeThenExplodes(), acquire="error") + result = await flow.check(on_acquire_failure="fail-open") + + assert result.allowed is True + assert result.reason == "lease_acquire_failed_fail_open" + assert result.reservation is None + + async def test_falls_back_when_the_probe_errors(self, clock: VirtualClock) -> None: + # A probe failure is a resolution miss, not the gate: the plain check + # has its own degradation, and no lease is acquired. + flow = make_flow(clock, engine=FlowEngine(probe_error=RuntimeError("wasm exploded"))) + result = await flow.check() + + assert flow.fallback.called is True + assert result.allowed is True + assert flow.wire.acquire_calls == [] + + async def test_falls_back_when_the_flag_is_not_cached(self, clock: VirtualClock) -> None: + flow = make_flow(clock, missing_flag=True) + result = await flow.check() + + assert flow.fallback.called is True + assert result.reservation is None + assert flow.engine.calls == [] + + async def test_zero_usage_falls_back_without_a_hold(self, clock: VirtualClock) -> None: + flow = make_flow(clock) + result = await flow.check(usage=0) + + assert flow.fallback.called is True + assert result.reservation is None + assert flow.wire.acquire_calls == [] + assert flow.engine.calls == [] + + async def test_falls_back_when_nothing_names_the_event_subtype(self, clock: VirtualClock) -> None: + # The hold settles into a track event named by the subtype; without one + # it could be consumed while billing nothing. + entitlement = RulesengineFeatureEntitlement( + feature_id="feat", + feature_key=FLAG_KEY, + value_type="credit", + credit_id=CREDIT_ID, + consumption_rate=10, + ) + flow = make_flow(clock, engine=FlowEngine(probe=_verdict(True, "probe", entitlement))) + result = await flow.check(event_subtype=None) + + assert flow.fallback.called is True + assert result.reservation is None + assert flow.wire.acquire_calls == [] + + async def test_rejects_a_nan_usage_without_touching_the_lease(self, clock: VirtualClock) -> None: + # An unguarded NaN debit poisons the shared lease balance into + # approving every later reserve, since NaN loses every comparison. + flow = make_flow(clock) + denied = await flow.check(usage=float("nan")) + + assert denied.allowed is False + assert denied.error == "invalid_usage" + assert denied.reservation is None + assert flow.wire.acquire_calls == [] + assert flow.engine.calls == [] + + allowed = await flow.check() + assert allowed.allowed is True + assert allowed.reservation is not None + assert allowed.reservation.credits_reserved == 500 + + async def test_resolves_an_invalid_usage_through_fail_open(self, clock: VirtualClock) -> None: + flow = make_flow(clock) + result = await flow.check(usage=-10, on_acquire_failure="fail-open") + + assert result.allowed is True + assert result.error == "invalid_usage" + assert result.reservation is None + assert flow.wire.acquire_calls == [] + + async def test_extends_once_and_retries_when_the_lease_is_short(self, clock: VirtualClock) -> None: + flow = make_flow(clock) + await flow.check() # draws the lease down to 500 + flow.wire.extend_responses.append( + {"lease": {"granted_total": 2000, "expires_at": clock() + LEASE_DURATION}} + ) + result = await flow.check(usage=90) # 900 credits, more than the 500 left + + assert result.allowed is True + assert result.reservation is not None + assert len(flow.wire.extend_calls) == 1 + assert await flow.remaining() == 600 + + async def test_denies_when_the_retry_after_a_failed_extend_is_still_short(self, clock: VirtualClock) -> None: + flow = make_flow(clock) + await flow.check() + flow.wire.extend_responses.append({"error": "wire down"}) + result = await flow.check(usage=90) + + assert result.allowed is False + assert result.reason == "insufficient_lease_balance" + assert result.error == "insufficient_lease_balance" + assert await flow.remaining() == 500 + assert await flow.reservations.count() == 1 + + +class TestEntityResolution: + async def test_threads_a_resolved_user_into_both_evaluations(self, clock: VirtualClock) -> None: + user = RulesengineUser( + id="user_1", account_id="acc_1", environment_id="env_1", keys={"id": "user_1"}, traits=[], rules=[] + ) + flow = make_flow(clock, user=user) + result = await check_with_lease( + flow.deps, + FLAG_KEY, + COMPANY, + {"id": "user_1"}, + CheckOptions(usage=50, event_subtype=EVENT_SUBTYPE), + flow.fallback, + ) + await flow.manager._drain_background() + + assert result.allowed is True + # Evaluating without the named user would silently skip user-targeted + # rules and overrides. + assert [call["user"] for call in flow.engine.calls] == [user, user] + + async def test_falls_back_when_the_user_cannot_be_resolved(self, clock: VirtualClock) -> None: + flow = make_flow(clock, user_error=RuntimeError("DataStream client is not connected")) + result = await check_with_lease( + flow.deps, + FLAG_KEY, + COMPANY, + {"id": "user_1"}, + CheckOptions(usage=50, event_subtype=EVENT_SUBTYPE), + flow.fallback, + ) + + assert flow.fallback.called is True + assert result.reservation is None + assert flow.wire.acquire_calls == [] + + async def test_falls_back_when_the_company_cannot_be_resolved(self, clock: VirtualClock) -> None: + flow = make_flow(clock, company_error=RuntimeError("DataStream client is not connected")) + result = await flow.check() + + assert flow.fallback.called is True + assert result.reservation is None + assert flow.wire.acquire_calls == [] + + async def test_falls_back_without_datastream(self, clock: VirtualClock) -> None: + flow = make_flow(clock) + flow.deps.datastream = None + result = await flow.check() + + assert flow.fallback.called is True + assert result.reservation is None + + +class TestStoreFailureContainment: + async def test_a_dead_store_resolves_fail_closed_rather_than_raising(self, clock: VirtualClock) -> None: + flow = make_flow(clock, lease_store=UnreachableLeaseStore(clock=clock)) + result = await flow.check() + + assert result.allowed is False + assert result.reason == "lease_store_error" + assert result.error == "lease_store_error" + assert result.reservation is None + + async def test_a_dead_store_honors_fail_open(self, clock: VirtualClock) -> None: + flow = make_flow(clock, lease_store=UnreachableLeaseStore(clock=clock)) + result = await flow.check(on_acquire_failure="fail-open") + + assert result.allowed is True + assert result.reservation is None + assert flow.engine.balance(1) == FAIL_OPEN_BALANCE + + +class TestCrashWindow: + async def test_the_debit_lands_before_the_record(self, clock: VirtualClock) -> None: + """A crash in the gap leaks the debit; it never leaves a record with no + debit, which a later consume would refund into a double spend.""" + flow = make_flow(clock) + seen: Dict[str, Any] = {} + + async def freeze(reservation: ReservationRecord) -> None: + # Freeze the flow where a process death would: reached, never done, + # so neither the persist nor the undo runs. + seen["record"] = reservation + seen["remaining"] = await flow.leases.get(COMPANY["id"], CREDIT_ID) + seen["reserved"] = await flow.reservations.reserved_credits(COMPANY["id"], CREDIT_ID) + raise asyncio.CancelledError() + + flow.reservations.add = freeze # type: ignore[method-assign] + with pytest.raises(asyncio.CancelledError): + await flow.check() + + assert seen["record"].credits_reserved == 500 + assert seen["remaining"].local_remaining_credits == 500 + # Nothing the sweeper could refund: the leak is bounded by this one + # hold and reclaimed when the lease expires server-side. + assert seen["reserved"] == 0 + assert await flow.reservations.count() == 0 + assert flow.fallback.called is False + + +class TestSettle: + async def _reservation(self, flow: Flow) -> Reservation: + result = await flow.check() + assert result.reservation is not None + return result.reservation + + async def test_underuse_refunds_the_unspent_slice(self, clock: VirtualClock) -> None: + flow = make_flow(clock) + reservation = await self._reservation(flow) + outcome = await consume_reservation_and_build_event(flow.reservations, reservation, 20) + + assert outcome.settled_locally is True + assert outcome.track.event == EVENT_SUBTYPE + assert outcome.track.quantity == 20 + assert outcome.track.lease_id == "lse_1" + assert outcome.track.reservation_id is None + assert outcome.track.company == COMPANY + assert await flow.remaining() == 800 + + async def test_overuse_bills_the_actual_but_clamps_the_local_debit(self, clock: VirtualClock) -> None: + flow = make_flow(clock) + reservation = await self._reservation(flow) + outcome = await consume_reservation_and_build_event(flow.reservations, reservation, 120) + + # The server is the source of truth for real consumption, so the event + # bills the unclamped quantity; only the lease's own view is clamped. + assert outcome.track.quantity == 120 + assert await flow.remaining() == 500 + + async def test_a_settle_after_the_sweep_is_a_recovery_emit(self, clock: VirtualClock) -> None: + flow = make_flow(clock) + reservation = await self._reservation(flow) + clock.advance_ms(60_001) + assert await flow.reservations.sweep_expired() == 1 + assert await flow.remaining() == LEASE_SIZE + + outcome = await consume_reservation_and_build_event(flow.reservations, reservation, 20) + + # The hold was already refunded, so nothing re-debits the consumed + # slice and the local balance reads high until the lease rolls over. + assert outcome.settled_locally is False + assert outcome.track.quantity == 20 + assert outcome.track.lease_id == "lse_1" + assert await flow.remaining() == LEASE_SIZE + + async def test_a_second_settle_finds_nothing_to_claim(self, clock: VirtualClock) -> None: + flow = make_flow(clock) + reservation = await self._reservation(flow) + await consume_reservation_and_build_event(flow.reservations, reservation, 20) + outcome = await consume_reservation_and_build_event(flow.reservations, reservation, 20) + + assert outcome.settled_locally is False + assert await flow.remaining() == 800 + + async def test_a_fractional_settle_bills_a_whole_unit(self, clock: VirtualClock) -> None: + flow = make_flow(clock) + reservation = await self._reservation(flow) + outcome = await consume_reservation_and_build_event(flow.reservations, reservation, 0.5) + + assert outcome.track.quantity == 1 + + +class TestFlagCheckEvents: + async def test_an_allowed_check_reports_the_engine_verdict(self, clock: VirtualClock) -> None: + flow = make_flow(clock) + result = await flow.check() + + assert len(flow.flag_checks.events) == 1 + event = flow.flag_checks.events[0] + assert event.flag_key == FLAG_KEY + assert event.value is True + assert event.reason == result.reason + assert event.flag_id == "flag_1" + assert event.company_id == COMPANY["id"] + assert event.req_company == COMPANY + assert event.error is None + + async def test_a_denied_check_reports_the_denial(self, clock: VirtualClock) -> None: + flow = make_flow(clock, engine=FlowEngine(gate=_verdict(False, "denied_by_targeting"))) + await flow.check() + + assert len(flow.flag_checks.events) == 1 + assert flow.flag_checks.events[0].value is False + assert flow.flag_checks.events[0].reason == "denied_by_targeting" + + async def test_a_lease_failure_reports_once(self, clock: VirtualClock) -> None: + flow = make_flow(clock, acquire="error") + await flow.check(on_acquire_failure="fail-closed") + + assert len(flow.flag_checks.events) == 1 + assert flow.flag_checks.events[0].value is False + assert flow.flag_checks.events[0].error == "lease_acquire_failed" + + async def test_a_fallback_exit_reports_nothing(self, clock: VirtualClock) -> None: + # The plain check the lease path defers to reports its own. + flow = make_flow(clock) + await flow.check(usage=0) + + assert flow.fallback.called is True + assert flow.flag_checks.events == [] + + async def test_a_reporting_failure_leaves_the_verdict_alone(self, clock: VirtualClock) -> None: + flow = make_flow(clock) + flow.flag_checks.explode = True + result = await flow.check() + + assert result.allowed is True + assert result.reservation is not None + assert await flow.remaining() == 500 + + +class TestPerCheckTimeout: + async def test_the_timeout_reaches_the_acquire(self, clock: VirtualClock) -> None: + flow = make_flow(clock) + await flow.check(timeout=2.5) + + assert flow.wire.acquire_calls[0]["timeout"] == 2.5 + + async def test_the_timeout_reaches_the_extend_the_reserve_triggers(self, clock: VirtualClock) -> None: + flow = make_flow(clock) + await flow.check() # draws the lease down to 500 + flow.wire.extend_responses.append( + {"lease": {"granted_total": 2000, "expires_at": clock() + LEASE_DURATION}} + ) + await flow.check(usage=90, timeout=2.5) # 900 credits, more than the 500 left + + assert flow.wire.extend_calls[0]["timeout"] == 2.5 + + async def test_no_timeout_leaves_the_wire_calls_alone(self, clock: VirtualClock) -> None: + flow = make_flow(clock) + await flow.check() + + assert flow.wire.acquire_calls[0]["timeout"] is None diff --git a/tests/leases/test_crash_windows.py b/tests/leases/test_crash_windows.py new file mode 100644 index 00000000..43f9face --- /dev/null +++ b/tests/leases/test_crash_windows.py @@ -0,0 +1,207 @@ +"""The two bounded-leak windows, on both backends. + +A crash between the two steps of a transition must strand locally held credits +(which the server reclaims at lease expiry) rather than enable a double-spend. +The debit and the claim are durable first; the record and the refund are what +may be lost. +""" + +from __future__ import annotations + +from typing import Any, Awaitable, Tuple, cast + +import pytest +from lease_support import CrashingRefundLeaseStore, VirtualClock, make_fake_redis, make_reservation + +from schematic.leases import ( + InMemoryLeaseStore, + InMemoryReservationStore, + LeaseStore, + RedisLeaseStore, + RedisReservationStore, + ReservationStore, +) + +BACKENDS = ("in_memory", "redis") + + +@pytest.fixture(autouse=True) +def _frozen(frozen_clock: VirtualClock) -> VirtualClock: + """fakeredis reads TIME from the process clock, so the virtual clock is it.""" + return frozen_clock + + +def _make_stores( + backend: str, clock: VirtualClock +) -> Tuple[LeaseStore, ReservationStore, CrashingRefundLeaseStore]: + leases: LeaseStore + reservations: ReservationStore + if backend == "in_memory": + leases = InMemoryLeaseStore(clock=clock) + crash = CrashingRefundLeaseStore(leases) + reservations = InMemoryReservationStore(crash, clock=clock) + else: + client = make_fake_redis() + leases = RedisLeaseStore(client, clock=clock) + crash = CrashingRefundLeaseStore(leases) + reservations = RedisReservationStore(client, crash, clock=clock) + return leases, reservations, crash + + +async def _seed(leases: LeaseStore, clock: VirtualClock, ttl: float = 60) -> None: + await leases.replace( + lease_id="lse_1", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=clock() + ttl, + ) + + +async def _balance(leases: LeaseStore) -> float: + entry = await leases.get("co_1", "ct_1") + assert entry is not None + return entry.local_remaining_credits + + +@pytest.mark.parametrize("backend", BACKENDS) +async def test_debit_without_record_leaks_only_the_hold(backend: str, frozen_clock: VirtualClock) -> None: + leases, reservations, _crash = _make_stores(backend, frozen_clock) + await _seed(leases, frozen_clock) + + # The crash: the atomic debit landed, the reservation record never did. + assert await leases.try_reserve("co_1", "ct_1", 100) == 900 + + # Exactly the reserved amount is stranded, and it is invisible to the + # reservation table, so no sweep can ever refund it. + assert await _balance(leases) == 900 + assert await reservations.reserved_credits("co_1", "ct_1") == 0 + assert await reservations.sweep_expired(frozen_clock() + 3600) == 0 + assert await _balance(leases) == 900 + + +@pytest.mark.parametrize("backend", BACKENDS) +async def test_debit_leak_is_reclaimed_at_lease_expiry(backend: str, frozen_clock: VirtualClock) -> None: + leases, _reservations, _crash = _make_stores(backend, frozen_clock) + await _seed(leases, frozen_clock) + assert await leases.try_reserve("co_1", "ct_1", 100) == 900 + + frozen_clock.advance_ms(60_001) + # The stale balance is never served again, and the successor installs at + # the full grant: the leak does not outlive the lease. + assert await leases.try_reserve("co_1", "ct_1", 1) is None + assert await leases.replace( + lease_id="lse_2", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=frozen_clock() + 120, + ) + assert await _balance(leases) == 1000 + + +@pytest.mark.parametrize("backend", BACKENDS) +async def test_a_retried_check_settles_independently(backend: str, frozen_clock: VirtualClock) -> None: + leases, reservations, _crash = _make_stores(backend, frozen_clock) + await _seed(leases, frozen_clock, ttl=3600) + # The crashed attempt, then the retry with a fresh reservation. + assert await leases.try_reserve("co_1", "ct_1", 100) == 900 + assert await leases.try_reserve("co_1", "ct_1", 100) == 800 + await reservations.add(make_reservation(id="res_retry", expires_at=frozen_clock() + 60)) + + assert await reservations.consume("res_retry", 40) == 40 + # 1000 less the 100 leaked and the 40 consumed: the retry's unspent 60 came + # back exactly once, the leaked 100 stayed leaked. + assert await _balance(leases) == 860 + + assert await reservations.consume("res_retry", 40) is None + assert await reservations.sweep_expired(frozen_clock() + 3600) == 0 + assert await _balance(leases) == 860 + + +@pytest.mark.parametrize("backend", BACKENDS) +async def test_crash_before_refund_keeps_the_claim(backend: str, frozen_clock: VirtualClock) -> None: + leases, reservations, crash = _make_stores(backend, frozen_clock) + await _seed(leases, frozen_clock) + assert await leases.try_reserve("co_1", "ct_1", 100) == 900 + await reservations.add(make_reservation(expires_at=frozen_clock() + 60)) + + crash.arm() + with pytest.raises(RuntimeError, match="simulated crash before refund"): + await reservations.consume("res_1", 30) + + # The claim survived, so the reservation is gone everywhere and nothing can + # double-spend; the 70-credit refund is lost, bounded by the hold. + assert await reservations.get("res_1") is None + assert await reservations.reserved_credits("co_1", "ct_1") == 0 + assert await _balance(leases) == 900 + + # Neither a retried settle nor the sweeper can refund a claimed hold. + assert await reservations.consume("res_1", 30) is None + assert await reservations.sweep_expired(frozen_clock() + 3600) == 0 + assert await _balance(leases) == 900 + + +@pytest.mark.parametrize("backend", BACKENDS) +async def test_crash_before_refund_never_leaks_into_a_successor( + backend: str, frozen_clock: VirtualClock +) -> None: + leases, reservations, crash = _make_stores(backend, frozen_clock) + await _seed(leases, frozen_clock) + assert await leases.try_reserve("co_1", "ct_1", 100) == 900 + await reservations.add(make_reservation(expires_at=frozen_clock() + 60)) + + crash.arm() + with pytest.raises(RuntimeError, match="simulated crash before refund"): + await reservations.consume("res_1", 30) + + frozen_clock.advance_ms(60_001) + assert await leases.try_reserve("co_1", "ct_1", 1) is None + assert await leases.replace( + lease_id="lse_2", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=frozen_clock() + 120, + ) + # A very late retried settle must not push the lost refund into lse_2. + assert await reservations.consume("res_1", 0) is None + assert await _balance(leases) == 1000 + + +async def test_crash_after_the_claim_is_reconciled_without_a_refund(frozen_clock: VirtualClock) -> None: + # Redis only: a process death after the claim but before the index cleanup + # leaves the per-tenant index over-counting. The sweeper reconciles it, and + # deliberately does not refund: without the reservation hash, exactly-once + # cannot be arbitrated across racing sweepers. + client = make_fake_redis() + leases = RedisLeaseStore(client, clock=frozen_clock) + reservations = RedisReservationStore(client, leases, clock=frozen_clock) + await _seed(leases, frozen_clock) + assert await leases.try_reserve("co_1", "ct_1", 100) == 900 + await reservations.add(make_reservation(expires_at=frozen_clock() - 0.001)) + + original_evalsha = client.evalsha + armed = True + + async def crash_after_claim(sha: str, numkeys: int, *args: Any) -> Any: + nonlocal armed + result = await cast(Awaitable[Any], original_evalsha(sha, numkeys, *args)) + if armed and isinstance(result, list): + armed = False + raise RuntimeError("simulated crash after claim") + return result + + client.evalsha = crash_after_claim # type: ignore[method-assign] + + with pytest.raises(RuntimeError, match="simulated crash after claim"): + await reservations.consume("res_1", 30) + client.evalsha = original_evalsha # type: ignore[method-assign] + + # The orphaned index field still counts the hold... + assert await reservations.reserved_credits("co_1", "ct_1") == 100 + # ...until the sweeper reconciles it, without refunding. + assert await reservations.sweep_expired() == 0 + assert await reservations.reserved_credits("co_1", "ct_1") == 0 + assert await reservations.count() == 0 + assert await _balance(leases) == 900 diff --git a/tests/leases/test_lease_manager.py b/tests/leases/test_lease_manager.py new file mode 100644 index 00000000..a4c0dbfc --- /dev/null +++ b/tests/leases/test_lease_manager.py @@ -0,0 +1,404 @@ +"""Lease manager behavior, ported from the Node SDK's manager tests. + +The cross-pod cases share one Redis between two managers, which is the shape +the store's convergence rules exist for. +""" + +from __future__ import annotations + +import asyncio +from typing import Any, List, Optional + +import pytest +from lease_support import ScriptedWireClient, VirtualClock, make_fake_redis + +from schematic.leases import ( + InMemoryLeaseStore, + InMemoryReservationStore, + LeaseConfig, + LeaseGrant, + LeaseManager, + LeaseState, + LeaseStore, + RedisLeaseStore, +) + +CONFIG = LeaseConfig(lease_duration=300, reservation_ttl=60, lease_size=1000, low_water_mark=0.25) + + +@pytest.fixture(autouse=True) +def _frozen(frozen_clock: VirtualClock) -> VirtualClock: + """The shared-store cases run against fakeredis, which reads TIME from the + process clock, so the virtual clock has to be that clock here.""" + return frozen_clock + + +def _lease(clock: VirtualClock, lease_id: str = "lse_1", granted: float = 1000, ttl: float = 300) -> dict: + return {"lease": {"lease_id": lease_id, "granted_amount": granted, "expires_at": clock() + ttl}} + + +def _make_manager(clock: VirtualClock) -> tuple[LeaseManager, InMemoryLeaseStore, ScriptedWireClient]: + store = InMemoryLeaseStore(clock=clock) + wire = ScriptedWireClient() + manager = LeaseManager(wire, store, config=CONFIG, clock=clock) + return manager, store, wire + + +async def test_acquire_installs_the_lease(clock: VirtualClock) -> None: + manager, store, wire = _make_manager(clock) + wire.acquire_responses.append(_lease(clock)) + + entry = await manager.acquire_if_needed("co_1", "ct_1") + assert len(wire.acquire_calls) == 1 + assert wire.acquire_calls[0]["requested_amount"] == 1000 + assert entry is not None and entry.lease_id == "lse_1" and entry.local_remaining_credits == 1000 + + +async def test_acquire_reuses_a_live_lease(clock: VirtualClock) -> None: + manager, _store, wire = _make_manager(clock) + wire.acquire_responses.append(_lease(clock)) + await manager.acquire_if_needed("co_1", "ct_1") + await manager.acquire_if_needed("co_1", "ct_1") + assert len(wire.acquire_calls) == 1 + + +async def test_acquire_is_single_flight(clock: VirtualClock) -> None: + manager, _store, wire = _make_manager(clock) + gate: "asyncio.Future[None]" = asyncio.get_running_loop().create_future() + original = wire.acquire + + async def slow_acquire(*args: Any, **kwargs: Any) -> LeaseGrant: + await gate + return await original(*args, **kwargs) + + wire.acquire = slow_acquire # type: ignore[method-assign] + wire.acquire_responses.append(_lease(clock)) + + pending = [asyncio.ensure_future(manager.acquire_if_needed("co_1", "ct_1")) for _ in range(3)] + await asyncio.sleep(0) + gate.set_result(None) + results = await asyncio.gather(*pending) + + assert len(wire.acquire_calls) == 1 + assert [r.lease_id for r in results if r] == ["lse_1"] * 3 + + +async def test_acquire_replaces_an_expired_slot_without_releasing(clock: VirtualClock) -> None: + manager, store, wire = _make_manager(clock) + await store.replace( + lease_id="lse_stale", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=clock() - 1, + ) + wire.acquire_responses.append(_lease(clock, "lse_fresh")) + + entry = await manager.acquire_if_needed("co_1", "ct_1") + await manager._drain_background() + assert entry is not None and entry.lease_id == "lse_fresh" and entry.local_remaining_credits == 1000 + # `replace` wrote rather than keeping a live lease, so nothing is redundant. + assert wire.release_calls == [] + + +async def test_extend_fires_below_the_water_mark(clock: VirtualClock) -> None: + manager, store, wire = _make_manager(clock) + wire.acquire_responses.append(_lease(clock)) + await manager.acquire_if_needed("co_1", "ct_1") + await store.try_reserve("co_1", "ct_1", 800) + + wire.extend_responses.append({"lease": {"granted_total": 2000, "expires_at": clock() + 600}}) + await manager.maybe_extend("co_1", "ct_1") + assert len(wire.extend_calls) == 1 + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.granted_amount == 2000 and entry.local_remaining_credits == 1200 + + +async def test_extend_fires_for_required_credits_above_the_water_mark(clock: VirtualClock) -> None: + manager, store, wire = _make_manager(clock) + wire.acquire_responses.append(_lease(clock)) + await manager.acquire_if_needed("co_1", "ct_1") + await store.try_reserve("co_1", "ct_1", 100) + + # Without the hint this is a no-op: 900/1000 is far above the water mark. + await manager.maybe_extend("co_1", "ct_1") + assert wire.extend_calls == [] + + wire.extend_responses.append({"lease": {"granted_total": 2000, "expires_at": clock() + 600}}) + await manager.maybe_extend("co_1", "ct_1", 1500) + assert len(wire.extend_calls) == 1 + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.granted_amount == 2000 + + +async def test_extend_is_sized_to_the_shortfall(clock: VirtualClock) -> None: + manager, store, wire = _make_manager(clock) + wire.acquire_responses.append(_lease(clock)) + await manager.acquire_if_needed("co_1", "ct_1") + await store.try_reserve("co_1", "ct_1", 100) + + # A check needing 5000 credits has a 4100 shortfall, above the configured + # 1000 tranche: a tranche-sized extend would leave its retry failing + # forever however much balance the server has. + wire.extend_responses.append({"lease": {"granted_total": 5100, "expires_at": clock() + 600}}) + await manager.maybe_extend("co_1", "ct_1", 5000) + assert wire.extend_calls[-1]["additional_amount"] == 4100 + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.local_remaining_credits == 5000 + + +async def test_never_extends_an_expired_lease(clock: VirtualClock) -> None: + manager, store, wire = _make_manager(clock) + await store.replace( + lease_id="lse_old", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=clock() - 1, + ) + assert await manager.maybe_extend("co_1", "ct_1", 1500) is None + assert wire.extend_calls == [] + + +async def test_store_failures_resolve_to_no_lease(clock: VirtualClock) -> None: + class BrokenStore(InMemoryLeaseStore): + async def get(self, company_id: str, credit_type_id: str) -> Optional[LeaseState]: + raise RuntimeError("redis down") + + wire = ScriptedWireClient() + manager = LeaseManager(wire, BrokenStore(clock=clock), config=CONFIG, clock=clock) + # Both are often called fire-and-forget, where a raised exception would + # surface as an unretrieved task exception. + assert await manager.acquire_if_needed("co_1", "ct_1") is None + assert await manager.maybe_extend("co_1", "ct_1") is None + assert wire.acquire_calls == [] + assert wire.extend_calls == [] + + +async def test_wire_failures_resolve_to_no_lease(clock: VirtualClock) -> None: + manager, store, wire = _make_manager(clock) + wire.acquire_responses.append({"error": "wire down"}) + assert await manager.acquire_if_needed("co_1", "ct_1") is None + assert await store.get("co_1", "ct_1") is None + + await store.replace( + lease_id="lse_1", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=clock() + 300, + ) + await store.try_reserve("co_1", "ct_1", 800) + wire.extend_responses.append({"error": "wire down"}) + assert await manager.maybe_extend("co_1", "ct_1") is None + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.granted_amount == 1000 and entry.local_remaining_credits == 200 + + +async def test_release_all_releases_live_and_skips_expired(clock: VirtualClock) -> None: + manager, store, wire = _make_manager(clock) + await store.replace( + lease_id="lse_live", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=clock() + 60, + ) + await store.replace( + lease_id="lse_expired", + company_id="co_2", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=clock() - 1, + ) + await manager.release_all_local_leases() + assert wire.release_calls == ["lse_live"] + # The released lease is dropped locally; the expired one is left to lazy + # expiry. + assert await store.get("co_1", "ct_1") is None + assert await store.get("co_2", "ct_1") is not None + + +async def test_release_all_skips_a_shared_store(clock: VirtualClock) -> None: + # A shared backend cannot enumerate: sibling pods still draw on its leases. + client = make_fake_redis() + store = RedisLeaseStore(client, clock=clock) + wire = ScriptedWireClient() + manager = LeaseManager(wire, store, config=CONFIG, clock=clock) + await store.replace( + lease_id="lse_shared", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=clock() + 60, + ) + await manager.release_all_local_leases() + assert wire.release_calls == [] + + +async def test_acquire_and_extend_do_not_share_inflight_state(clock: VirtualClock) -> None: + manager, store, wire = _make_manager(clock) + gate: "asyncio.Future[None]" = asyncio.get_running_loop().create_future() + original_extend = wire.extend + + async def slow_extend(*args: Any, **kwargs: Any) -> LeaseGrant: + await gate + return await original_extend(*args, **kwargs) + + wire.extend = slow_extend # type: ignore[method-assign] + wire.extend_responses.append({"lease": {"granted_total": 2000, "expires_at": clock() + 600}}) + wire.acquire_responses.append(_lease(clock, "lse_fresh")) + + await store.replace( + lease_id="lse_live", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=clock() + 300, + ) + await store.try_reserve("co_1", "ct_1", 800) + extending = asyncio.ensure_future(manager.maybe_extend("co_1", "ct_1")) + await asyncio.sleep(0) + + await store.drop("co_1", "ct_1") + acquired = await manager.acquire_if_needed("co_1", "ct_1") + assert acquired is not None and acquired.lease_id == "lse_fresh" + assert len(wire.acquire_calls) == 1 + + gate.set_result(None) + await extending + + +async def test_lost_acquire_race_releases_the_redundant_lease(clock: VirtualClock) -> None: + # Two managers on one Redis, as two pods would be. Both see an empty slot + # and acquire; only one lease can hold the slot, and the loser must release + # the one it minted rather than orphan it against the company balance. + shared: LeaseStore = RedisLeaseStore(make_fake_redis(), clock=clock) + pods = [] + for lease_id in ("lse_a", "lse_b"): + wire = ScriptedWireClient() + wire.acquire_responses.append(_lease(clock, lease_id)) + pods.append((LeaseManager(wire, shared, config=CONFIG, clock=clock), wire)) + + entries = await asyncio.gather(*(manager.acquire_if_needed("co_1", "ct_1") for manager, _ in pods)) + for manager, _ in pods: + await manager._drain_background() + + survivor = await shared.get("co_1", "ct_1") + assert survivor is not None and survivor.lease_id in ("lse_a", "lse_b") + assert [entry.lease_id for entry in entries if entry] == [survivor.lease_id] * 2 + + released: List[str] = [lease_id for _, wire in pods for lease_id in wire.release_calls] + loser = "lse_b" if survivor.lease_id == "lse_a" else "lse_a" + assert released == [loser] + + +async def test_lost_acquire_race_with_the_same_lease_releases_nothing(clock: VirtualClock) -> None: + # The server is idempotent for an active slot, so a racing acquire is + # handed back the SAME lease the sibling installed. Releasing it would pull + # the shared lease out from under every pod. + shared: LeaseStore = RedisLeaseStore(make_fake_redis(), clock=clock) + pods = [] + for _ in range(2): + wire = ScriptedWireClient() + wire.acquire_responses.append(_lease(clock, "lse_shared")) + pods.append((LeaseManager(wire, shared, config=CONFIG, clock=clock), wire)) + + entries = await asyncio.gather(*(manager.acquire_if_needed("co_1", "ct_1") for manager, _ in pods)) + for manager, _ in pods: + await manager._drain_background() + + assert [entry.lease_id for entry in entries if entry] == ["lse_shared", "lse_shared"] + assert [lease_id for _, wire in pods for lease_id in wire.release_calls] == [] + + +async def test_uncontended_acquire_releases_nothing(clock: VirtualClock) -> None: + manager, _store, wire = _make_manager(clock) + wire.acquire_responses.append(_lease(clock)) + await manager.acquire_if_needed("co_1", "ct_1") + await manager._drain_background() + assert wire.release_calls == [] + + +async def test_sweep_loop_runs_and_stops(clock: VirtualClock) -> None: + leases = InMemoryLeaseStore(clock=clock) + reservations = InMemoryReservationStore(leases, clock=clock) + manager = LeaseManager( + ScriptedWireClient(), + leases, + reservation_store=reservations, + # Short enough that the test does not idle, long enough that the loop + # cannot run twice before it is stopped. + config=LeaseConfig(sweep_interval=0.01), + clock=clock, + ) + swept: List[int] = [] + original = reservations.sweep_expired + + async def counting_sweep(now: Optional[float] = None) -> int: + result = await original(now) + swept.append(result) + return result + + reservations.sweep_expired = counting_sweep # type: ignore[method-assign] + + manager.start_sweep() + manager.start_sweep() # idempotent + await asyncio.sleep(0.03) + assert swept + manager.stop() + ticks = len(swept) + await asyncio.sleep(0.03) + assert len(swept) == ticks + + +async def test_sweep_loop_survives_a_failing_sweep(clock: VirtualClock) -> None: + leases = InMemoryLeaseStore(clock=clock) + reservations = InMemoryReservationStore(leases, clock=clock) + manager = LeaseManager( + ScriptedWireClient(), + leases, + reservation_store=reservations, + config=LeaseConfig(sweep_interval=0.01), + clock=clock, + ) + attempts: List[int] = [] + + async def failing_sweep(now: Optional[float] = None) -> int: + attempts.append(1) + raise RuntimeError("redis blip") + + reservations.sweep_expired = failing_sweep # type: ignore[method-assign] + manager.start_sweep() + await asyncio.sleep(0.05) + manager.stop() + # A transient failure must not kill the loop: the next tick retries. + assert len(attempts) > 1 + + +async def test_resolve_config_applies_overrides(clock: VirtualClock) -> None: + from schematic.leases import LeaseConfigOverride + + config = LeaseConfig( + lease_size=500, + overrides={"ct_special": LeaseConfigOverride(lease_size=25, low_water_mark=0.5)}, + ) + manager = LeaseManager(ScriptedWireClient(), InMemoryLeaseStore(clock=clock), config=config, clock=clock) + + plain = manager.resolve_config("ct_1") + assert plain.lease_size == 500 + assert plain.low_water_mark == 0.25 + assert plain.lease_duration == 300 + + special = manager.resolve_config("ct_special") + assert special.lease_size == 25 + assert special.low_water_mark == 0.5 + + +def test_resolve_lease_config_leaves_the_reservation_ttl_alone() -> None: + from schematic.leases import resolve_lease_config + + resolved = resolve_lease_config(LeaseConfig(reservation_ttl=7200.0), None, "ct_1") + # A client-mode TTL never reaches the server, so the server's cap does not + # apply to it: it only tells the local sweeper when to refund a hold. + assert resolved.reservation_ttl == 7200.0 diff --git a/tests/leases/test_lease_store.py b/tests/leases/test_lease_store.py new file mode 100644 index 00000000..574a65dd --- /dev/null +++ b/tests/leases/test_lease_store.py @@ -0,0 +1,229 @@ +"""In-memory lease slot semantics, ported from the Node SDK's store tests.""" + +from __future__ import annotations + +import asyncio +import math + +import pytest +from lease_support import VirtualClock + +from schematic.leases import InMemoryLeaseStore + + +async def _seed(store: InMemoryLeaseStore, clock: VirtualClock, *, granted: float = 100, ttl: float = 60) -> None: + await store.replace( + lease_id="lse_1", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=granted, + expires_at=clock() + ttl, + ) + + +@pytest.fixture +def store(clock: VirtualClock) -> InMemoryLeaseStore: + return InMemoryLeaseStore(clock=clock) + + +async def test_replace_installs_at_the_full_grant(store: InMemoryLeaseStore, clock: VirtualClock) -> None: + await _seed(store, clock) + entry = await store.get("co_1", "ct_1") + assert entry is not None + assert entry.granted_amount == 100 + assert entry.local_remaining_credits == 100 + + +async def test_try_reserve_returns_the_post_debit_balance(store: InMemoryLeaseStore, clock: VirtualClock) -> None: + await _seed(store, clock) + assert await store.try_reserve("co_1", "ct_1", 30) == 70 + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.local_remaining_credits == 70 + + +async def test_try_reserve_refuses_without_debiting(store: InMemoryLeaseStore, clock: VirtualClock) -> None: + await _seed(store, clock) + assert await store.try_reserve("co_1", "ct_1", 150) is None + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.local_remaining_credits == 100 + + +async def test_try_reserve_refuses_an_expired_lease(store: InMemoryLeaseStore, clock: VirtualClock) -> None: + await _seed(store, clock) + clock.advance_ms(60_001) + assert await store.try_reserve("co_1", "ct_1", 10) is None + # The balance is stale, not spendable: the server released the lease and + # refunded its remainder. + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.local_remaining_credits == 100 + + +async def test_try_reserve_rejects_nan_and_infinity(store: InMemoryLeaseStore, clock: VirtualClock) -> None: + await _seed(store, clock) + # NaN slips through every comparison, so an unguarded debit would set the + # balance to NaN and approve every later reserve. + assert await store.try_reserve("co_1", "ct_1", math.nan) is None + assert await store.try_reserve("co_1", "ct_1", -10) is None + assert await store.try_reserve("co_1", "ct_1", math.inf) is None + assert await store.try_reserve("co_1", "ct_1", 30) == 70 + assert await store.try_reserve("co_1", "ct_1", 80) is None + + +async def test_refund_caps_at_the_granted_amount(store: InMemoryLeaseStore, clock: VirtualClock) -> None: + await _seed(store, clock) + await store.try_reserve("co_1", "ct_1", 30) + await store.refund("co_1", "ct_1", 20) + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.local_remaining_credits == 90 + await store.refund("co_1", "ct_1", 9999) + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.local_remaining_credits == 100 + + +async def test_refund_pinned_to_a_stale_lease_is_dropped(store: InMemoryLeaseStore, clock: VirtualClock) -> None: + await store.replace( + lease_id="lse_b", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=100, + expires_at=clock() + 60, + ) + await store.try_reserve("co_1", "ct_1", 50) + await store.refund("co_1", "ct_1", 30, "lse_a") + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.local_remaining_credits == 50 + await store.refund("co_1", "ct_1", 30, "lse_b") + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.local_remaining_credits == 80 + + +async def test_concurrent_try_reserves_serialize_per_slot(store: InMemoryLeaseStore, clock: VirtualClock) -> None: + # Per-slot atomicity: three concurrent 40-credit reserves against a + # 100-credit lease must not oversell it. + await _seed(store, clock) + results = await asyncio.gather( + store.try_reserve("co_1", "ct_1", 40), + store.try_reserve("co_1", "ct_1", 40), + store.try_reserve("co_1", "ct_1", 40), + ) + assert sorted(r for r in results if r is not None) == [20, 60] + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.local_remaining_credits == 20 + + +async def test_extend_reconciles_to_the_server_total(store: InMemoryLeaseStore, clock: VirtualClock) -> None: + await _seed(store, clock, ttl=10) + await store.try_reserve("co_1", "ct_1", 30) + new_expiry = clock() + 60 + await store.extend("co_1", "ct_1", 150, new_expiry) + entry = await store.get("co_1", "ct_1") + assert entry is not None + assert entry.granted_amount == 150 + assert entry.local_remaining_credits == 120 + assert entry.expires_at == new_expiry + + +async def test_stale_extend_total_is_a_no_op(store: InMemoryLeaseStore, clock: VirtualClock) -> None: + await _seed(store, clock, ttl=10) + far_expiry = clock() + 120 + await store.extend("co_1", "ct_1", 200, far_expiry) + await store.extend("co_1", "ct_1", 150, clock() + 60) + entry = await store.get("co_1", "ct_1") + assert entry is not None + assert entry.granted_amount == 200 + assert entry.local_remaining_credits == 200 + assert entry.expires_at == far_expiry + + +async def test_extend_pinned_to_a_stale_lease_is_dropped(store: InMemoryLeaseStore, clock: VirtualClock) -> None: + await store.replace( + lease_id="lse_b", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=100, + expires_at=clock() + 60, + ) + await store.extend("co_1", "ct_1", 150, clock() + 120, "lse_a") + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.granted_amount == 100 and entry.local_remaining_credits == 100 + await store.extend("co_1", "ct_1", 150, clock() + 120, "lse_b") + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.granted_amount == 150 and entry.local_remaining_credits == 150 + + +async def test_drop_removes_the_slot(store: InMemoryLeaseStore, clock: VirtualClock) -> None: + await _seed(store, clock) + await store.drop("co_1", "ct_1") + assert await store.get("co_1", "ct_1") is None + + +async def test_replace_keeps_a_live_lease_with_a_different_id( + store: InMemoryLeaseStore, clock: VirtualClock +) -> None: + assert await store.replace( + lease_id="lse_1", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=100, + expires_at=clock() + 60, + ) + await store.try_reserve("co_1", "ct_1", 40) + wrote = await store.replace( + lease_id="lse_2", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=100, + expires_at=clock() + 60, + ) + assert wrote is False + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.lease_id == "lse_1" and entry.local_remaining_credits == 60 + + +async def test_replace_overwrites_an_expired_lease(store: InMemoryLeaseStore, clock: VirtualClock) -> None: + await _seed(store, clock, ttl=60) + clock.advance_ms(60_001) + wrote = await store.replace( + lease_id="lse_2", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=100, + expires_at=clock() + 60, + ) + assert wrote is True + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.lease_id == "lse_2" and entry.local_remaining_credits == 100 + + +async def test_replace_reconciles_an_expired_same_id_lease(store: InMemoryLeaseStore, clock: VirtualClock) -> None: + # A stale acquire response can hand back the lease already installed (the + # server is idempotent for an active slot) after the local row expired. + # Rewriting would reset the balance and erase debits whose reservations are + # still open. + await _seed(store, clock, ttl=60) + await store.try_reserve("co_1", "ct_1", 40) + clock.advance_ms(61_000) + later_expiry = clock() + 60 + wrote = await store.replace( + lease_id="lse_1", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=150, + expires_at=later_expiry, + ) + assert wrote is False + entry = await store.get("co_1", "ct_1") + assert entry is not None + assert entry.granted_amount == 150 + assert entry.local_remaining_credits == 110 + assert entry.expires_at == later_expiry + + +async def test_get_returns_a_snapshot(store: InMemoryLeaseStore, clock: VirtualClock) -> None: + # Mutating what `get` handed back must not reach the stored slot. + await _seed(store, clock) + entry = await store.get("co_1", "ct_1") + assert entry is not None + entry.local_remaining_credits = 0 + stored = await store.get("co_1", "ct_1") + assert stored is not None and stored.local_remaining_credits == 100 diff --git a/tests/leases/test_redis_lease_store.py b/tests/leases/test_redis_lease_store.py new file mode 100644 index 00000000..d184b553 --- /dev/null +++ b/tests/leases/test_redis_lease_store.py @@ -0,0 +1,298 @@ +"""Redis lease slot semantics, ported from the Node SDK's Redis store tests. + +The shared-backend cases here are the ones the vectors cannot express: two pods +on one Redis, and rows that are expired but not yet evicted. +""" + +from __future__ import annotations + +import asyncio +import math +from typing import Any + +import pytest +from lease_support import VirtualClock + +from schematic.leases import RedisLeaseStore +from schematic.leases.redis_lease_store import LEASE_TTL_GRACE_MS + + +@pytest.fixture +def store(redis_client: Any, frozen_clock: VirtualClock) -> RedisLeaseStore: + return RedisLeaseStore(redis_client, clock=frozen_clock) + + +async def _seed(store: RedisLeaseStore, clock: VirtualClock, *, granted: float = 100, ttl: float = 60) -> bool: + return await store.replace( + lease_id="lse_1", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=granted, + expires_at=clock() + ttl, + ) + + +async def test_replace_installs_a_fresh_lease(store: RedisLeaseStore, frozen_clock: VirtualClock) -> None: + assert await _seed(store, frozen_clock) is True + entry = await store.get("co_1", "ct_1") + assert entry is not None + assert entry.lease_id == "lse_1" + assert entry.granted_amount == 100 + assert entry.local_remaining_credits == 100 + + +async def test_replace_preserves_debits_when_the_same_live_lease_is_rewritten( + store: RedisLeaseStore, frozen_clock: VirtualClock +) -> None: + await _seed(store, frozen_clock, granted=1000) + await store.try_reserve("co_1", "ct_1", 400) + # A second pod acquires and is handed the same lease back. + assert await _seed(store, frozen_clock, granted=1000) is False + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.local_remaining_credits == 600 + + +async def test_replace_reconciles_an_expired_same_id_lease( + store: RedisLeaseStore, frozen_clock: VirtualClock +) -> None: + await _seed(store, frozen_clock, granted=1000) + await store.try_reserve("co_1", "ct_1", 400) + frozen_clock.advance_ms(61_000) + later_expiry = frozen_clock() + 60 + wrote = await store.replace( + lease_id="lse_1", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1500, + expires_at=later_expiry, + ) + assert wrote is False + entry = await store.get("co_1", "ct_1") + assert entry is not None + assert entry.granted_amount == 1500 + assert entry.local_remaining_credits == 1100 + assert entry.expires_at == later_expiry + + +async def test_replace_keeps_a_different_live_lease(store: RedisLeaseStore, frozen_clock: VirtualClock) -> None: + # Two pods race the first acquire: the loser must not clobber the winner's + # already-debited balance. + await store.replace( + lease_id="lse_winner", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=frozen_clock() + 60, + ) + await store.try_reserve("co_1", "ct_1", 400) + wrote = await store.replace( + lease_id="lse_loser", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=frozen_clock() + 60, + ) + assert wrote is False + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.lease_id == "lse_winner" and entry.local_remaining_credits == 600 + + +async def test_try_reserve_gates_the_shared_balance(store: RedisLeaseStore, frozen_clock: VirtualClock) -> None: + # Three concurrent 40-credit reserves against a 100-credit lease: the Lua + # check-and-debit must not oversell it. + await _seed(store, frozen_clock) + results = await asyncio.gather( + store.try_reserve("co_1", "ct_1", 40), + store.try_reserve("co_1", "ct_1", 40), + store.try_reserve("co_1", "ct_1", 40), + ) + assert sorted(r for r in results if r is not None) == [20, 60] + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.local_remaining_credits == 20 + + +async def test_try_reserve_refuses_an_expired_but_unevicted_row( + store: RedisLeaseStore, frozen_clock: VirtualClock +) -> None: + await _seed(store, frozen_clock, ttl=60) + # Past the declared expiry but inside the TTL grace, so the row is still + # readable; the script must still refuse it. + frozen_clock.advance_ms(60_000 + LEASE_TTL_GRACE_MS / 2) + assert await store.get("co_1", "ct_1") is not None + assert await store.try_reserve("co_1", "ct_1", 10) is None + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.local_remaining_credits == 100 + + +async def test_try_reserve_rejects_nan_before_it_reaches_the_script( + store: RedisLeaseStore, frozen_clock: VirtualClock +) -> None: + # The string form of NaN parses back to a Lua nan and would poison the + # SHARED balance for every pod. + await _seed(store, frozen_clock) + assert await store.try_reserve("co_1", "ct_1", math.nan) is None + assert await store.try_reserve("co_1", "ct_1", -10) is None + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.local_remaining_credits == 100 + assert await store.try_reserve("co_1", "ct_1", 30) == 70 + + +async def test_fractional_credits_survive_the_round_trip( + store: RedisLeaseStore, frozen_clock: VirtualClock +) -> None: + await _seed(store, frozen_clock, granted=10) + assert await store.try_reserve("co_1", "ct_1", 2.5) == 7.5 + await store.refund("co_1", "ct_1", 1.25) + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.local_remaining_credits == 8.75 + + +async def test_refund_caps_at_the_granted_amount(store: RedisLeaseStore, frozen_clock: VirtualClock) -> None: + await _seed(store, frozen_clock) + await store.try_reserve("co_1", "ct_1", 30) + await store.refund("co_1", "ct_1", 9999) + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.local_remaining_credits == 100 + + +async def test_refund_pinned_to_a_stale_lease_is_dropped( + store: RedisLeaseStore, frozen_clock: VirtualClock +) -> None: + await store.replace( + lease_id="lse_b", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=100, + expires_at=frozen_clock() + 60, + ) + await store.try_reserve("co_1", "ct_1", 50) + await store.refund("co_1", "ct_1", 30, "lse_a") + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.local_remaining_credits == 50 + await store.refund("co_1", "ct_1", 30, "lse_b") + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.local_remaining_credits == 80 + + +async def test_extend_falls_back_to_the_configured_duration( + redis_client: Any, frozen_clock: VirtualClock +) -> None: + store = RedisLeaseStore(redis_client, default_lease_duration=0.25, clock=frozen_clock) + # A short initial expiry so the fallback is a forward move. + await store.replace( + lease_id="lse_1", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=100, + expires_at=frozen_clock() + 0.1, + ) + await store.extend("co_1", "ct_1", 150) + entry = await store.get("co_1", "ct_1") + assert entry is not None + assert entry.granted_amount == 150 + assert entry.expires_at == pytest.approx(frozen_clock() + 0.25) + + +async def test_extend_pinned_to_a_stale_lease_is_dropped( + store: RedisLeaseStore, frozen_clock: VirtualClock +) -> None: + await store.replace( + lease_id="lse_b", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=100, + expires_at=frozen_clock() + 60, + ) + await store.extend("co_1", "ct_1", 150, frozen_clock() + 120, "lse_a") + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.granted_amount == 100 and entry.local_remaining_credits == 100 + await store.extend("co_1", "ct_1", 150, frozen_clock() + 120, "lse_b") + entry = await store.get("co_1", "ct_1") + assert entry is not None and entry.granted_amount == 150 and entry.local_remaining_credits == 150 + + +async def test_concurrent_sibling_extends_converge_on_the_server_total( + redis_client: Any, frozen_clock: VirtualClock +) -> None: + # Two pods on one Redis. Per-process single-flight cannot serialize them, + # so both wire calls go out against the same stale read (granted=100); the + # server lands B's total (150) then A's (200). Each pod applies the TOTAL + # it was handed, so the slot converges on 200, never 250. + pod_a = RedisLeaseStore(redis_client, clock=frozen_clock) + pod_b = RedisLeaseStore(redis_client, clock=frozen_clock) + await pod_a.replace( + lease_id="lse_1", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=100, + expires_at=frozen_clock() + 60, + ) + await pod_b.extend("co_1", "ct_1", 150, frozen_clock() + 90, "lse_1") + await pod_a.extend("co_1", "ct_1", 200, frozen_clock() + 120, "lse_1") + entry = await pod_a.get("co_1", "ct_1") + assert entry is not None and entry.granted_amount == 200 and entry.local_remaining_credits == 200 + + # Out-of-order arrival: the larger total lands first and the superseded one + # is a no-op that never pulls the expiry back. + await pod_a.replace( + lease_id="lse_2", + company_id="co_1", + credit_type_id="ct_2", + granted_amount=100, + expires_at=frozen_clock() + 60, + ) + far_expiry = frozen_clock() + 120 + await pod_a.extend("co_1", "ct_2", 200, far_expiry, "lse_2") + await pod_b.extend("co_1", "ct_2", 150, frozen_clock() + 90, "lse_2") + entry = await pod_a.get("co_1", "ct_2") + assert entry is not None + assert entry.granted_amount == 200 + assert entry.local_remaining_credits == 200 + assert entry.expires_at == far_expiry + + +async def test_drop_removes_the_hash(store: RedisLeaseStore, frozen_clock: VirtualClock) -> None: + await _seed(store, frozen_clock) + await store.drop("co_1", "ct_1") + assert await store.get("co_1", "ct_1") is None + + +async def test_key_layout_and_hash_fields_match_the_node_sdk( + store: RedisLeaseStore, redis_client: Any, frozen_clock: VirtualClock +) -> None: + # Node and Python pods share one Redis, so the key, the camelCase field + # names, and the millisecond instants have to match exactly. + expires_at = frozen_clock() + 60 + await store.replace( + lease_id="lse_1", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=expires_at, + ) + key = store.hash_key("co_1", "ct_1") + assert key == "schematic:credit-lease:co_1:ct_1" + raw = await redis_client.hgetall(key) + assert raw == { + "leaseId": "lse_1", + "companyId": "co_1", + "creditTypeId": "ct_1", + "grantedAmount": "1000", + "localRemainingCredits": "1000", + "expiresAt": str(int(round(expires_at * 1000))), + } + # The row outlives its expiry by the grace window so the sweeper can still + # read it. + ttl_ms = await redis_client.pttl(key) + assert 60_000 < ttl_ms <= 60_000 + LEASE_TTL_GRACE_MS + + +async def test_a_flushed_script_cache_falls_back_to_eval( + store: RedisLeaseStore, redis_client: Any, frozen_clock: VirtualClock +) -> None: + # A Redis that restarts (or is flushed) loses the cached script and answers + # NOSCRIPT; the store re-sends the body rather than failing the reserve. + await _seed(store, frozen_clock) + assert await store.try_reserve("co_1", "ct_1", 10) == 90 + await redis_client.script_flush() + assert await store.try_reserve("co_1", "ct_1", 10) == 80 diff --git a/tests/leases/test_redis_reservation_store.py b/tests/leases/test_redis_reservation_store.py new file mode 100644 index 00000000..c6271a45 --- /dev/null +++ b/tests/leases/test_redis_reservation_store.py @@ -0,0 +1,226 @@ +"""Redis reservation table semantics, ported from the Node SDK's tests.""" + +from __future__ import annotations + +from typing import Any + +import pytest +from lease_support import VirtualClock, make_reservation + +from schematic.leases import RedisLeaseStore, RedisReservationStore +from schematic.leases.redis_reservation_store import RES_TTL_GRACE_MS, SWEEP_BATCH_SIZE + + +@pytest.fixture +def leases(redis_client: Any, frozen_clock: VirtualClock) -> RedisLeaseStore: + return RedisLeaseStore(redis_client, clock=frozen_clock) + + +@pytest.fixture +def reservations(redis_client: Any, leases: RedisLeaseStore, frozen_clock: VirtualClock) -> RedisReservationStore: + return RedisReservationStore(redis_client, leases, clock=frozen_clock) + + +@pytest.fixture(autouse=True) +async def seeded_lease(leases: RedisLeaseStore, frozen_clock: VirtualClock) -> None: + await leases.replace( + lease_id="lse_1", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=frozen_clock() + 600, + ) + + +async def _balance(leases: RedisLeaseStore) -> float: + entry = await leases.get("co_1", "ct_1") + assert entry is not None + return entry.local_remaining_credits + + +async def test_add_round_trips_and_indexes( + reservations: RedisReservationStore, frozen_clock: VirtualClock +) -> None: + await reservations.add(make_reservation(expires_at=frozen_clock() + 60)) + fetched = await reservations.get("res_1") + assert fetched is not None + assert fetched.credits_reserved == 100 + assert fetched.lease_id == "lse_1" + assert fetched.company == {"id": "co_1"} + assert await reservations.count() == 1 + + +async def test_consume_refunds_the_unspent_slice( + leases: RedisLeaseStore, reservations: RedisReservationStore, frozen_clock: VirtualClock +) -> None: + await leases.try_reserve("co_1", "ct_1", 100) + await reservations.add(make_reservation(expires_at=frozen_clock() + 60)) + assert await _balance(leases) == 900 + + assert await reservations.consume("res_1", 30) == 30 + assert await _balance(leases) == 970 + assert await reservations.get("res_1") is None + + +async def test_double_consume_returns_null( + leases: RedisLeaseStore, reservations: RedisReservationStore, frozen_clock: VirtualClock +) -> None: + await leases.try_reserve("co_1", "ct_1", 100) + await reservations.add(make_reservation(expires_at=frozen_clock() + 60)) + await reservations.consume("res_1", 50) + assert await reservations.consume("res_1", 10) is None + + +async def test_sweep_returns_expired_holds_to_the_lease( + leases: RedisLeaseStore, reservations: RedisReservationStore, frozen_clock: VirtualClock +) -> None: + await leases.try_reserve("co_1", "ct_1", 100) + await reservations.add(make_reservation(expires_at=frozen_clock() - 0.001)) + assert await reservations.sweep_expired() == 1 + assert await _balance(leases) == 1000 + assert await reservations.get("res_1") is None + + +async def test_a_stale_lease_hold_never_inflates_its_successor( + leases: RedisLeaseStore, reservations: RedisReservationStore, frozen_clock: VirtualClock +) -> None: + await leases.try_reserve("co_1", "ct_1", 100) + await reservations.add(make_reservation(expires_at=frozen_clock() - 0.001)) + + await leases.drop("co_1", "ct_1") + await leases.replace( + lease_id="lse_2", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=frozen_clock() + 600, + ) + await leases.try_reserve("co_1", "ct_1", 200) + + assert await reservations.sweep_expired() == 1 + assert await _balance(leases) == 800 + assert await reservations.get("res_1") is None + + +async def test_reserved_credits_sums_open_holds( + leases: RedisLeaseStore, reservations: RedisReservationStore, frozen_clock: VirtualClock +) -> None: + await leases.try_reserve("co_1", "ct_1", 350) + await reservations.add(make_reservation(id="res_1", credits_reserved=100, expires_at=frozen_clock() + 60)) + await reservations.add(make_reservation(id="res_2", credits_reserved=250, expires_at=frozen_clock() + 60)) + await reservations.add( + make_reservation(id="res_3", credit_type_id="ct_2", credits_reserved=999, expires_at=frozen_clock() + 60) + ) + + assert await reservations.reserved_credits("co_1", "ct_1") == 350 + assert await reservations.reserved_credits("co_1", "ct_2") == 999 + + await reservations.consume("res_1", 40) + assert await reservations.reserved_credits("co_1", "ct_1") == 250 + + +async def test_every_lua_call_touches_one_key( + redis_client: Any, leases: RedisLeaseStore, reservations: RedisReservationStore, frozen_clock: VirtualClock +) -> None: + # A multi-key script whose keys hash to different slots raises CROSSSLOT on + # Redis Cluster, so every EVAL these stores send must touch exactly one. + key_counts = [] + original_eval = redis_client.eval + original_evalsha = redis_client.evalsha + + async def record_eval(script: str, numkeys: int, *args: Any) -> Any: + key_counts.append(numkeys) + return await original_eval(script, numkeys, *args) + + async def record_evalsha(sha: str, numkeys: int, *args: Any) -> Any: + key_counts.append(numkeys) + return await original_evalsha(sha, numkeys, *args) + + redis_client.eval = record_eval + redis_client.evalsha = record_evalsha + + await leases.try_reserve("co_1", "ct_1", 200) + await leases.refund("co_1", "ct_1", 50) + await leases.extend("co_1", "ct_1", 1100, frozen_clock() + 120) + await reservations.add(make_reservation(id="res_a", credits_reserved=100, expires_at=frozen_clock() + 60)) + await reservations.add( + make_reservation(id="res_b", credits_reserved=80, expires_at=frozen_clock() - 0.001) + ) + await reservations.consume("res_a", 40) + await reservations.sweep_expired() + + assert key_counts + assert set(key_counts) == {1} + + +async def test_sweep_reconciles_indexes_when_the_hash_has_evicted( + leases: RedisLeaseStore, reservations: RedisReservationStore, frozen_clock: VirtualClock +) -> None: + # A sweeper that goes silent long enough (deploy, restart, starvation) for + # Redis to evict the reservation hash must still clean both indexes, or the + # orphaned entry inflates reserved_credits forever. + await leases.try_reserve("co_1", "ct_1", 100) + await reservations.add(make_reservation(expires_at=frozen_clock() + 1)) + assert await reservations.count() == 1 + assert await reservations.reserved_credits("co_1", "ct_1") == 100 + + frozen_clock.advance_ms(1000 + RES_TTL_GRACE_MS + 1) + + # Nothing is swept in the refund sense: without the hash, exactly-once + # cannot be arbitrated across racing sweepers, so the slice waits for the + # lease to expire server-side. + assert await reservations.sweep_expired() == 0 + assert await reservations.count() == 0 + assert await reservations.reserved_credits("co_1", "ct_1") == 0 + assert await _balance(leases) == 900 + + +async def test_sweep_drops_an_unparseable_index_member( + redis_client: Any, reservations: RedisReservationStore, frozen_clock: VirtualClock +) -> None: + # Only `add` writes members, so this cannot happen; a member that does not + # decode must still be removed rather than re-read by every later sweep. + await redis_client.zadd( + "schematic:credit-reservations:byExpiry", {"garbage": int(frozen_clock() * 1000) - 1} + ) + assert await reservations.sweep_expired() == 0 + assert await reservations.count() == 0 + + +async def test_sweep_pages_through_a_backlog( + leases: RedisLeaseStore, reservations: RedisReservationStore, frozen_clock: VirtualClock +) -> None: + count = SWEEP_BATCH_SIZE + 44 + await leases.try_reserve("co_1", "ct_1", count) + for index in range(count): + await reservations.add( + make_reservation(id=f"res_{index}", credits_reserved=1, expires_at=frozen_clock() - 0.001) + ) + assert await reservations.sweep_expired() == count + assert await _balance(leases) == 1000 + assert await reservations.count() == 0 + assert await reservations.reserved_credits("co_1", "ct_1") == 0 + + +async def test_key_layout_and_hash_fields_match_the_node_sdk( + redis_client: Any, reservations: RedisReservationStore, frozen_clock: VirtualClock +) -> None: + expires_at = frozen_clock() + 60 + await reservations.add(make_reservation(expires_at=expires_at)) + raw = await redis_client.hgetall("schematic:credit-reservation:res_1") + assert raw == { + "id": "res_1", + "leaseId": "lse_1", + "companyId": "co_1", + "creditTypeId": "ct_1", + "eventSubtype": "inference_tokens", + "quantityReserved": "10", + "creditsReserved": "100", + "consumptionRate": "10", + "expiresAt": str(int(round(expires_at * 1000))), + "evalCtx": '{"company":{"id":"co_1"}}', + } + assert await redis_client.zrange("schematic:credit-reservations:byExpiry", 0, -1) == ["co_1|ct_1|res_1"] + assert await redis_client.hgetall("schematic:credit-reservations:byCredit:co_1:ct_1") == {"res_1": "100"} + ttl_ms = await redis_client.pttl("schematic:credit-reservation:res_1") + assert 60_000 < ttl_ms <= 60_000 + RES_TTL_GRACE_MS diff --git a/tests/leases/test_reservation_store.py b/tests/leases/test_reservation_store.py new file mode 100644 index 00000000..10c5d2f8 --- /dev/null +++ b/tests/leases/test_reservation_store.py @@ -0,0 +1,133 @@ +"""In-memory reservation table semantics, ported from the Node SDK's tests.""" + +from __future__ import annotations + +import pytest +from lease_support import VirtualClock, make_reservation + +from schematic.leases import InMemoryLeaseStore, InMemoryReservationStore + + +@pytest.fixture +def leases(clock: VirtualClock) -> InMemoryLeaseStore: + return InMemoryLeaseStore(clock=clock) + + +@pytest.fixture +def reservations(leases: InMemoryLeaseStore, clock: VirtualClock) -> InMemoryReservationStore: + return InMemoryReservationStore(leases, clock=clock) + + +@pytest.fixture(autouse=True) +async def seeded_lease(leases: InMemoryLeaseStore, clock: VirtualClock) -> None: + await leases.replace( + lease_id="lse_1", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=clock() + 60, + ) + + +async def _balance(leases: InMemoryLeaseStore) -> float: + entry = await leases.get("co_1", "ct_1") + assert entry is not None + return entry.local_remaining_credits + + +async def test_consume_refunds_the_unspent_slice( + leases: InMemoryLeaseStore, reservations: InMemoryReservationStore, clock: VirtualClock +) -> None: + await leases.try_reserve("co_1", "ct_1", 100) + await reservations.add(make_reservation(expires_at=clock() + 60)) + assert await _balance(leases) == 900 + + assert await reservations.consume("res_1", 30) == 30 + assert await _balance(leases) == 970 + assert await reservations.get("res_1") is None + + +async def test_consume_clamps_to_the_hold( + leases: InMemoryLeaseStore, reservations: InMemoryReservationStore, clock: VirtualClock +) -> None: + await leases.try_reserve("co_1", "ct_1", 100) + await reservations.add(make_reservation(expires_at=clock() + 60)) + assert await reservations.consume("res_1", 999) == 100 + assert await _balance(leases) == 900 + + +async def test_consume_of_a_missing_reservation_is_null(reservations: InMemoryReservationStore) -> None: + assert await reservations.consume("nope", 10) is None + + +async def test_consume_is_exactly_once( + leases: InMemoryLeaseStore, reservations: InMemoryReservationStore, clock: VirtualClock +) -> None: + await leases.try_reserve("co_1", "ct_1", 100) + await reservations.add(make_reservation(expires_at=clock() + 60)) + assert await reservations.consume("res_1", 30) == 30 + assert await reservations.consume("res_1", 30) is None + assert await _balance(leases) == 970 + + +async def test_a_stale_lease_hold_never_inflates_its_successor( + leases: InMemoryLeaseStore, reservations: InMemoryReservationStore, clock: VirtualClock +) -> None: + await leases.try_reserve("co_1", "ct_1", 100) + await reservations.add(make_reservation(expires_at=clock() - 0.001)) + + # lse_1 expires and lse_2 takes the slot, partially debited. + await leases.drop("co_1", "ct_1") + await leases.replace( + lease_id="lse_2", + company_id="co_1", + credit_type_id="ct_1", + granted_amount=1000, + expires_at=clock() + 60, + ) + await leases.try_reserve("co_1", "ct_1", 200) + + # Sweeping lse_1's hold must not credit lse_2: that slice went back to the + # company balance when lse_1 expired server-side. + assert await reservations.sweep_expired() == 1 + assert await _balance(leases) == 800 + + # An explicit consume of a stale-lease hold is dropped the same way. + await reservations.add(make_reservation(id="res_2", expires_at=clock() + 60)) + await reservations.consume("res_2", 0) + assert await _balance(leases) == 800 + + +async def test_sweep_refunds_expired_holds_only( + leases: InMemoryLeaseStore, reservations: InMemoryReservationStore, clock: VirtualClock +) -> None: + await leases.try_reserve("co_1", "ct_1", 100) + await reservations.add(make_reservation(expires_at=clock() + 10)) + assert await reservations.sweep_expired() == 0 + assert await reservations.count() == 1 + + clock.advance_ms(10_001) + assert await reservations.sweep_expired() == 1 + assert await reservations.get("res_1") is None + assert await _balance(leases) == 1000 + + +async def test_reserved_credits_sums_the_slot_only(reservations: InMemoryReservationStore) -> None: + await reservations.add(make_reservation(id="res_1", credits_reserved=100)) + await reservations.add(make_reservation(id="res_2", credits_reserved=250)) + await reservations.add(make_reservation(id="res_3", credit_type_id="ct_2", credits_reserved=999)) + await reservations.add(make_reservation(id="res_4", company_id="co_2", credits_reserved=999)) + + assert await reservations.reserved_credits("co_1", "ct_1") == 350 + assert await reservations.reserved_credits("co_1", "ct_2") == 999 + assert await reservations.reserved_credits("co_unknown", "ct_1") == 0 + + +async def test_reserved_credits_drops_a_consumed_hold( + leases: InMemoryLeaseStore, reservations: InMemoryReservationStore, clock: VirtualClock +) -> None: + await leases.try_reserve("co_1", "ct_1", 100) + await reservations.add(make_reservation(expires_at=clock() + 60)) + assert await reservations.reserved_credits("co_1", "ct_1") == 100 + await reservations.consume("res_1", 30) + assert await reservations.reserved_credits("co_1", "ct_1") == 0 diff --git a/tests/leases/test_wasm_credit_gate.py b/tests/leases/test_wasm_credit_gate.py new file mode 100644 index 00000000..90ea574c --- /dev/null +++ b/tests/leases/test_wasm_credit_gate.py @@ -0,0 +1,300 @@ +"""The credit gate against the real WASM rules engine. + +Every other lease test scripts the engine, so the contract that matters most +goes unexercised: resolving the matched credit entitlement from the probe, +substituting the lease balance into the company's credit balances, and letting +the engine's credit_cost gate decide. A drift in the option envelope or in the +entity shape the SDK feeds the engine fails here rather than mis-gating in +production. +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Optional + +import pytest +from lease_support import ScriptedWireClient, VirtualClock + +from schematic.client import CheckFlagOptions, CheckOptions, CheckResult, EventUsage +from schematic.datastream.rules_engine import RulesEngineClient +from schematic.leases import ( + CreditCheckDeps, + InMemoryLeaseStore, + InMemoryReservationStore, + LeaseConfig, + LeaseManager, + check_with_lease, +) +from schematic.types import ( + RulesengineCompany, + RulesengineCondition, + RulesengineFeatureEntitlement, + RulesengineFlag, + RulesengineRule, +) + +wasmtime = pytest.importorskip("wasmtime", reason="wasmtime not installed") + +FLAG_KEY = "infer" +CREDIT_ID = "credit-1" +EVENT_SUBTYPE = "inference_tokens" +COMPANY_ID = "co" +LEASE_SIZE = 10_000.0 + + +def _credit_condition() -> RulesengineCondition: + # A consumption rate of 1 keeps credits equal to quantity, so the + # arithmetic in the assertions is the engine's own. + return RulesengineCondition( + id="cond-credit", + account_id="acct", + environment_id="env", + condition_type="credit", + operator="lt", + resource_ids=[], + trait_value="", + metric_value=0, + credit_id=CREDIT_ID, + consumption_rate=1, + event_subtype=EVENT_SUBTYPE, + ) + + +def _company_condition(resource_ids: List[str]) -> RulesengineCondition: + """A membership condition to flip, so the engine can deny for a reason that + has nothing to do with the balance.""" + return RulesengineCondition( + id="cond-company", + account_id="acct", + environment_id="env", + condition_type="company", + operator="eq", + resource_ids=resource_ids, + trait_value="", + metric_value=0, + ) + + +def _credit_flag(extra_conditions: Optional[List[RulesengineCondition]] = None) -> RulesengineFlag: + return RulesengineFlag( + id="flag-infer", + account_id="acct", + environment_id="env", + key=FLAG_KEY, + default_value=False, + rules=[ + RulesengineRule( + id="rule-credit", + account_id="acct", + environment_id="env", + name="Credit", + rule_type="plan_entitlement", + priority=100, + value=True, + conditions=[_credit_condition(), *(extra_conditions or [])], + condition_groups=[], + ) + ], + ) + + +def _company( + credit_balance: float, entitlements: Optional[List[RulesengineFeatureEntitlement]] = None +) -> RulesengineCompany: + # The company carries its resolved entitlement for the feature, the shape + # the DataStream cache holds after a plan assignment. Entitlement-first + # resolution reads the credit, the rate, and the subtype off it. + default = RulesengineFeatureEntitlement( + feature_id="feat-infer", + feature_key=FLAG_KEY, + value_type="credit", + credit_id=CREDIT_ID, + consumption_rate=1, + event_subtype=EVENT_SUBTYPE, + credit_total=credit_balance, + credit_used=0, + credit_remaining=credit_balance, + ) + return RulesengineCompany( + id=COMPANY_ID, + account_id="acct", + environment_id="env", + keys={"id": COMPANY_ID}, + traits=[], + metrics=[], + rules=[], + entitlements=entitlements if entitlements is not None else [default], + billing_product_ids=[], + credit_balances={CREDIT_ID: credit_balance}, + plan_ids=[], + plan_version_ids=[], + ) + + +class RealEngineDataStream: + """Serves a fixed flag and company, and the real rules engine behind them.""" + + def __init__(self, engine: RulesEngineClient, flag: RulesengineFlag, company: RulesengineCompany) -> None: + self._engine = engine + self._flag = flag + self._company = company + + async def get_flag(self, flag_key: str) -> RulesengineFlag: + return self._flag + + async def get_company(self, keys: Dict[str, str]) -> RulesengineCompany: + return self._company + + async def get_user(self, keys: Dict[str, str]) -> None: + return None + + def evaluate_flag(self, flag: Any, company: Any, user: Any, options: Any = None) -> Any: + return self._engine.check_flag(flag, company, user, options) + + +def _deps( + engine: RulesEngineClient, flag: RulesengineFlag, company: RulesengineCompany, clock: VirtualClock +) -> CreditCheckDeps: + leases = InMemoryLeaseStore(clock=clock) + reservations = InMemoryReservationStore(leases, clock=clock) + wire = ScriptedWireClient() + wire.acquire_responses.append( + {"lease": {"lease_id": "lse-1", "granted_amount": LEASE_SIZE, "expires_at": clock() + 60}} + ) + manager = LeaseManager( + wire, + leases, + reservation_store=reservations, + config=LeaseConfig(lease_duration=60.0, reservation_ttl=60.0, lease_size=LEASE_SIZE), + clock=clock, + ) + async def enqueue_flag_check(body: Any) -> None: + """These tests pin the engine's verdict, not the analytics event.""" + + return CreditCheckDeps( + datastream=RealEngineDataStream(engine, flag, company), + lease_store=leases, + reservations=reservations, + manager=manager, + logger=logging.getLogger("wasm-credit-gate-test"), + enqueue_flag_check=enqueue_flag_check, + clock=clock, + ) + + +async def _fail_fallback() -> CheckResult: + raise AssertionError("the lease path should not have fallen back") + + +@pytest.fixture +async def engine() -> RulesEngineClient: + client = RulesEngineClient() + await client.initialize() + return client + + +class TestCreditGateAgainstTheRealEngine: + async def test_a_within_balance_usage_passes_and_holds( + self, engine: RulesEngineClient, clock: VirtualClock + ) -> None: + deps = _deps(engine, _credit_flag(), _company(100), clock) + result = await check_with_lease( + deps, + FLAG_KEY, + {"id": COMPANY_ID}, + None, + CheckOptions(usage=50, event_subtype=EVENT_SUBTYPE), + _fail_fallback, + ) + await deps.manager._drain_background() + + assert result.allowed is True + assert result.value is True + assert result.reservation is not None + assert result.reservation.credit_type_id == CREDIT_ID + assert result.reservation.credits_reserved == 50 + entry = await deps.lease_store.get(COMPANY_ID, CREDIT_ID) + assert entry is not None and entry.local_remaining_credits == LEASE_SIZE - 50 + assert await deps.reservations.count() == 1 + + async def test_a_non_credit_denial_refunds_the_hold( + self, engine: RulesEngineClient, clock: VirtualClock + ) -> None: + # Credits are plentiful, but the membership condition excludes this + # company, so the hold taken before the gate has to come back. + flag = _credit_flag([_company_condition(["some-other-company"])]) + deps = _deps(engine, flag, _company(10_000), clock) + result = await check_with_lease( + deps, + FLAG_KEY, + {"id": COMPANY_ID}, + None, + CheckOptions(usage=50, event_subtype=EVENT_SUBTYPE), + _fail_fallback, + ) + + assert result.allowed is False + assert result.value is False + assert result.reservation is None + entry = await deps.lease_store.get(COMPANY_ID, CREDIT_ID) + assert entry is not None and entry.local_remaining_credits == LEASE_SIZE + assert await deps.reservations.count() == 0 + + async def test_an_override_granted_company_never_touches_the_lease( + self, engine: RulesEngineClient, clock: VirtualClock + ) -> None: + # A company override grants the feature outright, so the company's + # effective entitlement is boolean rather than credit-metered: no + # reserve-then-cancel, and no credits billed for usage the override + # grants for free. + flag = _credit_flag() + override = RulesengineRule( + id="rule-override", + account_id="acct", + environment_id="env", + name="Override", + rule_type="company_override", + priority=1, + value=True, + conditions=[_company_condition([COMPANY_ID])], + condition_groups=[], + ) + flag = flag.model_copy(update={"rules": [override, *flag.rules]}) + company = _company( + 100, + [RulesengineFeatureEntitlement(feature_id="feat-infer", feature_key=FLAG_KEY, value_type="boolean")], + ) + deps = _deps(engine, flag, company, clock) + fell_back = False + + async def fallback() -> CheckResult: + nonlocal fell_back + fell_back = True + return CheckResult(allowed=True, value=True, reason="override", flag_key=FLAG_KEY) + + result = await check_with_lease( + deps, FLAG_KEY, {"id": COMPANY_ID}, None, CheckOptions(usage=50, event_subtype=EVENT_SUBTYPE), fallback + ) + + assert fell_back is True + assert result.allowed is True + assert result.reservation is None + assert await deps.lease_store.get(COMPANY_ID, CREDIT_ID) is None + assert await deps.reservations.count() == 0 + + async def test_the_preflight_envelope_gates_at_the_balance_boundary(self, engine: RulesEngineClient) -> None: + # The contract the lease path leans on: the SDK's event_usage option + # reaches the engine as the envelope it gates on. + flag = _credit_flag() + company = _company(100) + + under = engine.check_flag( + flag, company, None, CheckFlagOptions(event_usage=EventUsage(event_subtype=EVENT_SUBTYPE, quantity=50)) + ) + over = engine.check_flag( + flag, company, None, CheckFlagOptions(event_usage=EventUsage(event_subtype=EVENT_SUBTYPE, quantity=150)) + ) + + assert under.value is True + assert over.value is False diff --git a/tests/leases/test_wire_client.py b/tests/leases/test_wire_client.py new file mode 100644 index 00000000..6f499794 --- /dev/null +++ b/tests/leases/test_wire_client.py @@ -0,0 +1,131 @@ +"""The adapter between the manager and the generated credits client.""" + +from __future__ import annotations + +import datetime as dt +from typing import Any, Dict, List, Optional + +from schematic.credits.types.acquire_credit_lease_response import AcquireCreditLeaseResponse +from schematic.credits.types.extend_credit_lease_response import ExtendCreditLeaseResponse +from schematic.leases import CreditsWireClient, LeaseWireClient +from schematic.types.credit_lease_response_data import CreditLeaseResponseData + +EXPIRES_AT = dt.datetime(2026, 1, 1, 0, 5, tzinfo=dt.timezone.utc) + + +def _lease_data(lease_id: str = "lse_1", granted_amount: float = 1000) -> CreditLeaseResponseData: + now = dt.datetime(2026, 1, 1, tzinfo=dt.timezone.utc) + return CreditLeaseResponseData( + id=lease_id, + company_id="co_1", + credit_type_id="ct_1", + granted_amount=granted_amount, + tracked_amount=0, + expires_at=EXPIRES_AT, + created_at=now, + updated_at=now, + ) + + +class StubCreditsClient: + def __init__(self) -> None: + self.acquire_calls: List[Dict[str, Any]] = [] + self.extend_calls: List[Dict[str, Any]] = [] + self.release_calls: List[str] = [] + + async def acquire_credit_lease(self, **kwargs: Any) -> AcquireCreditLeaseResponse: + self.acquire_calls.append(kwargs) + return AcquireCreditLeaseResponse(data=_lease_data(), params={}) + + async def extend_credit_lease(self, lease_id: str, **kwargs: Any) -> ExtendCreditLeaseResponse: + self.extend_calls.append({"lease_id": lease_id, **kwargs}) + return ExtendCreditLeaseResponse(data=_lease_data(granted_amount=2000), params={}) + + async def release_credit_lease(self, lease_id: str, **kwargs: Any) -> None: + self.release_calls.append(lease_id) + + +async def test_acquire_maps_the_request_and_the_response() -> None: + stub = StubCreditsClient() + wire: LeaseWireClient = CreditsWireClient(stub) + + grant = await wire.acquire("co_1", "ct_1", 1000, EXPIRES_AT.timestamp()) + + assert stub.acquire_calls[0]["company_id"] == "co_1" + assert stub.acquire_calls[0]["credit_type_id"] == "ct_1" + assert stub.acquire_calls[0]["requested_amount"] == 1000 + assert stub.acquire_calls[0]["expires_at"] == EXPIRES_AT + assert grant.lease_id == "lse_1" + assert grant.granted_amount == 1000 + assert grant.expires_at == EXPIRES_AT.timestamp() + + +async def test_extend_sends_the_additional_amount_and_reads_back_the_total() -> None: + stub = StubCreditsClient() + wire: LeaseWireClient = CreditsWireClient(stub) + + grant = await wire.extend("lse_1", 500, EXPIRES_AT.timestamp()) + + assert stub.extend_calls[0]["lease_id"] == "lse_1" + assert stub.extend_calls[0]["additional_amount"] == 500 + # The response carries the server-authoritative TOTAL, not the increment. + assert grant.granted_amount == 2000 + + +async def test_release_passes_the_lease_id() -> None: + stub = StubCreditsClient() + wire: LeaseWireClient = CreditsWireClient(stub) + await wire.release("lse_1") + assert stub.release_calls == ["lse_1"] + + +async def test_a_naive_expiry_is_read_as_utc() -> None: + # A naive timestamp from the API is UTC; reading it as local time would + # shift every expiry by the pod's offset. + class NaiveClient(StubCreditsClient): + async def acquire_credit_lease(self, **kwargs: Any) -> AcquireCreditLeaseResponse: + data = _lease_data().model_copy(update={"expires_at": EXPIRES_AT.replace(tzinfo=None)}) + return AcquireCreditLeaseResponse(data=data, params={}) + + wire: LeaseWireClient = CreditsWireClient(NaiveClient()) + grant = await wire.acquire("co_1", "ct_1", 1000, EXPIRES_AT.timestamp()) + assert grant.expires_at == EXPIRES_AT.timestamp() + + +async def test_request_options_are_threaded_through() -> None: + stub = StubCreditsClient() + options: Optional[Any] = {"timeout_in_seconds": 2} + wire: LeaseWireClient = CreditsWireClient(stub, request_options=options) + await wire.acquire("co_1", "ct_1", 1000, EXPIRES_AT.timestamp()) + assert stub.acquire_calls[0]["request_options"] == options + + +async def test_a_per_call_timeout_becomes_request_options() -> None: + stub = StubCreditsClient() + wire: LeaseWireClient = CreditsWireClient(stub) + + await wire.acquire("co_1", "ct_1", 1000, EXPIRES_AT.timestamp(), 1.5) + await wire.extend("lse_1", 500, EXPIRES_AT.timestamp(), 1.5) + + assert stub.acquire_calls[0]["request_options"] == {"timeout": 1.5} + assert stub.extend_calls[0]["request_options"] == {"timeout": 1.5} + + +async def test_no_timeout_sends_no_request_options() -> None: + stub = StubCreditsClient() + wire: LeaseWireClient = CreditsWireClient(stub) + + await wire.acquire("co_1", "ct_1", 1000, EXPIRES_AT.timestamp()) + await wire.extend("lse_1", 500, EXPIRES_AT.timestamp()) + + assert stub.acquire_calls[0]["request_options"] is None + assert stub.extend_calls[0]["request_options"] is None + + +async def test_a_per_call_timeout_wins_over_the_client_wide_options() -> None: + stub = StubCreditsClient() + wire: LeaseWireClient = CreditsWireClient(stub, request_options={"timeout": 30}) + + await wire.acquire("co_1", "ct_1", 1000, EXPIRES_AT.timestamp(), 1.5) + + assert stub.acquire_calls[0]["request_options"] == {"timeout": 1.5}