mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
43c07a0f12
* fix(router): check and count deployment itpm and otpm on every Responses WebSocket turn The success hook only reconciled an ITPM/OTPM reservation found in top-level metadata or the standard logging object, but on the proxy the reservation lives in litellm_params.metadata, so no reservation was ever settled there. The hook now always calls the reconcile, which already returns early without a reservation and clears what it settled, so a second success event is a no-op. A Responses WebSocket turn after the first now has its deployment's current-minute ITPM and OTPM counters checked and is refused with a 429 frame once a counter is at its limit. An admitted turn carries a zero reservation on the current minute's keys, so its actual tokens land on those counters on both the native and the managed relay. * fix(responses): settle the socket open's itpm/otpm reservation on the first turn only After the first turn takes its logging copy, both websocket relays drop the open's ITPM/OTPM reservation keys from the connection metadata, the same way the budget reservation is detached. Later turns, unsolicited backend events, turns after limits are removed mid-socket, and sockets with no turn admission no longer settle or refund the open's reservation a second time. * refactor(responses): read the itpm/otpm reservation keys from core metadata helpers instead of router_utils * fix(responses): keep a socket open's hold for its first real turn and check tokens before request limits A turn that will settle the socket open's itpm/otpm reservation no longer counts that reservation against itself or swaps it for a zero one, so Codex's warmup frame on a managed provider, or a refused first frame on the native socket, no longer blocks the first real prompt or strands the hold for the minute. The token check now runs before the rpm/rpd checks, so a turn refused for tokens never uses up a request slot. Token-check tests pin the clock, and the managed relay settle test drives a real streamed call through a custom provider. * test(router): pin the router clock in the socket turn admission tests Both tests seeded this minute and the next to dodge a minute boundary. They now use the file's pinned-minute fixture and seed only that minute, so neither depends on when the test runs. * fix(router): check a responses socket's open without reserving and read shared token counters first A Responses WebSocket open now runs the same itpm/otpm check a turn runs and reserves nothing. It used to reserve an estimate that falls back to the model's max output, so a socket could not open under any otpm below that. The open still records a zero reservation for its own minute, so the first frame, which the relay forwards without a turn check, counts its actual tokens like every later turn. This drops the plumbing that only carried the open's hold into the first turn. The turn admission callable takes no arguments again, and the native relay detaches only the budget reservation, as before this change. The itpm/otpm turn check reads Redis first and falls back to this instance's in-memory count only without Redis or when the circuit breaker is open. These counters go down when a call settles, so an in-memory copy at the limit could refuse turns after the shared count dropped. The tpm check keeps its in-memory shortcut, since its counter only grows within a minute. Responses usage that reports input_tokens_details.cached_tokens now leaves cached input out of itpm, as chat usage already did. * fix(router): leave a socket turn's cached input out of itpm The native Responses WebSocket relay logs a turn with its list of events, so the deployment rate limit hook finds no usage on the response and falls back to the standard logging payload. That fallback read cached input from metadata.cache_read_input_tokens, a key the payload's cleaned metadata never carries, so every cached token of a socket turn was counted against itpm. The fallback now reads the cached count from the payload's usage_object, the same usage the logging path already built for cost, so a native socket turn counts only its uncached input. The managed relay and HTTP /v1/responses already counted correctly through the response's own usage * fix(router): mark only responses socket opens as check-only and count a turn on the minute it finishes The check-only branch for a socket open keyed on a websocket value in the request kwargs, which any HTTP body can carry and which the realtime route also sets. The router now passes an opens_responses_socket flag into set_io_token_rate_limit_request_kwargs only for its _aresponses_websocket call, and the flag lives in the same context value as the request kwargs, so no request body can set it and realtime opens reserve and get refused as before A zero reservation holds nothing to refund, so a socket turn now counts its real tokens on the minute the call finishes, read from the clock at settle time. A non-zero reservation still settles on the keys it reserved on so its refund hits the same counter When the shared Redis read returns nothing, the check now falls back to this instance's in-memory count instead of admitting the turn with no count at all * fix(router): keep io token reservation keys out of the metadata sent to providers The itpm/otpm pre-call check wrote its reservation keys into both metadata and litellm_metadata. On the routes the router serves through its generic helper (Responses, Anthropic Messages, and the rest), metadata is the client's own field and goes to the provider as is. So _litellm_itpm_cache_key and _litellm_otpm_cache_key, which carry the deployment id and the model name, reached the provider, came back in the echoed response metadata, and counted toward the provider's 16-key cap. The managed Responses socket relay merged each turn's reservation into the frame's metadata the same way, and a native socket opened with metadata sent the keys on every turn through its request defaults The router now passes the limiter the metadata field it uses for its own fields on that call (metadata for chat, litellm_metadata for the generic routes), and the reservation goes only there. metadata_channel is a required keyword, so a new caller has to name the field it writes. The managed relay puts each turn's reservation only on litellm_metadata With the reservation only on litellm_metadata, a retried call on the generic routes refunded its first attempt twice and pushed the deployment's counters below zero. The failure hook read the reservation from the logging copy in litellm_params.metadata and cleared only that copy, so the router's litellm_metadata still held it and the retry refunded it again. The hooks now also read and clear litellm_params.litellm_metadata, which is the router's own dict The proxy does not strip client keys that start with _litellm_ on the way in, so dropping that prefix at the provider boundary would also drop keys clients send today The io token tests now run Redis in process with fakeredis instead of a TCP fake server * test(responses): count socket turns through real logging callbacks instead of a Logging subclass The native relay test now hands the rate limit check to a real Logging object as its success and failure callback, so the turn's tokens go through the same dispatch path production uses. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> |
||
|---|---|---|
| .. | ||
| _prompt_cache_pin_support.py | ||
| _reasoning_effort_levels_support.py | ||
| either_role.json | ||
| test_advisor_failure_cooldown.py | ||
| test_audit_any_sweep_chaos.py | ||
| test_auto_router_management.py | ||
| test_azure_passthrough_spend_row.py | ||
| test_base_model_context_window_owned_proxy.py | ||
| test_batch_mode_deployments_wire.py | ||
| test_complexity_router_codex_compaction.py | ||
| test_complexity_router_databricks_classifier.py | ||
| test_complexity_router_llm_classifier.py | ||
| test_declared_provider_bare_model_routes.py | ||
| test_deployment_mode_routing_wire.py | ||
| test_deployment_num_retries_generic_routes_wire.py | ||
| test_encrypted_content_affinity_no_peer.py | ||
| test_end_user_region_routing.py | ||
| test_enforce_model_rate_limits_daily_limits.py | ||
| test_fallback_hop_foreign_encrypted_reasoning.py | ||
| test_fallback_retry_header_contracts.py | ||
| test_include_fallback_errors_wire.py | ||
| test_key_max_parallel_requests.py | ||
| test_key_tpm_reservation.py | ||
| test_messages_stream_retry_budget_sources_owned_proxy.py | ||
| test_model_capability_flags.py | ||
| test_non_chat_model_endpoint_gate.py | ||
| test_observed_routing.py | ||
| test_otpm_default_reservation_fraction.py | ||
| test_payment_required_cooldown.py | ||
| test_payment_required_mapping.py | ||
| test_pdf_document_pre_call_checks_owned_proxy.py | ||
| test_per_model_rate_limit_fallbacks.py | ||
| test_per_model_rate_limit_fallbacks_lifecycle.py | ||
| test_pre_call_response_format_filter.py | ||
| test_priority_model_tpm_enforcement.py | ||
| test_priority_rate_limit_headers.py | ||
| test_priority_scheduler_queue_cleanup.py | ||
| test_prompt_cache_pin_memory_routes.py | ||
| test_prompt_cache_pin_ttl.py | ||
| test_prompt_cache_pin_ttl_chaos.py | ||
| test_prompt_caching_affinity_tool_call_marks.py | ||
| test_ptu_shares.py | ||
| test_ptu_shares_regex_timing.py | ||
| test_reasoning_effort_levels_chaos.py | ||
| test_reasoning_effort_levels_wire.py | ||
| test_redis_increment_with_floor.py | ||
| test_redis_recovery.py | ||
| test_responses_websocket_turn_admission.py | ||
| test_router_strategies_count_async_requests_once.py | ||
| test_silent_model_shadow_wire.py | ||
| test_stale_cost_map_boot.py | ||
| test_tag_routing_prefix.py | ||
| test_team_tag_routing.py | ||
| test_team_wildcard_precedence.py | ||
| test_usage_based_routing_redis_reads.py | ||
| test_user_config_routing.py | ||
| test_video_poll_routing_wire.py | ||
| test_vllm_passthrough_spend_row.py | ||
| test_vllm_passthrough_stream_relay.py | ||
| test_vllm_passthrough_stream_usage_settings.py | ||
| test_wildcard_model_listing.py | ||