« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Relay reservation and quota admission.
3Implements in-memory admission control for relay requests: an async
4``RelayReservationManager`` holds sliding token/charge windows per
5configured scope dimension (tenant, account, user, model, provider,
6channel) and reserves capacity before upstream admission. The reserve
7operation is atomic across every applicable scope under a single async
8lock, so concurrent requests cannot oversubscribe a configured limit.
10Prompt estimation for the reservation uses an injected token counter
11when available; otherwise it falls back to the explicit character
12estimate used by the LLM pricing conventions (~4 chars per token).
13Requested max output tokens are included only in the admission
14reservation and are never copied into actual usage.
15"""
17from __future__ import annotations
19import asyncio
20from dataclasses import dataclass
21from datetime import datetime, timedelta
22from decimal import Decimal
23from typing import TYPE_CHECKING
25from lexigram.ai.governance.budget import SlidingWindowCounter
26from lexigram.contracts.ai.governance import (
27 RelayBillingError,
28 RelayUsageReservation,
29 RelayUsageScope,
30 invalid_usage,
31 quota_exhausted,
32 reservation_expired,
33)
34from lexigram.contracts.ai.relay import (
35 ClaudeRequest,
36 GeminiRequest,
37 OpenAIChatRequest,
38 RelayRequestPayload,
39 ResponsesRequest,
40)
41from lexigram.contracts.core.result import Err, Ok, Result
42from lexigram.identity import ambient as identity
43from lexigram.logging import get_logger
44from lexigram.primitives import clock
45from lexigram.serialization import dumps_str
47if TYPE_CHECKING:
48 from lexigram.contracts.ai.llm import TokenCounterProtocol
50logger = get_logger(__name__)
52__all__ = [
53 "DEFAULT_RESERVATION_TTL",
54 "DEFAULT_WINDOW_SECONDS",
55 "RelayQuotaEntry",
56 "RelayQuotaSnapshot",
57 "RelayReservationLimits",
58 "RelayReservationManager",
59 "RelayScopeLimit",
60 "estimate_prompt_tokens",
61 "requested_max_output_tokens",
62]
64RELAY_DIMENSIONS = ("tenant", "account", "user", "model", "provider", "channel")
65DEFAULT_WINDOW_SECONDS = 60.0
66DEFAULT_RESERVATION_TTL = 60.0
67_CHARS_PER_TOKEN = 4
70@dataclass(frozen=True, slots=True)
71class RelayScopeLimit:
72 """Admission quota for one scope dimension.
74 Attributes:
75 max_tokens: Maximum tokens admitted in the sliding window.
76 max_charge: Maximum charge admitted in the sliding window.
77 window_seconds: Sliding window size in seconds.
78 """
80 max_tokens: int
81 max_charge: Decimal
82 window_seconds: float = DEFAULT_WINDOW_SECONDS
84 def __post_init__(self) -> None:
85 """Reject negative token/charge limits or a non-positive window."""
86 if self.max_tokens < 0:
87 raise ValueError("max_tokens must be non-negative")
88 if self.max_charge < 0:
89 raise ValueError("max_charge must be non-negative")
90 if self.window_seconds <= 0:
91 raise ValueError("window_seconds must be positive")
94@dataclass(frozen=True, slots=True)
95class RelayReservationLimits:
96 """Configured per-dimension admission limits.
98 Attributes:
99 tenant: Tenant-scope limit, when enforced.
100 account: Account-scope limit, when enforced.
101 user: User-scope limit, when enforced.
102 model: Model-scope limit, when enforced.
103 provider: Provider-scope limit, when enforced.
104 channel: Channel-scope limit, when enforced.
105 """
107 tenant: RelayScopeLimit | None = None
108 account: RelayScopeLimit | None = None
109 user: RelayScopeLimit | None = None
110 model: RelayScopeLimit | None = None
111 provider: RelayScopeLimit | None = None
112 channel: RelayScopeLimit | None = None
115@dataclass(frozen=True, slots=True)
116class _ReservationState:
117 """In-memory reservation bookkeeping.
119 Attributes:
120 reservation: The public reservation value.
121 window_keys: Sliding-window keys amounts were reserved on.
122 expires_at: Reservation expiry instant (timezone-aware).
123 """
125 reservation: RelayUsageReservation
126 window_keys: tuple[str, ...]
127 expires_at: datetime
130@dataclass(frozen=True, slots=True)
131class RelayQuotaEntry:
132 """One dimension's quota configuration and current usage.
134 Attributes:
135 dimension: Scope dimension the entry applies to.
136 value: Scope values currently tracked on the dimension.
137 max_tokens: Configured token limit for the window.
138 max_charge: Configured charge limit for the window.
139 window_seconds: Sliding window size in seconds.
140 used_tokens: Tokens currently held by live reservations.
141 used_charge: Charge currently held by live reservations.
142 """
144 dimension: str
145 value: str
146 max_tokens: int
147 max_charge: Decimal
148 window_seconds: float
149 used_tokens: int
150 used_charge: Decimal
152 def remaining_tokens(self) -> int:
153 """Return tokens still available in the window (never negative)."""
154 return max(0, self.max_tokens - self.used_tokens)
156 def remaining_charge(self) -> Decimal:
157 """Return charge still available in the window (never negative)."""
158 return max(Decimal("0"), self.max_charge - self.used_charge)
161@dataclass(frozen=True, slots=True)
162class RelayQuotaSnapshot:
163 """Read-only quota usage per configured scope dimension.
165 Attributes:
166 tenant: Tenant-dimension quota entry, when limits are configured.
167 account: Account-dimension quota entry, when limits are configured.
168 user: User-dimension quota entry, when limits are configured.
169 model: Model-dimension quota entry, when limits are configured.
170 provider: Provider-dimension quota entry, when limits are configured.
171 channel: Channel-dimension quota entry, when limits are configured.
172 """
174 tenant: RelayQuotaEntry | None = None
175 account: RelayQuotaEntry | None = None
176 user: RelayQuotaEntry | None = None
177 model: RelayQuotaEntry | None = None
178 provider: RelayQuotaEntry | None = None
179 channel: RelayQuotaEntry | None = None
182def _payload_text(payload: RelayRequestPayload) -> str:
183 """Serialize a request payload to a JSON text for estimation."""
184 return dumps_str(payload.to_dict())
187def estimate_prompt_tokens(
188 payload: RelayRequestPayload,
189 token_counter: TokenCounterProtocol | None = None,
190) -> int:
191 """Estimate the prompt token count for a request payload.
193 Args:
194 payload: Relay request payload to estimate.
195 token_counter: Optional model-aware token counter; when provided
196 it counts the serialized payload, otherwise the character
197 estimate (~4 chars per token) is used.
199 Returns:
200 A non-negative, admission-only token estimate.
201 """
202 text = _payload_text(payload)
203 if token_counter is not None:
204 return max(0, int(token_counter.count(text)))
205 return max(1, len(text) // _CHARS_PER_TOKEN)
208def requested_max_output_tokens(payload: RelayRequestPayload) -> int:
209 """Return the requested max output tokens, or 0 when not set.
211 Args:
212 payload: Relay request payload to inspect.
214 Returns:
215 The requested output budget, or 0 when the request does not
216 carry one.
217 """
218 if isinstance(payload, OpenAIChatRequest):
219 value = payload.max_completion_tokens or payload.max_tokens
220 return value or 0
221 if isinstance(payload, ResponsesRequest):
222 return payload.max_output_tokens or 0
223 if isinstance(payload, ClaudeRequest):
224 return payload.max_tokens
225 if isinstance(payload, GeminiRequest):
226 value = payload.generation_config.get("maxOutputTokens")
227 if value is None:
228 value = payload.generation_config.get("max_output_tokens")
229 return int(value) if value is not None else 0
230 return 0
233class RelayReservationManager:
234 """In-memory admission reservations across configured scope quotas.
236 Reservations hold estimated token and charge amounts in sliding
237 windows so concurrent requests cannot oversubscribe a configured
238 limit. The :meth:`reserve` operation is atomic across every
239 applicable scope under one async lock, so concurrent requests cannot
240 double-spend capacity. Expired reservations are released before a
241 new admission check; releasing is idempotent and harmless.
243 Args:
244 limits: Per-dimension quota configuration.
245 ttl_seconds: Default reservation lifetime in seconds.
246 """
248 def __init__(
249 self,
250 limits: RelayReservationLimits | None = None,
251 *,
252 ttl_seconds: float = DEFAULT_RESERVATION_TTL,
253 ) -> None:
254 self._limits = limits or RelayReservationLimits()
255 self._ttl_seconds = ttl_seconds
256 self._token_windows: dict[str, SlidingWindowCounter] = {}
257 self._charge_windows: dict[str, SlidingWindowCounter] = {}
258 self._reservations: dict[str, _ReservationState] = {}
259 self._started: set[str] = set()
260 self._lock = asyncio.Lock()
262 def _token_window(self, key: str) -> SlidingWindowCounter:
263 """Return the token window for *key*, creating it lazily."""
264 window = self._token_windows.get(key)
265 if window is None:
266 window = SlidingWindowCounter()
267 self._token_windows[key] = window
268 return window
270 def _charge_window(self, key: str) -> SlidingWindowCounter:
271 """Return the charge window for *key*, creating it lazily."""
272 window = self._charge_windows.get(key)
273 if window is None:
274 window = SlidingWindowCounter()
275 self._charge_windows[key] = window
276 return window
278 def _scope_value(self, scope: RelayUsageScope, dimension: str) -> str:
279 """Return the value of a dimension in *scope* (``""`` if unset)."""
280 if dimension == "tenant":
281 return scope.tenant_id
282 if dimension == "model":
283 return scope.model
284 if dimension == "provider":
285 return scope.provider
286 if dimension == "channel":
287 return scope.channel
288 if dimension == "account":
289 return scope.account_id or ""
290 return scope.user_id or ""
292 def _scope_keys(self, scope: RelayUsageScope) -> list[str]:
293 """Collect internal window keys for the configured scope.
295 Args:
296 scope: The scope to map onto configured dimensions.
298 Returns:
299 Window keys ``"<dim>:<value>"`` for every configured
300 dimension the scope carries a value for.
301 """
302 keys: list[str] = []
303 for dimension in RELAY_DIMENSIONS:
304 limit = getattr(self._limits, dimension)
305 if limit is None:
306 continue
307 value = self._scope_value(scope, dimension)
308 if not value:
309 logger.debug(
310 "relay_reservation_scope_empty",
311 dimension=dimension,
312 tenant_id=scope.tenant_id,
313 )
314 continue
315 keys.append(f"{dimension}:{value}")
316 return keys
318 async def reserve(
319 self,
320 request_id: str,
321 scope: RelayUsageScope,
322 estimated_tokens: int,
323 estimated_charge: Decimal,
324 *,
325 ttl_seconds: float | None = None,
326 ) -> Result[RelayUsageReservation, RelayBillingError]:
327 """Atomically reserve capacity across every applicable scope.
329 Args:
330 request_id: Gateway request identifier for the reservation.
331 scope: Accounting scope of the request.
332 estimated_tokens: Prompt estimate used for admission only,
333 including requested max output tokens folded in by the
334 caller.
335 estimated_charge: Maximum charge the reservation covers.
336 ttl_seconds: Reservation lifetime override; defaults to the
337 manager's configured TTL.
339 Returns:
340 ``Ok(reservation)`` when every applicable scope has room,
341 otherwise ``Err`` with code ``invalid_usage`` or
342 ``quota_exhausted``.
343 """
344 if estimated_tokens < 0:
345 return Err(
346 invalid_usage(
347 message="estimated_tokens must be non-negative",
348 request_id=request_id,
349 tenant_id=scope.tenant_id,
350 )
351 )
352 if estimated_charge < 0:
353 return Err(
354 invalid_usage(
355 message="estimated_charge must be non-negative",
356 request_id=request_id,
357 tenant_id=scope.tenant_id,
358 )
359 )
360 ttl = ttl_seconds if ttl_seconds is not None else self._ttl_seconds
361 if ttl <= 0:
362 return Err(
363 invalid_usage(
364 message="reservation TTL must be positive",
365 request_id=request_id,
366 tenant_id=scope.tenant_id,
367 )
368 )
370 async with self._lock:
371 now = clock.now()
372 await self._release_expired(now)
374 window_keys: list[str] = []
375 for key in self._scope_keys(scope):
376 dimension, _, _ = key.partition(":")
377 limit = getattr(self._limits, dimension)
378 if limit is None:
379 continue
380 token_total = await self._token_window(key).total()
381 if token_total + estimated_tokens > limit.max_tokens:
382 return Err(
383 quota_exhausted(
384 message=(
385 f"{dimension} token quota {limit.max_tokens} "
386 f"exceeded ({token_total} used + {estimated_tokens})"
387 ),
388 request_id=request_id,
389 tenant_id=scope.tenant_id,
390 )
391 )
392 charge_total = await self._charge_window(key).total()
393 if charge_total + float(estimated_charge) > float(limit.max_charge):
394 return Err(
395 quota_exhausted(
396 message=(
397 f"{dimension} credit quota {limit.max_charge} "
398 f"exceeded ({charge_total} used)"
399 ),
400 request_id=request_id,
401 tenant_id=scope.tenant_id,
402 )
403 )
404 window_keys.append(key)
406 reservation_id = identity.new_uuid()
407 expires_at = now + timedelta(seconds=ttl)
408 reservation = RelayUsageReservation(
409 reservation_id=reservation_id,
410 request_id=request_id,
411 estimated_tokens=estimated_tokens,
412 estimated_charge=estimated_charge,
413 expires_at=expires_at,
414 )
415 for key in window_keys:
416 await self._token_window(key).reserve(
417 reservation_id, float(estimated_tokens)
418 )
419 await self._charge_window(key).reserve(
420 reservation_id, float(estimated_charge)
421 )
422 self._reservations[reservation_id] = _ReservationState(
423 reservation=reservation,
424 window_keys=tuple(window_keys),
425 expires_at=expires_at,
426 )
427 logger.info(
428 "relay_reservation_created",
429 reservation_id=reservation_id,
430 request_id=request_id,
431 estimated_tokens=estimated_tokens,
432 estimated_charge=str(estimated_charge),
433 )
434 return Ok(reservation)
436 async def mark_started(self, reservation_id: str) -> None:
437 """Record that the upstream attempt started for a reservation.
439 A reservation marked started may be settled even after expiry.
441 Args:
442 reservation_id: Reservation identifier.
443 """
444 async with self._lock:
445 self._started.add(reservation_id)
447 async def release(self, reservation_id: str) -> None:
448 """Release a reservation and its reserved capacity.
450 Idempotent and harmless for unknown reservation IDs.
452 Args:
453 reservation_id: Reservation identifier.
454 """
455 async with self._lock:
456 await self._release_locked(reservation_id)
458 async def settle(self, reservation_id: str) -> Result[None, RelayBillingError]:
459 """Settle a reservation exactly once, freeing its capacity.
461 A reservation cannot be settled after expiry unless the upstream
462 attempt was marked started via :meth:`mark_started`. Actual
463 usage is recorded separately by the billing service.
465 Args:
466 reservation_id: Reservation identifier.
468 Returns:
469 Ok(None) on settlement, or an Err of code
470 ``reservation_expired`` when it expired before starting.
471 """
472 async with self._lock:
473 state = self._reservations.get(reservation_id)
474 if state is None:
475 return Ok(None)
476 now = clock.now()
477 if state.expires_at <= now and reservation_id not in self._started:
478 await self._release_locked(reservation_id)
479 return Err(
480 reservation_expired(
481 message="reservation expired before starting",
482 request_id=state.reservation.request_id,
483 )
484 )
485 for key in state.window_keys:
486 await self._token_window(key).release_reservation(reservation_id)
487 await self._charge_window(key).release_reservation(reservation_id)
488 self._reservations.pop(reservation_id, None)
489 self._started.discard(reservation_id)
490 logger.info("relay_reservation_settled", reservation_id=reservation_id)
491 return Ok(None)
493 async def quota_snapshot(self) -> RelayQuotaSnapshot:
494 """Report configured limits and current usage per dimension.
496 Expired reservations are released first so the snapshot reflects
497 live capacity only. Dimensions without a configured limit never
498 appear in the snapshot.
500 Returns:
501 Per-dimension quota entries aggregating every tracked window
502 value; ``None`` for unconfigured dimensions.
503 """
504 async with self._lock:
505 await self._release_expired(clock.now())
506 entries: dict[str, RelayQuotaEntry] = {}
507 for dimension in RELAY_DIMENSIONS:
508 limit = getattr(self._limits, dimension)
509 if limit is None:
510 continue
511 prefix = f"{dimension}:"
512 keys = sorted(
513 key for key in self._token_windows if key.startswith(prefix)
514 )
515 if not keys:
516 continue
517 used_tokens = 0
518 used_charge = Decimal("0")
519 for key in keys:
520 used_tokens += int(await self._token_window(key).total())
521 used_charge += Decimal(str(await self._charge_window(key).total()))
522 entries[dimension] = RelayQuotaEntry(
523 dimension=dimension,
524 value=", ".join(key.partition(":")[2] for key in keys),
525 max_tokens=limit.max_tokens,
526 max_charge=limit.max_charge,
527 window_seconds=limit.window_seconds,
528 used_tokens=used_tokens,
529 used_charge=used_charge,
530 )
531 return RelayQuotaSnapshot(
532 tenant=entries.get("tenant"),
533 account=entries.get("account"),
534 user=entries.get("user"),
535 model=entries.get("model"),
536 provider=entries.get("provider"),
537 channel=entries.get("channel"),
538 )
540 async def _release_expired(self, now: datetime) -> None:
541 """Release every reservation that expired at or before *now*."""
542 for reservation_id, state in list(self._reservations.items()):
543 if state.expires_at <= now:
544 await self._release_locked(reservation_id)
546 async def _release_locked(self, reservation_id: str) -> None:
547 """Release a reservation assuming the manager lock is held."""
548 state = self._reservations.pop(reservation_id, None)
549 if state is None:
550 return
551 for key in state.window_keys:
552 await self._token_window(key).release_reservation(reservation_id)
553 await self._charge_window(key).release_reservation(reservation_id)
554 self._started.discard(reservation_id)
555 logger.debug("relay_reservation_released", reservation_id=reservation_id)