From 47ad2c6471d8072ec3f5d1545732dd0af026467f Mon Sep 17 00:00:00 2001 From: Ben Papillon Date: Tue, 15 Sep 2026 10:24:26 -0700 Subject: [PATCH 1/8] add server-mode credit reservations and preflight options --- README.md | 90 ++ src/schematic/client.py | 710 +++++++++++++++- src/schematic/datastream/datastream_client.py | 20 +- src/schematic/datastream/rules_engine.py | 34 +- tests/custom/test_client.py | 773 +++++++++++++++++- tests/datastream/test_datastream_client.py | 60 ++ tests/datastream/test_rules_engine.py | 72 ++ 7 files changed, 1738 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 87b72c2f..23565c74 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,6 +566,68 @@ 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. The server evaluates the flag against +the company's real balance and takes the hold in one call, then refunds the +unspent slice when the settling event arrives. + +Opt in with `credit_leases`: + +```python +from schematic.client import CreditLeaseConfig, Schematic, SchematicConfig + +config = SchematicConfig( + credit_leases=CreditLeaseConfig( + default_reservation_ttl=60.0, # seconds the server holds unsettled credits, max 1 hour + ), +) +client = Schematic("YOUR_API_KEY", config) +``` + +Then reserve the operation's upper bound, do the work, and report what it +actually used: + +```python +from schematic.client import CheckOptions + +result = client.check( + "inference", + company={"id": "your-company-id"}, + options=CheckOptions( + usage=max_tokens, # upper bound for this operation + event_subtype="inference_tokens", # the metered event + ), +) +if not result.allowed: + raise RuntimeError("credit balance exceeded") + +inference = run_inference() + +client.track_with_reservation(result.reservation, inference.tokens_used) +``` + +`AsyncSchematic` mirrors both methods: `await client.check(...)` and +`await client.track_with_reservation(...)`. + +A check that cannot gate, because the API is unreachable or errored, fails +closed by default: `allowed` is False and no hold is taken. Pass +`on_acquire_failure="fail-open"` for callers where letting traffic through +beats denying it, and the check returns your default value +(`CheckOptions.default_value`, else the client's flag default) instead. A 402 +is different: the server knows the credits are not there, so the check denies +whatever `on_acquire_failure` says. + +`mode` defaults to `auto`, which means server mode: every check with `usage` +is one API call. + +If nothing settles a reservation, its hold is refunded at +`default_reservation_ttl`. The settling event carries an idempotency key +derived from the reservation ID, so a retried or duplicated settle is billed +once. + ## 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. diff --git a/src/schematic/client.py b/src/schematic/client.py index adf07b9b..4487133f 100644 --- a/src/schematic/client.py +++ b/src/schematic/client.py @@ -2,12 +2,14 @@ import datetime as dt import logging 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 .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 @@ -22,6 +24,8 @@ EventBodyIdentifyCompany, EventBodyTrack, FeatureEntitlement, + PreflightEventUsageRequestBody, + PreflightRequestBody, RulesengineCheckFlagResult, ) @@ -32,6 +36,35 @@ 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 + +# 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": local leases carved up in-process. Not implemented in this SDK +# yet; see CreditLeaseConfig.mode. +# - "auto" (default): "server", until client mode exists here. +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 +72,229 @@ 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. + """ + + # Where the hold lives. "server" and "auto" both take the hold over the + # check-and-reserve API. "client" (local leases carved out in-process) is + # not implemented in this SDK yet: it leaves checks ungated and warns at + # construction. + mode: CreditLeaseMode = "auto" + # How long the server holds credits for an unsettled reservation, in + # seconds. Size it above the longest expected gap between check() and + # track_with_reservation(). The server caps a hold at one hour. + default_reservation_ttl: float = DEFAULT_RESERVATION_TTL + + +@dataclass +class CheckOptions: + """Options accepted by ``check``.""" + + # Units of the feature this operation will consume. The check holds + # usage * consumption_rate credits. A check takes at most one hold. + usage: Optional[int] = 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 _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. + """ + if options is None: + return None + flag_options = CheckFlagOptions(default_value=options.default_value, timeout=options.timeout) + if options.usage is not None: + if options.event_subtype is not None: + flag_options.event_usage = EventUsage(event_subtype=options.event_subtype, quantity=options.usage) + else: + flag_options.usage = options.usage + return flag_options + + +def _is_valid_quantity(value: Any) -> bool: + """Whether a caller-supplied usage can size a credit hold. + + A bool is an int in Python and a float can be NaN, and the server would + size a hold from either without any comparison rejecting it. + """ + return isinstance(value, int) and not isinstance(value, bool) and value >= 0 + + +def _resolve_lease_mode( + credit_leases: Optional[CreditLeaseConfig], offline: bool, +) -> Optional[Literal["client", "server"]]: + """Which reservation mode a ``check`` with usage resolves to right now. + + None means no credit gating at all: credit leases are not configured, the + client is offline, or the caller asked for a mode this SDK cannot serve. + """ + if credit_leases is None or offline: + return None + if credit_leases.mode == "client": + # Client-side leases do not exist in this SDK yet, and gating on a + # lease that was never taken would be a lie, so the check stays plain. + # A later PR returns "client" here without moving the public surface. + return None + # Other SDKs let "auto" pick client mode when DataStream is ready. With no + # client mode to pick, both "auto" and "server" mean server. + return "server" + + +def _warn_credit_lease_config( + logger: logging.Logger, credit_leases: CreditLeaseConfig, offline: bool, +) -> None: + """Say once, at construction, when the configured credit leases will not + gate anything.""" + if offline: + logger.warning( + "credit_leases is configured but the client is offline; check() returns flag defaults " + "and holds no credits." + ) + if credit_leases.mode == "client": + logger.warning( + "credit_leases.mode is 'client', which this SDK does not support yet; check() falls back to a " + "plain, ungated flag check. Use 'server' (or the 'auto' default) to gate on credits." + ) + + +def _reservation_request_kwargs(options: CheckOptions) -> Dict[str, Any]: + """Preflight body and per-check request options for a check-and-reserve + call, each omitted when the caller set nothing.""" + kwargs: Dict[str, Any] = {} + preflight = _build_preflight(_check_options_to_flag_options(options)) + if preflight is not None: + kwargs["preflight"] = preflight + if options.timeout is not None: + request_options: RequestOptions = {"timeout": options.timeout} + kwargs["request_options"] = request_options + return kwargs + + +def _payment_required_message(error: PaymentRequiredError) -> str: + """The server's own explanation for a 402, when the body carries one.""" + message = getattr(error.body, "error", None) + if isinstance(message, str) and message: + return message + return str(error) + + +def _build_reservation_track_event( + reservation: Reservation, + actual_quantity: int, + options: Optional[TrackWithReservationOptions] = None, +) -> EventBodyTrack: + """Build the track event that settles a reservation.""" + return EventBodyTrack( + company=reservation.company, + event=reservation.event_subtype, + # In server mode the hold lives on the server and settles by id. Never + # send lease_id as well: the server prefers it when both are set, 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, + ) @dataclass @@ -128,6 +384,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 +418,14 @@ 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 = ( + config.credit_leases.default_reservation_ttl + if config.credit_leases is not None + else DEFAULT_RESERVATION_TTL + ) + if config.credit_leases is not None: + _warn_credit_lease_config(self.logger, config.credit_leases, self.offline) atexit.register(self.shutdown) @@ -309,23 +574,193 @@ 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"]]: + return _resolve_lease_mode(self._credit_leases, self.offline) + + 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. + """ + 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 non-negative integer" + ) + return self._server_failure_result(flag_key, options, "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 PaymentRequiredError as e: + # 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", + flag_key=flag_key, + error=_payment_required_message(e), + ) + except Exception as e: + self.logger.error(f"Server reservation: check-and-reserve for flag {flag_key} failed: {e}") + return self._server_failure_result(flag_key, options, "server_reservation_failed") + + result = 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, + ) + + # 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 self._server_failure_result(flag_key, options, "missing_event_subtype") + + result.reservation = 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, + ) + return result + + def _server_failure_result(self, flag_key: str, options: CheckOptions, reason: str) -> CheckResult: + if options.on_acquire_failure == "fail-closed": + return CheckResult(allowed=False, value=False, reason=reason, flag_key=flag_key, error=reason) + value = self._resolve_default(flag_key, _check_options_to_flag_options(options)) + return CheckResult( + allowed=value, value=value, reason=f"{reason}_fail_open", flag_key=flag_key, error=reason, + ) + def identify( self, keys: Dict[str, str], @@ -366,6 +801,37 @@ def track( options=options, ) + def track_with_reservation( + self, + reservation: Reservation, + actual_quantity: int, + 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 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 non-negative integer. Skipping the settle, the hold is " + "refunded at its TTL" + ) + return + self._enqueue_event( + "track", + _build_reservation_track_event(reservation, actual_quantity, options), + options=TrackOptions(idempotency_key=f"{RESERVATION_TRACK_IDEMPOTENCY_PREFIX}{reservation.id}"), + ) + def _enqueue_event( self, event_type: str, @@ -418,6 +884,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 +949,14 @@ 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 = ( + config.credit_leases.default_reservation_ttl + if config.credit_leases is not None + else DEFAULT_RESERVATION_TTL + ) + if config.credit_leases is not None: + _warn_credit_lease_config(self.logger, config.credit_leases, self.offline) # DataStream client self._datastream_client: Optional[DataStreamClient] = None @@ -564,6 +1039,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 +1226,192 @@ 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"]]: + return _resolve_lease_mode(self._credit_leases, self.offline) + + 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_fallback(flag_key, company, user, options) + + 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. + """ + 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 non-negative integer" + ) + return self._server_failure_result(flag_key, options, "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 PaymentRequiredError as e: + # 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", + flag_key=flag_key, + error=_payment_required_message(e), + ) + except Exception as e: + self.logger.error(f"Server reservation: check-and-reserve for flag {flag_key} failed: {e}") + return self._server_failure_result(flag_key, options, "server_reservation_failed") + + result = 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, + ) + + # 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 self._server_failure_result(flag_key, options, "missing_event_subtype") + + result.reservation = 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, + ) + return result + + def _server_failure_result(self, flag_key: str, options: CheckOptions, reason: str) -> CheckResult: + if options.on_acquire_failure == "fail-closed": + return CheckResult(allowed=False, value=False, reason=reason, flag_key=flag_key, error=reason) + value = self._resolve_default(flag_key, _check_options_to_flag_options(options)) + return CheckResult( + allowed=value, value=value, reason=f"{reason}_fail_open", flag_key=flag_key, error=reason, + ) + async def identify( self, keys: Dict[str, str], @@ -808,6 +1453,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 +1469,40 @@ async def track( except Exception as e: self.logger.error(f"Failed to update company metrics: {e}") + async def track_with_reservation( + self, + reservation: Reservation, + actual_quantity: int, + 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 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 non-negative integer. Skipping the settle, the hold is " + "refunded at its TTL" + ) + return + await self._enqueue_event( + "track", + _build_reservation_track_event(reservation, actual_quantity, options), + 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, actual_quantity) + async def _enqueue_event( self, event_type: str, diff --git a/src/schematic/datastream/datastream_client.py b/src/schematic/datastream/datastream_client.py index 90faa8e3..019161d3 100644 --- a/src/schematic/datastream/datastream_client.py +++ b/src/schematic/datastream/datastream_client.py @@ -19,6 +19,10 @@ from .types import DataStreamBaseReq, DataStreamReq, DataStreamResp, EntityType, KeyConflictError, MessageType, RulesEngineError from .websocket_client import MAX_MESSAGE_SIZE, ClientOptions as WSClientOptions, DatastreamWSClient +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 +388,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 +424,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 +446,7 @@ 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) 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 +1016,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 +1029,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/tests/custom/test_client.py b/tests/custom/test_client.py index 668ffb4d..a1bb4938 100644 --- a/tests/custom/test_client.py +++ b/tests/custom/test_client.py @@ -1,3 +1,4 @@ +import datetime as dt import time import unittest from unittest.mock import AsyncMock, MagicMock, patch @@ -12,12 +13,27 @@ AsyncSchematic, AsyncSchematicConfig, CheckFlagOptions, + CheckOptions, + CreditLeaseConfig, + EventUsage, IdentifyOptions, + Reservation, Schematic, SchematicConfig, TrackOptions, + TrackWithReservationOptions, +) +from schematic.errors import PaymentRequiredError +from schematic.types import ( + ApiError, + CheckAndReserveFlagResponseData, + CheckFlagResponseData, + FeatureEntitlement, + FlagCheckReservationResponseData, + PreflightEventUsageRequestBody, + PreflightRequestBody, + RulesengineCheckFlagResult, ) -from schematic.types import CheckFlagResponseData, FeatureEntitlement class TestSchematic(unittest.TestCase): @@ -1530,5 +1546,760 @@ 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 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.assertNotIn("request_options", kwargs) + + # 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"], {"timeout": 2.5}) + + 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") + 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, 1.5, 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_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, 2.5, True): + self.schematic.track_with_reservation(reservation, quantity) # type: ignore[arg-type] + mock_push.assert_not_called() + + 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"] == {"timeout": 2.5} + + 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" + assert result.error == "credit balance exhausted" + + 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, 1.5, 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_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, 2.5, True): + await self.client.track_with_reservation(reservation, quantity) # type: ignore[arg-type] + mock_push.assert_not_called() + + 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() + + 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..7a9e3502 100644 --- a/tests/datastream/test_rules_engine.py +++ b/tests/datastream/test_rules_engine.py @@ -239,6 +239,78 @@ 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 TestRulesEngineFileNotFound: async def test_missing_wasm_raises(self) -> None: engine = RulesEngineClient(wasm_path="/nonexistent/rulesengine.wasm") From be5ca837432f7a605d899506664fe1985806079f Mon Sep 17 00:00:00 2001 From: Ben Papillon Date: Tue, 15 Sep 2026 12:30:42 -0700 Subject: [PATCH 2/8] guard a settle with no hold and match a plain 402 --- README.md | 18 +- src/schematic/client.py | 356 ++++++++++++++++---------- tests/custom/test_client.py | 168 +++++++++++- tests/datastream/test_rules_engine.py | 104 ++++++++ 4 files changed, 506 insertions(+), 140 deletions(-) diff --git a/README.md b/README.md index 23565c74..d2c8bb9e 100644 --- a/README.md +++ b/README.md @@ -597,7 +597,7 @@ result = client.check( "inference", company={"id": "your-company-id"}, options=CheckOptions( - usage=max_tokens, # upper bound for this operation + usage=1000, # upper bound for this operation event_subtype="inference_tokens", # the metered event ), ) @@ -606,7 +606,16 @@ if not result.allowed: inference = run_inference() -client.track_with_reservation(result.reservation, inference.tokens_used) +# 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: + client.track_with_reservation(result.reservation, inference.tokens_used) +else: + client.track( + "inference_tokens", + company={"id": "your-company-id"}, + quantity=inference.tokens_used, + ) ``` `AsyncSchematic` mirrors both methods: `await client.check(...)` and @@ -620,8 +629,9 @@ beats denying it, and the check returns your default value is different: the server knows the credits are not there, so the check denies whatever `on_acquire_failure` says. -`mode` defaults to `auto`, which means server mode: every check with `usage` -is one API call. +`mode` defaults to `auto`, which picks client mode, where leases are carved up +locally over DataStream, when DataStream is enabled, and server mode otherwise. +Client mode lands in this same release. If nothing settles a reservation, its hold is refunded at `default_reservation_ttl`. The settling event carries an idempotency key diff --git a/src/schematic/client.py b/src/schematic/client.py index 4487133f..2096ba1c 100644 --- a/src/schematic/client.py +++ b/src/schematic/client.py @@ -1,12 +1,14 @@ import atexit import datetime as dt import logging +import math from dataclasses import dataclass 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 .core.api_error import ApiError from .core.request_options import RequestOptions from .datastream import DataStreamClient, DataStreamClientOptions from .errors import PaymentRequiredError @@ -15,6 +17,7 @@ from .http_client import AsyncOfflineHTTPClient, OfflineHTTPClient from .logging import DEFAULT_LOG_LEVEL, LogLevel, get_default_logger from .types import ( + CheckAndReserveFlagResponseData, CheckFlagRequestBody, CheckFlagResponseData, CreateEventRequestBody, @@ -24,6 +27,7 @@ EventBodyIdentifyCompany, EventBodyTrack, FeatureEntitlement, + FlagCheckReservationResponseData, PreflightEventUsageRequestBody, PreflightRequestBody, RulesengineCheckFlagResult, @@ -45,6 +49,10 @@ # seconds, like every other duration on this client. DEFAULT_RESERVATION_TTL = 60.0 +# 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 + # 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. @@ -100,7 +108,8 @@ class CreditLeaseConfig: mode: CreditLeaseMode = "auto" # How long the server holds credits for an unsettled reservation, in # seconds. Size it above the longest expected gap between check() and - # track_with_reservation(). The server caps a hold at one hour. + # track_with_reservation(). Anything above the server's one hour cap is + # clamped to MAX_RESERVATION_TTL. default_reservation_ttl: float = DEFAULT_RESERVATION_TTL @@ -108,9 +117,10 @@ class CreditLeaseConfig: class CheckOptions: """Options accepted by ``check``.""" - # Units of the feature this operation will consume. The check holds - # usage * consumption_rate credits. A check takes at most one hold. - usage: Optional[int] = None + # 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 @@ -192,33 +202,49 @@ def _build_preflight(options: Optional[CheckFlagOptions]) -> Optional[PreflightR ) +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 the server would size a hold from any + of them with nothing rejecting it. + """ + if isinstance(value, bool) or not isinstance(value, (int, float)): + return False + return math.isfinite(value) and value >= 0 + + +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. + 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: + 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=options.usage) + flag_options.event_usage = EventUsage(event_subtype=options.event_subtype, quantity=quantity) else: - flag_options.usage = options.usage + flag_options.usage = quantity return flag_options -def _is_valid_quantity(value: Any) -> bool: - """Whether a caller-supplied usage can size a credit hold. - - A bool is an int in Python and a float can be NaN, and the server would - size a hold from either without any comparison rejecting it. - """ - return isinstance(value, int) and not isinstance(value, bool) and value >= 0 - - def _resolve_lease_mode( credit_leases: Optional[CreditLeaseConfig], offline: bool, ) -> Optional[Literal["client", "server"]]: @@ -256,6 +282,21 @@ def _warn_credit_lease_config( ) +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.""" + if credit_leases is None: + return DEFAULT_RESERVATION_TTL + ttl = credit_leases.default_reservation_ttl + if ttl > MAX_RESERVATION_TTL: + logger.warning( + f"credit_leases.default_reservation_ttl of {ttl}s is above the server's one hour cap; " + f"holds expire after {MAX_RESERVATION_TTL}s" + ) + return MAX_RESERVATION_TTL + return ttl + + def _reservation_request_kwargs(options: CheckOptions) -> Dict[str, Any]: """Preflight body and per-check request options for a check-and-reserve call, each omitted when the caller set nothing.""" @@ -269,14 +310,118 @@ def _reservation_request_kwargs(options: CheckOptions) -> Dict[str, Any]: return kwargs -def _payment_required_message(error: PaymentRequiredError) -> str: +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.""" - message = getattr(error.body, "error", None) + 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", + 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 + + +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) + + def _build_reservation_track_event( reservation: Reservation, actual_quantity: int, @@ -419,11 +564,7 @@ def __init__(self, api_key: str, config: Optional[SchematicConfig] = None): ) self.offline = config.offline self._credit_leases = config.credit_leases - self._reservation_ttl = ( - config.credit_leases.default_reservation_ttl - if config.credit_leases is not None - else DEFAULT_RESERVATION_TTL - ) + 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) @@ -666,12 +807,17 @@ def _check_with_server_reservation( 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 non-negative integer" + "must be a finite, non-negative number" ) - return self._server_failure_result(flag_key, options, "invalid_usage") + return failure("invalid_usage") if options.usage == 0: self.logger.debug( @@ -689,29 +835,13 @@ def _check_with_server_reservation( **_reservation_request_kwargs(options), ) data = resp.data - except PaymentRequiredError as e: - # 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", - flag_key=flag_key, - error=_payment_required_message(e), - ) 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 self._server_failure_result(flag_key, options, "server_reservation_failed") - - result = 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, - ) + 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. @@ -735,31 +865,10 @@ def _check_with_server_reservation( self.logger.warning( f"Server reservation: failed to release {held.id} ({e}); its hold is refunded when it expires" ) - return self._server_failure_result(flag_key, options, "missing_event_subtype") - - result.reservation = 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, - ) - return result + return _missing_event_subtype_result(options, result) - def _server_failure_result(self, flag_key: str, options: CheckOptions, reason: str) -> CheckResult: - if options.on_acquire_failure == "fail-closed": - return CheckResult(allowed=False, value=False, reason=reason, flag_key=flag_key, error=reason) - value = self._resolve_default(flag_key, _check_options_to_flag_options(options)) - return CheckResult( - allowed=value, value=value, reason=f"{reason}_fail_open", flag_key=flag_key, error=reason, - ) + result.reservation = _server_reservation(held, event_subtype, company, user) + return result def identify( self, @@ -803,8 +912,8 @@ def track( def track_with_reservation( self, - reservation: Reservation, - actual_quantity: int, + reservation: Optional[Reservation], + actual_quantity: float, options: Optional[TrackWithReservationOptions] = None, ) -> None: """Settle a reservation issued by ``check`` with the actual usage. @@ -817,18 +926,27 @@ def track_with_reservation( """ 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 non-negative integer. Skipping the settle, the hold is " + 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, actual_quantity, options), + _build_reservation_track_event(reservation, _settled_quantity(actual_quantity), options), options=TrackOptions(idempotency_key=f"{RESERVATION_TRACK_IDEMPOTENCY_PREFIX}{reservation.id}"), ) @@ -950,11 +1068,7 @@ def __init__(self, api_key: str, config: Optional[AsyncSchematicConfig] = None): self._shutdown_requested = False self._is_shutting_down = False self._credit_leases = config.credit_leases - self._reservation_ttl = ( - config.credit_leases.default_reservation_ttl - if config.credit_leases is not None - else DEFAULT_RESERVATION_TTL - ) + 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) @@ -1317,12 +1431,17 @@ async def _check_with_server_reservation( 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 non-negative integer" + "must be a finite, non-negative number" ) - return self._server_failure_result(flag_key, options, "invalid_usage") + return failure("invalid_usage") if options.usage == 0: self.logger.debug( @@ -1340,29 +1459,13 @@ async def _check_with_server_reservation( **_reservation_request_kwargs(options), ) data = resp.data - except PaymentRequiredError as e: - # 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", - flag_key=flag_key, - error=_payment_required_message(e), - ) 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 self._server_failure_result(flag_key, options, "server_reservation_failed") - - result = 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, - ) + 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. @@ -1386,31 +1489,10 @@ async def _check_with_server_reservation( self.logger.warning( f"Server reservation: failed to release {held.id} ({e}); its hold is refunded when it expires" ) - return self._server_failure_result(flag_key, options, "missing_event_subtype") - - result.reservation = 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, - ) - return result + return _missing_event_subtype_result(options, result) - def _server_failure_result(self, flag_key: str, options: CheckOptions, reason: str) -> CheckResult: - if options.on_acquire_failure == "fail-closed": - return CheckResult(allowed=False, value=False, reason=reason, flag_key=flag_key, error=reason) - value = self._resolve_default(flag_key, _check_options_to_flag_options(options)) - return CheckResult( - allowed=value, value=value, reason=f"{reason}_fail_open", flag_key=flag_key, error=reason, - ) + result.reservation = _server_reservation(held, event_subtype, company, user) + return result async def identify( self, @@ -1471,8 +1553,8 @@ async def _update_company_metrics( async def track_with_reservation( self, - reservation: Reservation, - actual_quantity: int, + reservation: Optional[Reservation], + actual_quantity: float, options: Optional[TrackWithReservationOptions] = None, ) -> None: """Settle a reservation issued by ``check`` with the actual usage. @@ -1485,23 +1567,33 @@ async def track_with_reservation( """ 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 non-negative integer. Skipping the settle, the hold is " + 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) await self._enqueue_event( "track", - _build_reservation_track_event(reservation, actual_quantity, options), + _build_reservation_track_event(reservation, quantity, options), 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, actual_quantity) + await self._update_company_metrics(reservation.company, reservation.event_subtype, quantity) async def _enqueue_event( self, diff --git a/tests/custom/test_client.py b/tests/custom/test_client.py index a1bb4938..569bad7f 100644 --- a/tests/custom/test_client.py +++ b/tests/custom/test_client.py @@ -8,6 +8,7 @@ from schematic.cache import LocalCache from schematic.client import ( + MAX_RESERVATION_TTL, REASON_FLAG_NOT_FOUND, REASON_OFFLINE, AsyncSchematic, @@ -22,7 +23,9 @@ SchematicConfig, TrackOptions, TrackWithReservationOptions, + _is_valid_quantity, ) +from schematic.core.api_error import ApiError as CoreApiError from schematic.errors import PaymentRequiredError from schematic.types import ( ApiError, @@ -1651,6 +1654,20 @@ def test_plain_check_still_caches(self): 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.""" @@ -1738,6 +1755,38 @@ def test_forwards_the_per_check_timeout(self): kwargs = self.schematic.features.check_and_reserve_flag.call_args.kwargs self.assertEqual(kwargs["request_options"], {"timeout": 2.5}) + 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: + self.assertEqual(client._reservation_ttl, MAX_RESERVATION_TTL) + warning = " ".join(str(call.args[0]) for call in client.logger.warning.call_args_list) + self.assertIn("one hour cap", 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, @@ -1779,6 +1828,23 @@ def test_payment_required_denies_even_with_fail_open(self): 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") + 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)) @@ -1828,7 +1894,7 @@ def test_zero_usage_falls_back_to_a_plain_check(self): self.assertIsNone(result.reservation) def test_invalid_usage_resolves_through_the_failure_contract(self): - for usage in (-5, 1.5, True): + 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] @@ -1867,6 +1933,24 @@ def test_a_failed_release_is_swallowed(self): 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: @@ -1937,10 +2021,23 @@ def test_track_with_reservation_settles_by_reservation_id(self): 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, 2.5, True): + 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 @@ -2104,6 +2201,21 @@ async def test_forwards_the_per_check_timeout(self): kwargs = self.client.features.check_and_reserve_flag.call_args.kwargs assert kwargs["request_options"] == {"timeout": 2.5} + 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 + warning = " ".join(str(call.args[0]) for call in client.logger.warning.call_args_list) + assert "one hour cap" 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, @@ -2129,6 +2241,23 @@ async def test_payment_required_denies_even_with_fail_open(self): assert result.reason == "insufficient_credits" 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" + 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)) @@ -2174,7 +2303,7 @@ async def test_zero_usage_falls_back_to_a_plain_check(self): assert result.reservation is None async def test_invalid_usage_resolves_through_the_failure_contract(self): - for usage in (-5, 1.5, True): + 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] ) @@ -2212,6 +2341,24 @@ async def test_a_failed_release_is_swallowed(self): 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: @@ -2289,10 +2436,23 @@ async def test_track_with_reservation_updates_datastream_company_metrics(self): 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, 2.5, True): + 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 diff --git a/tests/datastream/test_rules_engine.py b/tests/datastream/test_rules_engine.py index 7a9e3502..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, ) @@ -311,6 +313,108 @@ async def test_event_usage_and_credit_cost_are_carried_snake_cased(self, engine: } +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") From ba589f1d21f6d96a29ab059626dab24adc910902 Mon Sep 17 00:00:00 2001 From: Ben Papillon Date: Tue, 15 Sep 2026 12:40:28 -0700 Subject: [PATCH 3/8] add client-mode lease stores, manager, and conformance runner --- conformance/README.md | 10 + conformance/SPEC.md | 446 +++++++++++++++ conformance/vectors/check-flow.json | 527 ++++++++++++++++++ conformance/vectors/crash-windows.json | 271 +++++++++ conformance/vectors/expiry.json | 198 +++++++ conformance/vectors/lease-lifecycle.json | 447 +++++++++++++++ conformance/vectors/lease-manager.json | 417 ++++++++++++++ .../vectors/reservation-lifecycle.json | 353 ++++++++++++ conformance/vectors/track-settle.json | 194 +++++++ poetry.lock | 181 +++--- pyproject.toml | 1 + src/schematic/leases/__init__.py | 59 ++ src/schematic/leases/lease_manager.py | 405 ++++++++++++++ src/schematic/leases/lease_store.py | 270 +++++++++ src/schematic/leases/redis_lease_store.py | 349 ++++++++++++ .../leases/redis_reservation_store.py | 278 +++++++++ src/schematic/leases/reservation_store.py | 125 +++++ src/schematic/leases/types.py | 151 +++++ tests/conformance/test_vectors.py | 344 ++++++++++++ tests/lease_support.py | 215 +++++++ tests/leases/__init__.py | 0 tests/leases/conftest.py | 29 + tests/leases/test_crash_windows.py | 207 +++++++ tests/leases/test_lease_manager.py | 404 ++++++++++++++ tests/leases/test_lease_store.py | 229 ++++++++ tests/leases/test_redis_lease_store.py | 298 ++++++++++ tests/leases/test_redis_reservation_store.py | 226 ++++++++ tests/leases/test_reservation_store.py | 133 +++++ tests/leases/test_wire_client.py | 100 ++++ 29 files changed, 6805 insertions(+), 62 deletions(-) create mode 100644 conformance/README.md create mode 100644 conformance/SPEC.md create mode 100644 conformance/vectors/check-flow.json create mode 100644 conformance/vectors/crash-windows.json create mode 100644 conformance/vectors/expiry.json create mode 100644 conformance/vectors/lease-lifecycle.json create mode 100644 conformance/vectors/lease-manager.json create mode 100644 conformance/vectors/reservation-lifecycle.json create mode 100644 conformance/vectors/track-settle.json create mode 100644 src/schematic/leases/__init__.py create mode 100644 src/schematic/leases/lease_manager.py create mode 100644 src/schematic/leases/lease_store.py create mode 100644 src/schematic/leases/redis_lease_store.py create mode 100644 src/schematic/leases/redis_reservation_store.py create mode 100644 src/schematic/leases/reservation_store.py create mode 100644 src/schematic/leases/types.py create mode 100644 tests/conformance/test_vectors.py create mode 100644 tests/lease_support.py create mode 100644 tests/leases/__init__.py create mode 100644 tests/leases/conftest.py create mode 100644 tests/leases/test_crash_windows.py create mode 100644 tests/leases/test_lease_manager.py create mode 100644 tests/leases/test_lease_store.py create mode 100644 tests/leases/test_redis_lease_store.py create mode 100644 tests/leases/test_redis_reservation_store.py create mode 100644 tests/leases/test_reservation_store.py create mode 100644 tests/leases/test_wire_client.py 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/leases/__init__.py b/src/schematic/leases/__init__.py new file mode 100644 index 00000000..aa4188f3 --- /dev/null +++ b/src/schematic/leases/__init__.py @@ -0,0 +1,59 @@ +"""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 .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 .types import ( + DEFAULT_LEASE_DURATION, + DEFAULT_LEASE_SIZE, + DEFAULT_LOW_WATER_MARK, + DEFAULT_PREWARM_RESOLVE_TIMEOUT, + DEFAULT_RESERVATION_TTL, + DEFAULT_SWEEP_INTERVAL, + MAX_RESERVATION_TTL, + Clock, + LeaseConfig, + LeaseConfigOverride, + LeaseState, + ReservationRecord, + ResolvedLeaseConfig, + resolve_lease_config, +) + +__all__ = [ + "Clock", + "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", + "MAX_RESERVATION_TTL", + "RedisLeaseStore", + "RedisReservationStore", + "ReservationRecord", + "ReservationStore", + "ResolvedLeaseConfig", + "lease_key", + "resolve_lease_config", +] diff --git a/src/schematic/leases/lease_manager.py b/src/schematic/leases/lease_manager.py new file mode 100644 index 00000000..33114767 --- /dev/null +++ b/src/schematic/leases/lease_manager.py @@ -0,0 +1,405 @@ +"""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, + ) -> LeaseGrant: ... + + async def extend(self, lease_id: str, additional_amount: float, expires_at: float) -> 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 + + async def acquire( + self, + company_id: str, + credit_type_id: str, + requested_amount: float, + expires_at: float, + ) -> 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._request_options, + ) + return _grant_from_response(response) + + async def extend(self, lease_id: str, additional_amount: float, expires_at: float) -> LeaseGrant: + response = await self._credits.extend_credit_lease( + lease_id, + additional_amount=additional_amount, + expires_at=_to_datetime(expires_at), + request_options=self._request_options, + ) + 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) -> Optional[LeaseState]: + """The slot's live lease, acquiring one over the wire if none is live.""" + 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) + ) + + async def _acquire(self, company_id: str, credit_type_id: str) -> 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, + ) + 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, + ) -> 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) + ) + + async def _extend( + self, + entry: LeaseState, + resolved: ResolvedLeaseConfig, + required_credits: Optional[float], + ) -> 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, + ) + # 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 + + 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/types.py b/src/schematic/leases/types.py new file mode 100644 index 00000000..030f58e3 --- /dev/null +++ b/src/schematic/leases/types.py @@ -0,0 +1,151 @@ +"""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 + +from dataclasses import dataclass, field +from typing import 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 +# The server refuses to hold credits for longer than an hour, so a larger +# configured TTL would have the local sweeper trail the server's own release. +MAX_RESERVATION_TTL = 3600.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), + # Clamped rather than rejected: a TTL past the server's cap would have + # the sweeper refund a hold the server already released. + reservation_ttl=min(pick("reservation_ttl", DEFAULT_RESERVATION_TTL), MAX_RESERVATION_TTL), + lease_size=pick("lease_size", DEFAULT_LEASE_SIZE), + low_water_mark=pick("low_water_mark", DEFAULT_LOW_WATER_MARK), + ) diff --git a/tests/conformance/test_vectors.py b/tests/conformance/test_vectors.py new file mode 100644 index 00000000..52b0ee17 --- /dev/null +++ b/tests/conformance/test_vectors.py @@ -0,0 +1,344 @@ +"""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. + +Flow-level vectors (``check`` / ``track``) skip until the check/track port +lands: fill in ``_op_check`` and ``_op_track``, drop them from ``FLOW_OPS``, +and the same vectors start running. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional +from unittest import mock + +import pytest +from lease_support import CrashingRefundLeaseStore, ScriptedWireClient, VirtualClock, make_fake_redis + +from schematic.leases import ( + InMemoryLeaseStore, + InMemoryReservationStore, + LeaseConfig, + LeaseManager, + LeaseState, + LeaseStore, + RedisLeaseStore, + RedisReservationStore, + ReservationRecord, + ReservationStore, +) + +VECTORS_DIR = Path(__file__).resolve().parents[2] / "conformance" / "vectors" +BACKENDS = ("in_memory", "redis") +# Ops belonging to the check/track orchestration, which this SDK has yet to +# port. A vector using one of them skips rather than half-runs. +FLOW_OPS = {"check", "track"} + + +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, ReservationRecord] = {} + 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: + if any(op["op"] in FLOW_OPS for op in vector["operations"]): + pytest.skip("flow ops land with the check/track port") + 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: + # Part 2 (the check/track port) fills this in and drops "check" from + # FLOW_OPS; the vectors then run unchanged. + raise NotImplementedError("check flow not ported yet") + + +async def _op_track(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: + # See _op_check. + raise NotImplementedError("track flow not ported yet") + + +_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/lease_support.py b/tests/lease_support.py new file mode 100644 index 00000000..5ed09352 --- /dev/null +++ b/tests/lease_support.py @@ -0,0 +1,215 @@ +"""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 a scriptable stand-in for the lease wire API. +""" + +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 + +# 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, + ) -> LeaseGrant: + self.acquire_calls.append( + { + "company_id": company_id, + "credit_type_id": credit_type_id, + "requested_amount": requested_amount, + "expires_at": expires_at, + } + ) + 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) -> LeaseGrant: + self.extend_calls.append( + {"lease_id": lease_id, "additional_amount": additional_amount, "expires_at": expires_at} + ) + 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) 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_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..bbb26d09 --- /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_clamps_the_reservation_ttl() -> None: + from schematic.leases import MAX_RESERVATION_TTL, resolve_lease_config + + resolved = resolve_lease_config(LeaseConfig(reservation_ttl=MAX_RESERVATION_TTL * 2), None, "ct_1") + # The server refuses to hold credits longer than this, so a bigger TTL + # would have the sweeper trail the server's own release. + assert resolved.reservation_ttl == MAX_RESERVATION_TTL 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_wire_client.py b/tests/leases/test_wire_client.py new file mode 100644 index 00000000..4ee120c5 --- /dev/null +++ b/tests/leases/test_wire_client.py @@ -0,0 +1,100 @@ +"""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 From 4620c25f87e135ddd137b6a2e19bfb9293399fc1 Mon Sep 17 00:00:00 2001 From: Ben Papillon Date: Tue, 15 Sep 2026 13:23:49 -0700 Subject: [PATCH 4/8] add client-mode leases, prewarm, and the lease-gated check flow --- README.md | 143 +++- src/schematic/cache/redis.py | 10 + src/schematic/client.py | 514 +++++++++++--- src/schematic/datastream/datastream_client.py | 33 +- src/schematic/event_buffer.py | 10 + src/schematic/leases/__init__.py | 16 + src/schematic/leases/check.py | 524 ++++++++++++++ src/schematic/leases/lease_manager.py | 13 + src/schematic/leases/track.py | 82 +++ src/schematic/leases/types.py | 16 +- tests/conformance/test_vectors.py | 145 +++- tests/custom/test_client.py | 411 ++++++++++- tests/lease_support.py | 135 +++- tests/leases/test_check_and_track.py | 651 ++++++++++++++++++ tests/leases/test_wasm_credit_gate.py | 300 ++++++++ 15 files changed, 2868 insertions(+), 135 deletions(-) create mode 100644 src/schematic/leases/check.py create mode 100644 src/schematic/leases/track.py create mode 100644 tests/leases/test_check_and_track.py create mode 100644 tests/leases/test_wasm_credit_gate.py diff --git a/README.md b/README.md index d2c8bb9e..f15a3bcf 100644 --- a/README.md +++ b/README.md @@ -570,30 +570,72 @@ client = Schematic( 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. The server evaluates the flag against -the company's real balance and takes the hold in one call, then refunds the -unspent slice when the settling event arrives. +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. -Opt in with `credit_leases`: +> 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 unsettled credits, max 1 hour + default_reservation_ttl=60.0, # seconds the server holds credits, max 1 hour ), ) client = Schematic("YOUR_API_KEY", config) ``` -Then reserve the operation's upper bound, do the work, and report what it -actually used: +### 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 = client.check( +result = await client.check( "inference", company={"id": "your-company-id"}, options=CheckOptions( @@ -604,45 +646,88 @@ result = client.check( if not result.allowed: raise RuntimeError("credit balance exceeded") -inference = run_inference() +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: - client.track_with_reservation(result.reservation, inference.tokens_used) + await client.track_with_reservation(result.reservation, inference.tokens_used) else: - client.track( + await client.track( "inference_tokens", company={"id": "your-company-id"}, quantity=inference.tokens_used, ) ``` -`AsyncSchematic` mirrors both methods: `await client.check(...)` and -`await client.track_with_reservation(...)`. +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. -A check that cannot gate, because the API is unreachable or errored, fails -closed by default: `allowed` is False and no hold is taken. Pass -`on_acquire_failure="fail-open"` for callers where letting traffic through -beats denying it, and the check returns your default value -(`CheckOptions.default_value`, else the client's flag default) instead. A 402 -is different: the server knows the credits are not there, so the check denies -whatever `on_acquire_failure` says. +### Pre-warming -`mode` defaults to `auto`, which picks client mode, where leases are carved up -locally over DataStream, when DataStream is enabled, and server mode otherwise. -Client mode lands in this same release. +A session's first check pays the lease acquire round trip. Warm the lease when +the user is identified instead: -If nothing settles a reservation, its hold is refunded at -`default_reservation_ttl`. The settling event carries an idempotency key -derived from the reservation ID, so a retried or duplicated settle is billed -once. +```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. | +| `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. | +| `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). + +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 @@ -740,7 +825,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/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 2096ba1c..ef8cd639 100644 --- a/src/schematic/client.py +++ b/src/schematic/client.py @@ -1,13 +1,22 @@ +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, 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 @@ -15,6 +24,27 @@ 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, @@ -49,6 +79,9 @@ # 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 @@ -99,18 +132,40 @@ 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" and "auto" both take the hold over the - # check-and-reserve API. "client" (local leases carved out in-process) is - # not implemented in this SDK yet: it leaves checks ungated and warns at - # construction. + # 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 the server holds credits for an unsettled reservation, in - # seconds. Size it above the longest expected gap between check() and - # track_with_reservation(). Anything above the server's one hour cap is - # clamped to MAX_RESERVATION_TTL. + # 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 to MAX_RESERVATION_TTL. 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 @@ -202,18 +257,6 @@ def _build_preflight(options: Optional[CheckFlagOptions]) -> Optional[PreflightR ) -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 the server would size a hold from any - of them with nothing rejecting it. - """ - if isinstance(value, bool) or not isinstance(value, (int, float)): - return False - return math.isfinite(value) and value >= 0 - - def _preflight_quantity(usage: float) -> int: """Cast a usage onto the integer the preflight body carries. @@ -245,40 +288,74 @@ def _check_options_to_flag_options(options: Optional[CheckOptions]) -> Optional[ return flag_options -def _resolve_lease_mode( - credit_leases: Optional[CreditLeaseConfig], offline: bool, -) -> Optional[Literal["client", "server"]]: - """Which reservation mode a ``check`` with usage resolves to right now. +# 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", +) - None means no credit gating at all: credit leases are not configured, the - client is offline, or the caller asked for a mode this SDK cannot serve. - """ - if credit_leases is None or offline: - return None - if credit_leases.mode == "client": - # Client-side leases do not exist in this SDK yet, and gating on a - # lease that was never taken would be a lie, so the check stays plain. - # A later PR returns "client" here without moving the public surface. - return None - # Other SDKs let "auto" pick client mode when DataStream is ready. With no - # client mode to pick, both "auto" and "server" mean server. - return "server" + +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, + 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 anything.""" + 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 + if credit_leases.mode == "server": + ignored = [name for name in _CLIENT_ONLY_LEASE_OPTIONS if getattr(credit_leases, name) is not None] + if ignored: + logger.warning( + f"credit_leases.mode is 'server', so {', '.join(ignored)} will be ignored; those options only " + "apply to client mode, where leases are carved up locally over DataStream." + ) + 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', which this SDK does not support yet; check() falls back to a " - "plain, ungated flag check. Use 'server' (or the 'auto' default) to gate on credits." + "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." ) @@ -413,35 +490,6 @@ def _missing_event_subtype_result(options: CheckOptions, result: CheckResult) -> return result -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) - - -def _build_reservation_track_event( - reservation: Reservation, - actual_quantity: int, - options: Optional[TrackWithReservationOptions] = None, -) -> EventBodyTrack: - """Build the track event that settles a reservation.""" - return EventBodyTrack( - company=reservation.company, - event=reservation.event_subtype, - # In server mode the hold lives on the server and settles by id. Never - # send lease_id as well: the server prefers it when both are set, 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, - ) - - @dataclass class TrackOptions: """Optional metadata for a track event. @@ -476,6 +524,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( @@ -566,7 +618,9 @@ def __init__(self, api_key: str, config: Optional[SchematicConfig] = None): 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) + _warn_credit_lease_config( + self.logger, config.credit_leases, self.offline, supports_client_mode=False, + ) atexit.register(self.shutdown) @@ -741,7 +795,17 @@ def _check_flag_via_api( return self._default_response(flag_key, options, f"{REASON_ERROR}: {e}") def _effective_lease_mode(self) -> Optional[Literal["client", "server"]]: - return _resolve_lease_mode(self._credit_leases, self.offline) + """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, @@ -1069,8 +1133,14 @@ def __init__(self, api_key: str, config: Optional[AsyncSchematicConfig] = None): self._is_shutting_down = False 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) + # 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 @@ -1103,6 +1173,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): @@ -1112,6 +1195,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: @@ -1119,6 +1257,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() @@ -1366,7 +1509,24 @@ async def _check_flag_via_api( return self._default_response(flag_key, options, f"{REASON_ERROR}: {e}") def _effective_lease_mode(self) -> Optional[Literal["client", "server"]]: - return _resolve_lease_mode(self._credit_leases, self.offline) + """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, @@ -1393,7 +1553,108 @@ async def check( 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_fallback(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 skips the wait, which then needs the + company id up front. + """ + company_id = company.get("id") + if company_id: + return company_id + datastream = self._datastream_client + if datastream is None or 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, @@ -1512,6 +1773,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, @@ -1559,11 +1851,15 @@ async def track_with_reservation( ) -> 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. + 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 @@ -1586,15 +1882,58 @@ async def track_with_reservation( ) 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", - _build_reservation_track_event(reservation, quantity, options), + 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, @@ -1654,6 +1993,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 019161d3..756167db 100644 --- a/src/schematic/datastream/datastream_client.py +++ b/src/schematic/datastream/datastream_client.py @@ -1,23 +1,31 @@ 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. @@ -448,6 +456,21 @@ async def check_flag( results: list = await asyncio.gather(*tasks) 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 update_company_metrics(self, keys: Dict[str, str], event: str, quantity: int) -> None: """Update company metrics locally in cache (for track events).""" company = await self._get_company_from_cache(keys) 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 index aa4188f3..5698ef80 100644 --- a/src/schematic/leases/__init__.py +++ b/src/schematic/leases/__init__.py @@ -8,11 +8,18 @@ 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, @@ -27,11 +34,14 @@ LeaseState, ReservationRecord, ResolvedLeaseConfig, + is_valid_quantity, resolve_lease_config, ) __all__ = [ + "CheckDataStream", "Clock", + "CreditCheckDeps", "CreditsWireClient", "DEFAULT_LEASE_DURATION", "DEFAULT_LEASE_SIZE", @@ -51,9 +61,15 @@ "MAX_RESERVATION_TTL", "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..bd8111e0 --- /dev/null +++ b/src/schematic/leases/check.py @@ -0,0 +1,524 @@ +"""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, + ) + + lease = await deps.manager.acquire_if_needed(resolved_company.id, credit_id) + 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) + 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 index 33114767..fa75f568 100644 --- a/src/schematic/leases/lease_manager.py +++ b/src/schematic/leases/lease_manager.py @@ -275,6 +275,19 @@ async def _extend( 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. 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 index 030f58e3..9786448a 100644 --- a/src/schematic/leases/types.py +++ b/src/schematic/leases/types.py @@ -9,8 +9,9 @@ from __future__ import annotations +import math from dataclasses import dataclass, field -from typing import Callable, Dict, Mapping, Optional +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. @@ -149,3 +150,16 @@ def pick(name: str, default: float) -> float: 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 index 52b0ee17..3c3693e1 100644 --- a/tests/conformance/test_vectors.py +++ b/tests/conformance/test_vectors.py @@ -5,22 +5,32 @@ runner is the only language-specific piece; every port reimplements it and must pass the same vectors, on every store backend it ships. -Flow-level vectors (``check`` / ``track``) skip until the check/track port -lands: fill in ``_op_check`` and ``_op_track``, drop them from ``FLOW_OPS``, -and the same vectors start running. +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, ScriptedWireClient, VirtualClock, make_fake_redis +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, @@ -31,13 +41,15 @@ RedisReservationStore, ReservationRecord, ReservationStore, + check_with_lease, + consume_reservation_and_build_event, ) VECTORS_DIR = Path(__file__).resolve().parents[2] / "conformance" / "vectors" BACKENDS = ("in_memory", "redis") -# Ops belonging to the check/track orchestration, which this SDK has yet to -# port. A vector using one of them skips rather than half-runs. -FLOW_OPS = {"check", "track"} +# 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]: @@ -64,7 +76,7 @@ class Harness: def __init__(self, backend: str, config: Dict[str, Any]) -> None: self.clock = VirtualClock() - self.handles: Dict[str, ReservationRecord] = {} + self.handles: Dict[str, Reservation] = {} self.wire = ScriptedWireClient() self.leases: LeaseStore self.reservations: ReservationStore @@ -110,8 +122,6 @@ async def drain(self) -> None: @pytest.mark.parametrize("backend,vector", _load_cases()) async def test_vector(backend: str, vector: Dict[str, Any]) -> None: - if any(op["op"] in FLOW_OPS for op in vector["operations"]): - pytest.skip("flow ops land with the check/track port") 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 @@ -289,14 +299,119 @@ async def _op_release_all_local_leases(h: Harness, op: Dict[str, Any], expect: D async def _op_check(h: Harness, op: Dict[str, Any], expect: Dict[str, Any]) -> None: - # Part 2 (the check/track port) fills this in and drops "check" from - # FLOW_OPS; the vectors then run unchanged. - raise NotImplementedError("check flow not ported yet") + 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: - # See _op_check. - raise NotImplementedError("track flow not ported yet") + 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] diff --git a/tests/custom/test_client.py b/tests/custom/test_client.py index 569bad7f..a24c893d 100644 --- a/tests/custom/test_client.py +++ b/tests/custom/test_client.py @@ -1,12 +1,15 @@ +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 ( MAX_RESERVATION_TTL, REASON_FLAG_NOT_FOUND, @@ -16,6 +19,7 @@ CheckFlagOptions, CheckOptions, CreditLeaseConfig, + DataStreamConfig, EventUsage, IdentifyOptions, Reservation, @@ -27,10 +31,12 @@ ) 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, @@ -2461,5 +2467,408 @@ async def test_track_with_reservation_is_a_no_op_when_offline(self): 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_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_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_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() + + if __name__ == "__main__": unittest.main() diff --git a/tests/lease_support.py b/tests/lease_support.py index 5ed09352..5f1ae7bf 100644 --- a/tests/lease_support.py +++ b/tests/lease_support.py @@ -2,7 +2,8 @@ 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 a scriptable stand-in for the lease wire API. +need, and scriptable stand-ins for the lease wire API, the rules engine, and +DataStream. """ from __future__ import annotations @@ -14,6 +15,12 @@ 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() @@ -213,3 +220,129 @@ def make_reservation(**overrides: Any) -> ReservationRecord: } 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, + ) -> 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 + + 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_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/test_check_and_track.py b/tests/leases/test_check_and_track.py new file mode 100644 index 00000000..945dad89 --- /dev/null +++ b/tests/leases/test_check_and_track.py @@ -0,0 +1,651 @@ +"""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 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 From cc4aecfd2d01fd2eb34611a48f93130b71096d66 Mon Sep 17 00:00:00 2001 From: Ben Papillon Date: Tue, 15 Sep 2026 14:27:37 -0700 Subject: [PATCH 5/8] thread the check timeout into lease calls and warm from cache at timeout 0 --- README.md | 2 +- src/schematic/client.py | 16 ++++-- src/schematic/datastream/datastream_client.py | 9 ++++ src/schematic/leases/check.py | 6 ++- src/schematic/leases/lease_manager.py | 52 +++++++++++++++---- tests/custom/test_client.py | 26 ++++++++++ tests/lease_support.py | 22 +++++++- tests/leases/test_check_and_track.py | 24 +++++++++ tests/leases/test_wire_client.py | 31 +++++++++++ 9 files changed, 171 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index f15a3bcf..f12f9a85 100644 --- a/README.md +++ b/README.md @@ -700,7 +700,7 @@ All fields live on `CreditLeaseConfig`. Durations are seconds. Everything below | `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. | +| `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. | diff --git a/src/schematic/client.py b/src/schematic/client.py index ef8cd639..041d82d0 100644 --- a/src/schematic/client.py +++ b/src/schematic/client.py @@ -1633,14 +1633,24 @@ async def _resolve_company_id_with_wait(self, company: Dict[str, str]) -> Option 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 skips the wait, which then needs the - company id up front. + 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 or self._prewarm_resolve_timeout <= 0: + 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: diff --git a/src/schematic/datastream/datastream_client.py b/src/schematic/datastream/datastream_client.py index 756167db..2578cb96 100644 --- a/src/schematic/datastream/datastream_client.py +++ b/src/schematic/datastream/datastream_client.py @@ -471,6 +471,15 @@ def evaluate_flag( """ 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).""" company = await self._get_company_from_cache(keys) diff --git a/src/schematic/leases/check.py b/src/schematic/leases/check.py index bd8111e0..5793cfa1 100644 --- a/src/schematic/leases/check.py +++ b/src/schematic/leases/check.py @@ -199,7 +199,9 @@ async def failure(reason: str) -> "CheckResult": user_id=resolved_user.id if resolved_user is not None else None, ) - lease = await deps.manager.acquire_if_needed(resolved_company.id, credit_id) + # 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") @@ -210,7 +212,7 @@ async def failure(reason: str) -> "CheckResult": 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) + 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}") diff --git a/src/schematic/leases/lease_manager.py b/src/schematic/leases/lease_manager.py index fa75f568..141b93be 100644 --- a/src/schematic/leases/lease_manager.py +++ b/src/schematic/leases/lease_manager.py @@ -54,9 +54,16 @@ async def acquire( 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) -> 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: ... @@ -68,28 +75,42 @@ def __init__(self, credits_client: Any, *, request_options: Optional[Any] = 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._request_options, + request_options=self._options(timeout), ) return _grant_from_response(response) - async def extend(self, lease_id: str, additional_amount: float, expires_at: float) -> LeaseGrant: + 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._request_options, + request_options=self._options(timeout), ) return _grant_from_response(response) @@ -137,8 +158,15 @@ def resolve_config(self, credit_type_id: str) -> ResolvedLeaseConfig: 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) -> Optional[LeaseState]: - """The slot's live lease, acquiring one over the wire if none is live.""" + 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: @@ -158,10 +186,12 @@ async def acquire_if_needed(self, company_id: str, credit_type_id: str) -> Optio 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) + self._inflight_acquire, key, self._acquire(company_id, credit_type_id, timeout) ) - async def _acquire(self, company_id: str, credit_type_id: str) -> Optional[LeaseState]: + 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( @@ -169,6 +199,7 @@ async def _acquire(self, company_id: str, credit_type_id: str) -> Optional[Lease credit_type_id, resolved.lease_size, self._clock() + resolved.lease_duration, + timeout, ) wrote = await self._lease_store.replace( lease_id=grant.lease_id, @@ -207,6 +238,7 @@ async def maybe_extend( 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. @@ -238,7 +270,7 @@ async def maybe_extend( 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) + self._inflight_extend, key, self._extend(entry, resolved, required_credits, timeout) ) async def _extend( @@ -246,6 +278,7 @@ async def _extend( 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 @@ -257,6 +290,7 @@ async def _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 diff --git a/tests/custom/test_client.py b/tests/custom/test_client.py index a24c893d..f61af6fe 100644 --- a/tests/custom/test_client.py +++ b/tests/custom/test_client.py @@ -2709,6 +2709,32 @@ async def test_prewarm_gives_up_when_the_company_never_surfaces(self): 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: diff --git a/tests/lease_support.py b/tests/lease_support.py index 5f1ae7bf..cb5a49fe 100644 --- a/tests/lease_support.py +++ b/tests/lease_support.py @@ -155,6 +155,7 @@ async def acquire( credit_type_id: str, requested_amount: float, expires_at: float, + timeout: Optional[float] = None, ) -> LeaseGrant: self.acquire_calls.append( { @@ -162,6 +163,7 @@ async def acquire( "credit_type_id": credit_type_id, "requested_amount": requested_amount, "expires_at": expires_at, + "timeout": timeout, } ) during = self.during_acquire @@ -178,9 +180,20 @@ async def acquire( expires_at=lease["expires_at"], ) - async def extend(self, lease_id: str, additional_amount: float, expires_at: float) -> LeaseGrant: + 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} + { + "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") @@ -285,6 +298,7 @@ def __init__( 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 @@ -293,6 +307,7 @@ def __init__( 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: @@ -311,6 +326,9 @@ async def get_company(self, keys: Dict[str, str]) -> RulesengineCompany: 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 diff --git a/tests/leases/test_check_and_track.py b/tests/leases/test_check_and_track.py index 945dad89..e7a7b555 100644 --- a/tests/leases/test_check_and_track.py +++ b/tests/leases/test_check_and_track.py @@ -649,3 +649,27 @@ async def test_a_reporting_failure_leaves_the_verdict_alone(self, clock: Virtual 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_wire_client.py b/tests/leases/test_wire_client.py index 4ee120c5..6f499794 100644 --- a/tests/leases/test_wire_client.py +++ b/tests/leases/test_wire_client.py @@ -98,3 +98,34 @@ async def test_request_options_are_threaded_through() -> None: 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} From 442b81626762ace657f1e66994a9136166554349 Mon Sep 17 00:00:00 2001 From: Ben Papillon Date: Tue, 15 Sep 2026 14:39:10 -0700 Subject: [PATCH 6/8] disable retries on check-and-reserve and leave skew room in the ttl cap --- README.md | 2 +- src/schematic/client.py | 45 +++++++++++++++++++------ src/schematic/leases/__init__.py | 2 -- src/schematic/leases/types.py | 9 ++--- tests/custom/test_client.py | 54 +++++++++++++++++++++++++----- tests/leases/test_lease_manager.py | 12 +++---- 6 files changed, 90 insertions(+), 34 deletions(-) diff --git a/README.md b/README.md index f12f9a85..56bdde23 100644 --- a/README.md +++ b/README.md @@ -695,7 +695,7 @@ All fields live on `CreditLeaseConfig`. Durations are seconds. Everything below | 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. | +| `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. | diff --git a/src/schematic/client.py b/src/schematic/client.py index 041d82d0..301d6408 100644 --- a/src/schematic/client.py +++ b/src/schematic/client.py @@ -85,6 +85,15 @@ # 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 @@ -142,7 +151,8 @@ class CreditLeaseConfig: 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 to MAX_RESERVATION_TTL. + # 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 @@ -361,29 +371,44 @@ def _warn_credit_lease_config( 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.""" + 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 ttl > MAX_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"holds expire after {MAX_RESERVATION_TTL}s" + f"holds expire after {effective}s" ) - return MAX_RESERVATION_TTL + return effective return ttl def _reservation_request_kwargs(options: CheckOptions) -> Dict[str, Any]: - """Preflight body and per-check request options for a check-and-reserve - call, each omitted when the caller set nothing.""" + """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: RequestOptions = {"timeout": options.timeout} - kwargs["request_options"] = request_options + request_options["timeout"] = options.timeout + kwargs["request_options"] = request_options return kwargs @@ -414,7 +439,7 @@ def _payment_required_result(flag_key: str, error: Exception) -> CheckResult: return CheckResult( allowed=False, value=False, - reason="insufficient_credits", + reason=INSUFFICIENT_CREDITS_REASON, flag_key=flag_key, error=_payment_required_message(error), ) diff --git a/src/schematic/leases/__init__.py b/src/schematic/leases/__init__.py index 5698ef80..6dcb0832 100644 --- a/src/schematic/leases/__init__.py +++ b/src/schematic/leases/__init__.py @@ -27,7 +27,6 @@ DEFAULT_PREWARM_RESOLVE_TIMEOUT, DEFAULT_RESERVATION_TTL, DEFAULT_SWEEP_INTERVAL, - MAX_RESERVATION_TTL, Clock, LeaseConfig, LeaseConfigOverride, @@ -58,7 +57,6 @@ "LeaseState", "LeaseStore", "LeaseWireClient", - "MAX_RESERVATION_TTL", "RedisLeaseStore", "RedisReservationStore", "ReservationConsumeResult", diff --git a/src/schematic/leases/types.py b/src/schematic/leases/types.py index 9786448a..4a9d8a89 100644 --- a/src/schematic/leases/types.py +++ b/src/schematic/leases/types.py @@ -32,9 +32,6 @@ # 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 -# The server refuses to hold credits for longer than an hour, so a larger -# configured TTL would have the local sweeper trail the server's own release. -MAX_RESERVATION_TTL = 3600.0 @dataclass @@ -144,9 +141,9 @@ def pick(name: str, default: float) -> float: return ResolvedLeaseConfig( lease_duration=pick("lease_duration", DEFAULT_LEASE_DURATION), - # Clamped rather than rejected: a TTL past the server's cap would have - # the sweeper refund a hold the server already released. - reservation_ttl=min(pick("reservation_ttl", DEFAULT_RESERVATION_TTL), MAX_RESERVATION_TTL), + # 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), ) diff --git a/tests/custom/test_client.py b/tests/custom/test_client.py index f61af6fe..b584b848 100644 --- a/tests/custom/test_client.py +++ b/tests/custom/test_client.py @@ -11,9 +11,11 @@ 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, @@ -1746,7 +1748,7 @@ def test_returns_a_reservation_handle_built_from_the_response(self): ttl = dt.timedelta(seconds=TTL_SECONDS) self.assertGreaterEqual(kwargs["expires_at"], before + ttl) self.assertLessEqual(kwargs["expires_at"], after + ttl) - self.assertNotIn("request_options", kwargs) + self.assertEqual(kwargs["request_options"], {"max_retries": 0}) # The server logs the flag check for check-and-reserve itself. mock_push.assert_not_called() @@ -1759,7 +1761,14 @@ def test_sends_the_generic_usage_preflight_without_an_event_subtype(self): 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"], {"timeout": 2.5}) + 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)) @@ -1787,9 +1796,12 @@ def test_an_integral_float_usage_reaches_the_preflight_unchanged(self): def test_a_reservation_ttl_above_the_cap_is_clamped(self): client = self._client(credit_leases=CreditLeaseConfig(default_reservation_ttl=7200.0)) try: - self.assertEqual(client._reservation_ttl, MAX_RESERVATION_TTL) + # 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(str(MAX_RESERVATION_TTL - RESERVATION_TTL_SKEW_ALLOWANCE), warning) finally: client.event_buffer.stop() @@ -1830,7 +1842,7 @@ def test_payment_required_denies_even_with_fail_open(self): ) self.assertFalse(result.allowed) self.assertFalse(result.value) - self.assertEqual(result.reason, "insufficient_credits") + self.assertEqual(result.reason, INSUFFICIENT_CREDITS_REASON) self.assertEqual(result.error, "credit balance exhausted") self.assertIsNone(result.reservation) @@ -1847,7 +1859,7 @@ def test_a_402_api_error_denies_even_with_fail_open(self): ) self.assertFalse(result.allowed) self.assertFalse(result.value) - self.assertEqual(result.reason, "insufficient_credits") + self.assertEqual(result.reason, INSUFFICIENT_CREDITS_REASON) self.assertEqual(result.error, "credit balance exhausted") self.assertIsNone(result.reservation) @@ -2205,7 +2217,14 @@ async def test_sends_the_generic_usage_preflight_without_an_event_subtype(self): 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"] == {"timeout": 2.5} + 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)) @@ -2216,9 +2235,10 @@ async def test_a_fractional_usage_sizes_the_hold_and_rounds_the_preflight_up(sel 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 + 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 str(MAX_RESERVATION_TTL - RESERVATION_TTL_SKEW_ALLOWANCE) in warning finally: await client.event_buffer.stop() @@ -2244,7 +2264,7 @@ async def test_payment_required_denies_even_with_fail_open(self): options=CheckOptions(usage=50, on_acquire_failure="fail-open", default_value=True), ) assert result.allowed is False - assert result.reason == "insufficient_credits" + 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): @@ -2260,7 +2280,7 @@ async def test_a_402_api_error_denies_even_with_fail_open(self): ) assert result.allowed is False assert result.value is False - assert result.reason == "insufficient_credits" + assert result.reason == INSUFFICIENT_CREDITS_REASON assert result.error == "credit balance exhausted" assert result.reservation is None @@ -2829,6 +2849,22 @@ async def test_a_lease_gated_check_reports_one_flag_check_event(self): 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]) diff --git a/tests/leases/test_lease_manager.py b/tests/leases/test_lease_manager.py index bbb26d09..a4c0dbfc 100644 --- a/tests/leases/test_lease_manager.py +++ b/tests/leases/test_lease_manager.py @@ -395,10 +395,10 @@ async def test_resolve_config_applies_overrides(clock: VirtualClock) -> None: assert special.low_water_mark == 0.5 -def test_resolve_lease_config_clamps_the_reservation_ttl() -> None: - from schematic.leases import MAX_RESERVATION_TTL, resolve_lease_config +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=MAX_RESERVATION_TTL * 2), None, "ct_1") - # The server refuses to hold credits longer than this, so a bigger TTL - # would have the sweeper trail the server's own release. - assert resolved.reservation_ttl == MAX_RESERVATION_TTL + 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 From e7ceb7d9b93a8c2e1c1d6019e081bf7c4485499b Mon Sep 17 00:00:00 2001 From: Ben Papillon Date: Tue, 15 Sep 2026 15:17:26 -0700 Subject: [PATCH 7/8] warn about client-only options when auto resolves to server --- README.md | 2 ++ src/schematic/client.py | 12 ++++++--- tests/custom/test_client.py | 51 +++++++++++++++++++++++++++++++++++-- 3 files changed, 59 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 56bdde23..a88f2413 100644 --- a/README.md +++ b/README.md @@ -717,6 +717,8 @@ 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. diff --git a/src/schematic/client.py b/src/schematic/client.py index 301d6408..b025f97e 100644 --- a/src/schematic/client.py +++ b/src/schematic/client.py @@ -337,13 +337,17 @@ def _warn_credit_lease_config( "and holds no credits." ) return - if credit_leases.mode == "server": + # 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.mode is 'server', so {', '.join(ignored)} will be ignored; those options only " - "apply to client mode, where leases are carved up locally over DataStream." + 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": @@ -386,7 +390,7 @@ def _resolve_reservation_ttl(logger: logging.Logger, credit_leases: Optional[Cre if ttl > effective: logger.warning( f"credit_leases.default_reservation_ttl of {ttl}s is above the server's one hour cap; " - f"holds expire after {effective}s" + f"server-mode holds will be clamped to {effective}s" ) return effective return ttl diff --git a/tests/custom/test_client.py b/tests/custom/test_client.py index b584b848..9f76216d 100644 --- a/tests/custom/test_client.py +++ b/tests/custom/test_client.py @@ -1801,7 +1801,10 @@ def test_a_reservation_ttl_above_the_cap_is_clamped(self): 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(str(MAX_RESERVATION_TTL - RESERVATION_TTL_SKEW_ALLOWANCE), 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() @@ -2238,7 +2241,10 @@ async def test_a_reservation_ttl_above_the_cap_is_clamped(self): 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 str(MAX_RESERVATION_TTL - RESERVATION_TTL_SKEW_ALLOWANCE) 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() @@ -2624,6 +2630,26 @@ async def test_server_mode_warns_about_the_client_only_options(self): 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: @@ -2932,5 +2958,26 @@ def test_client_mode_points_at_the_async_client(self): 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() From 629da5e10a647b1ba1c13c75de26acebca0048b4 Mon Sep 17 00:00:00 2001 From: Ben Papillon Date: Tue, 15 Sep 2026 15:58:29 -0700 Subject: [PATCH 8/8] chore(docs): fix the stale credit lease mode comment --- src/schematic/client.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/schematic/client.py b/src/schematic/client.py index b025f97e..45d0bcf4 100644 --- a/src/schematic/client.py +++ b/src/schematic/client.py @@ -98,9 +98,10 @@ # 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": local leases carved up in-process. Not implemented in this SDK -# yet; see CreditLeaseConfig.mode. -# - "auto" (default): "server", until client mode exists here. +# - "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