fix(proxy): iam refresh engine watcher race - #30183
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Greptile SummaryThis PR fixes a race condition in the RDS IAM token refresh path where a planned engine restart (killing the old Prisma query-engine subprocess) was indistinguishable from a crash, causing cascading reconnect storms. The fix introduces four coordinating mechanisms: planned-death tracking (
Confidence Score: 5/5Safe to merge. The change is a targeted fix to an IAM token refresh race condition with no production-path regressions identified. Both issues flagged in prior review rounds (the heavy reconnect path missing its generation snapshot, and the frozenset/discard mismatch) are resolved in this version. The four new coordination mechanisms are individually correct and interact safely: No files require special attention.
|
| Filename | Overview |
|---|---|
| litellm/proxy/db/prisma_client.py | Core fix: adds _expected_engine_deaths, _engine_generation, and on_engine_replaced to PrismaWrapper; splits recreate_prisma_client into a public lock-acquiring wrapper and a locked inner method; adds refresh coalescing via _token_refresh_not_needed. |
| litellm/proxy/utils.py | Adds _consume_expected_death to all four watcher death detectors; adds _handle_writer_engine_replaced hook; captures engine-generation snapshot before both reconnect closures; adds probe-first logic to _do_direct_reconnect; wires on_engine_replaced in watchdog start. |
| litellm/proxy/db/routing_prisma_wrapper.py | Forwards expected_generation to the writer's recreate call and short-circuits reader recreation when the writer recreate was a no-op; return type changed from None to bool. |
| tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py | New test file: covers planned-death marking, generation guard, refresh coalescing, routing wrapper guard forwarding, and the 64-entry set cap. |
| tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py | Adds end-to-end race repro using event synchronization to deterministically interleave a mid-recreate refresh with an in-flight transport-error reconnect; asserts exactly one Prisma construction and no kill of the fresh engine. |
| tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_engine_watcher.py | Adds planned-death coverage across all four watcher paths; tests _handle_writer_engine_replaced re-arming and hook wiring in watchdog start. |
| tests/litellm/proxy/test_prisma_engine_watchdog.py | Updated existing tests to match the new probe-first direct path and to pass expected_generation=ANY to the recreate assertion; added probe-fail setup for the escalation test. |
| tests/test_litellm/proxy/db/test_prisma_self_heal.py | Updated to reflect probe-first behavior: query_raw uses a side-effect list, assertions updated to match two query_raw calls and expected_generation=0 keyword arg. |
| tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py | Minimal update: assertion updated to include expected_generation=None in the writer recreate call to match the new signature. |
Reviews (4): Last reviewed commit: "fix(proxy): stop IAM-refresh engine rest..." | Re-trigger Greptile
f0b1cbe to
d18f8e9
Compare
PR overviewAll previously flagged issues have been addressed. No open security concerns remain on this pull request. Security reviewNo open security issues remain on this pull request. Fixed/addressed: 1 · PR risk: 0/10 |
ad04344 to
fe9f9f7
Compare
|
@greptile-apps please review |
fe9f9f7 to
886ff21
Compare
|
@greptile-apps review please |
…BerriAI#29176) An RDS IAM token refresh recreates the Prisma client, which SIGKILLs the running query-engine and spawns a new one. That planned kill was indistinguishable from a crash, and three reconnect paths used two uncoordinated locks, so a single refresh triggered a cascade of engine kill/respawn cycles: 1. `_safe_refresh_token` (holds `_reconnection_lock`) -> recreate -> kill old engine, spawn new one. 2. The engine-death watcher sees that kill, assumes a crash, and calls `attempt_db_reconnect(force=True)` (a different lock, `_db_reconnect_lock`) -> recreate again -> kills the fresh engine. 3. In-flight queries failing during the swap are classified as transport errors and trigger their own `attempt_db_reconnect` -> recreate again. Fix coordinates planned restarts across the wrapper and the watcher: - PrismaWrapper records the old engine PID in `_expected_engine_deaths` before killing it; all four watcher death-detectors (waitpid thread, pidfd, already-dead probe, os.kill poll) consume that PID and skip the reconnect instead of treating it as a crash. - `recreate_prisma_client` now serializes through `_reconnection_lock` and bumps a monotonic `_engine_generation`. Callers pass `expected_generation` as an optimistic-lock token, so racing/cascading recreates collapse into a single restart (losers no-op). This closes the two-lock gap. - The direct reconnect path probes the writer with SELECT 1 before recreating; a healthy connection (e.g. engine already replaced by a refresh) skips the recreate entirely. - `_safe_refresh_token` coalesces: it skips when the current token still has more than the refresh buffer of runway, so stacked triggers (proactive loop + __getattr__ fallback) don't each restart the engine. An `on_engine_replaced` hook re-arms the watcher on the new PID. RoutingPrismaWrapper forwards `expected_generation` and skips recreating the reader when the writer recreate was skipped.
886ff21 to
98b9914
Compare
|
@greptile-apps plz review |
|
This is looking great — Greptile 5/5 and CI is green! One last thing: could you fill in the "Screenshots / Proof of Fix" section with some evidence of the fix working (e.g. logs showing a clean IAM token refresh without the engine kill loop, or the repro test output)? It really speeds up maintainer review. Thanks! |
Screenshot/ ProofThis is a concurrency race in the IAM-refresh path (no UI, no steady-state log line to screenshot), so the proof is the before/after of the end-to-end repro test, test_iam_refresh_racing_reconnect_recreates_engine_only_once. It reproduces the exact reported scenario: an RDS IAM token refresh is mid-recreate — it has already spawned a fresh query-engine (PID 222) and is holding _reconnection_lock — when an in-flight transport-error reconnect (attempt_db_reconnect(force=True)) fires. The test asserts the engine is recreated exactly once (prisma_constructed == 1) and that the refresh's freshly-spawned engine (PID 222) is never killed (fresh_engine_killed == False). ❌ Before the fix — source reverted to the base branch, the repro test kept. The racing reconnect undoes the refresh, reproducing the kill/respawn storm from the issue: $ pytest test_iam_refresh_racing_reconnect_recreates_engine_only_once $ pytest test_iam_refresh_racing_reconnect_recreates_engine_only_once |
|
@Sameerlite added proof of fix. Could you please take a look? |
ryan-crabbe-berri
left a comment
There was a problem hiding this comment.
lgtm, out of curiousity what does your deployment look like? how many pods/workers per pod?
|
Thanks for the PR! A couple of things to get this over the finish line:
Once those are addressed we'll take another look — appreciate the contribution! |
|
@Sameerlite fixed it @ryan-crabbe-berri we currently have 2 pods and 2 workers per pod. But we also have horizontal autoscaling enabled based on in_flight_request count. |
85828da
into
BerriAI:litellm_oss_staging
…#29176) (#30183) An RDS IAM token refresh recreates the Prisma client, which SIGKILLs the running query-engine and spawns a new one. That planned kill was indistinguishable from a crash, and three reconnect paths used two uncoordinated locks, so a single refresh triggered a cascade of engine kill/respawn cycles: 1. `_safe_refresh_token` (holds `_reconnection_lock`) -> recreate -> kill old engine, spawn new one. 2. The engine-death watcher sees that kill, assumes a crash, and calls `attempt_db_reconnect(force=True)` (a different lock, `_db_reconnect_lock`) -> recreate again -> kills the fresh engine. 3. In-flight queries failing during the swap are classified as transport errors and trigger their own `attempt_db_reconnect` -> recreate again. Fix coordinates planned restarts across the wrapper and the watcher: - PrismaWrapper records the old engine PID in `_expected_engine_deaths` before killing it; all four watcher death-detectors (waitpid thread, pidfd, already-dead probe, os.kill poll) consume that PID and skip the reconnect instead of treating it as a crash. - `recreate_prisma_client` now serializes through `_reconnection_lock` and bumps a monotonic `_engine_generation`. Callers pass `expected_generation` as an optimistic-lock token, so racing/cascading recreates collapse into a single restart (losers no-op). This closes the two-lock gap. - The direct reconnect path probes the writer with SELECT 1 before recreating; a healthy connection (e.g. engine already replaced by a refresh) skips the recreate entirely. - `_safe_refresh_token` coalesces: it skips when the current token still has more than the refresh buffer of runway, so stacked triggers (proactive loop + __getattr__ fallback) don't each restart the engine. An `on_engine_replaced` hook re-arms the watcher on the new PID. RoutingPrismaWrapper forwards `expected_generation` and skips recreating the reader when the writer recreate was skipped.
* fix(proxy): bump health-check max_tokens default to 16 for GPT-5 compatibility (#30708) OpenAI GPT-5 models require max_completion_tokens >= 16. Health checks were using 5 (proxy/health_check.py) and 10 (health_check_helpers.py), causing failures on GPT-5 models. Fixes #23836 * fix: increase health check max_tokens from 5 to 16 (#23836) (#26610) GPT-5 models enforce a minimum of 16 for max_output_tokens. The current default of 5 still causes health checks to fail for these models. Bump the non-wildcard default to 16 — the smallest value that satisfies all known provider minimums while keeping health checks lightweight. Also tightens the wildcard test assertion from a weak disjunctive check to strict key-absence. Co-authored-by: Sameer Kankute <sameer@berri.ai> * fix: ensure checks show gemini-3-flash-preview supports responseJsonS… (#30696) * fix: ensure checks show gemini-3-flash-preview supports responseJsonSchema. * fix: remove async keyword from test. * fix: make Bedrock Mantle Responses routing data-driven per model (#30700) * Make Bedrock Mantle Responses routing data-driven per model Route Bedrock Mantle models to the native Responses API based on each model's price-map capability signal instead of a hardcoded model-name heuristic, and derive the OpenAI-compatible base path segment per model. Responses dispatch now selects the native config when the model advertises responses support (/v1/responses in supported_endpoints, or mode=responses), both overridable via register_model and proxy model_info. This enables native Responses for gpt-oss-120b/20b and the gemma-4 family while keeping chat-only models (gpt-oss safeguard, nvidia, mistral, ...) on the existing chat-completions emulation. Capability is per-model, so gpt-oss-120b routes natively while gpt-oss-safeguard-120b does not despite sharing the gpt-oss substring. The wire path is a separate concern, driven by the existing use_openai_responses_path flag rather than a model-name match: gpt-5.x and gemma-4-* on /openai/v1, everything else (incl. gpt-oss) on /v1. The chat config now derives its base from the same flag, fixing gemma-4 chat-completions requests that previously went to /v1 instead of /openai/v1. Cost maps: add supported_endpoints to the gpt-oss entries (responses for the non-safeguard variants, chat-only for safeguard) and supported_endpoints + use_openai_responses_path to all three gemma-4 entries. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Address review: move capability helper into bedrock_mantle package Move the Responses capability check out of utils.py into litellm/llms/bedrock_mantle/common_utils.py as mantle_supports_responses, alongside its companion wire-path helper mantle_base_segment. Both are now pure functions of (model, model_cost): the price-map mode/supported_endpoints read replaces the get_model_info call, so the rules are unit-testable without patching global state and the Bedrock Mantle package is self-contained. Use str | None instead of Optional[str] on the new signatures to satisfy the ruff UP045 strict-rule gate. Add direct unit tests for both helpers. Fix test_register_model_restore_undoes_existing_key_overwrite: gpt-oss-120b now legitimately supports Responses, so it can no longer be the "None after restore" vehicle; use the chat-only safeguard variant, which isolates the register/restore effect from the model's own capability. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Sameer Kankute <sameer@berri.ai> * fix(proxy): fail fast on non-PostgreSQL DATABASE_URL instead of hanging on startup (#30366) * fix(proxy): fail fast on non-PostgreSQL DATABASE_URL instead of hanging on startup LiteLLM's Prisma datasource is pinned to provider = 'postgresql', so a sqlite:// or mysql:// DATABASE_URL can never connect. Today that surfaces as an opaque startup stall where the port never binds, and a separate 'DB not connected' 500 on /key/generate when no DATABASE_URL is set at all leaves operators guessing what to configure. Validate the DATABASE_URL / DIRECT_URL scheme in run_server before any Prisma call and exit with an actionable message naming the unsupported scheme. Also reword CommonProxyErrors.db_not_connected_error to tell the operator to set DATABASE_URL to a postgresql:// connection string. Add regression tests covering postgres acceptance and sqlite/mysql/mssql rejection. * fix: resolve CI failures and proxy DB URL typing issue * fix(dashscope): treat an explicit 0.0 tier cost as a real price, not missing (#30653) The tiered cost calculator resolved a tier's per-token cost with `tier.get(cost_key) or tier.get(fallback_cost_key, 0)`. Because `or` short-circuits on any falsy value, a tier that legitimately prices a component at 0.0 (e.g. a free-cache-read tier with cache_read_input_token_cost: 0.0, or a free-reasoning tier) is treated as missing and silently billed at the full fallback rate (input_cost_per_token / output_cost_per_token). The flat-pricing path in the same module already handles this correctly with an `is None` guard. Resolve tier costs through a small helper that mirrors it, so 0.0 is honored at both the in-range and overflow sites. No shipped model currently has a 0.0 tier cost, so this is a latent defect; the fix makes the tiered path consistent with the flat path and prevents over-charging the first time such a tier appears. Adds unit tests covering the in-range and overflow paths, and drops an unused import flagged by ruff in the touched test file. * feat(proxy): show session-aggregate cost and duration in request logs (#25708) (#30507) * fix(anthropic): don't leak tool 'type' into OpenAI function parameters schema (#30618) In the messages->chat/completions bridge, translate_anthropic_tools_to_openai merged every non-mapped tool key into the function parameters dict. The Anthropic tool 'type' (e.g. 'custom') thus overwrote parameters.type ('object' -> 'custom'), and providers reject it ('custom' is not a valid JSON-Schema type). Exclude 'type' from the passthrough. Fixes #30557. * fix(proxy): stop IAM-refresh engine restart from cascading reconnects (#29176) (#30183) An RDS IAM token refresh recreates the Prisma client, which SIGKILLs the running query-engine and spawns a new one. That planned kill was indistinguishable from a crash, and three reconnect paths used two uncoordinated locks, so a single refresh triggered a cascade of engine kill/respawn cycles: 1. `_safe_refresh_token` (holds `_reconnection_lock`) -> recreate -> kill old engine, spawn new one. 2. The engine-death watcher sees that kill, assumes a crash, and calls `attempt_db_reconnect(force=True)` (a different lock, `_db_reconnect_lock`) -> recreate again -> kills the fresh engine. 3. In-flight queries failing during the swap are classified as transport errors and trigger their own `attempt_db_reconnect` -> recreate again. Fix coordinates planned restarts across the wrapper and the watcher: - PrismaWrapper records the old engine PID in `_expected_engine_deaths` before killing it; all four watcher death-detectors (waitpid thread, pidfd, already-dead probe, os.kill poll) consume that PID and skip the reconnect instead of treating it as a crash. - `recreate_prisma_client` now serializes through `_reconnection_lock` and bumps a monotonic `_engine_generation`. Callers pass `expected_generation` as an optimistic-lock token, so racing/cascading recreates collapse into a single restart (losers no-op). This closes the two-lock gap. - The direct reconnect path probes the writer with SELECT 1 before recreating; a healthy connection (e.g. engine already replaced by a refresh) skips the recreate entirely. - `_safe_refresh_token` coalesces: it skips when the current token still has more than the refresh buffer of runway, so stacked triggers (proactive loop + __getattr__ fallback) don't each restart the engine. An `on_engine_replaced` hook re-arms the watcher on the new PID. RoutingPrismaWrapper forwards `expected_generation` and skips recreating the reader when the writer recreate was skipped. * feat(bedrock): support file content retrieval for batch output files (#30595) Implements transform_file_content_request and transform_file_content_response in BedrockFilesConfig so GET /v1/files/{id}/content works for Bedrock batch files. The request transform resolves the file id (direct s3:// URI or base64 unified id) to its S3 object, validates bucket and key prefix against the server-configured bucket, and SigV4-signs an S3 GetObject using the same credential and region resolution as the existing upload path. The credential and region params are validated into a typed model at the boundary, so the only untyped values left are the botocore signing primitives. Also fixes the proxy managed-files path: CredentialLiteLLMParams now carries s3_bucket_name (previously dropped when building deployment credentials) and the managed-files hook passes the deployment credential snapshot when routing afile_content, so unified-id content retrieval works with per-model bucket config instead of only the AWS_S3_BUCKET_NAME env var. Preserves managed-file access control: the proxy file-content endpoint now rejects raw cloud-storage ids (s3://, gs://), which would otherwise skip the owner/team check that only runs for unified ids and let a caller read another tenant's batch output by its object key. Managed outputs are reachable only through their unified file id. The afile_content "not found" error now reports the caller's unified id rather than the resolved internal S3 URI. Fixes #16186, #15563 * fix(oci): make Cohere {{trace}} judges work (tool param types + agentic tool-calling continuation) (#30646) * fix(oci): map Cohere tool array/object params to lowercase builtins OCI's Cohere backend returns HTTP 500 on a tool parameter typed as a bare "List", which is what OCI_JSON_TO_PYTHON_TYPES produced for JSON-schema arrays. MLflow {{trace}} judges trip this: their tools (get_root_span, get_span) take an attributes_to_fetch array. The lowercase builtins list/dict are accepted; only the bare "List" 500s ("Dict" happens to be tolerated, but both are lowercased for consistency). Verified live against us-chicago-1 (cohere.command-a-03-2025 and command-latest). Adds a unit regression on the transformed parameterDefinitions plus a gated integration test exercising an array-param tool end to end. * fix(oci): make Cohere agentic tool-calling continuation work Two bugs broke the OCI Cohere tool-calling loop that MLflow {{trace}} judges drive once a tool has been executed and its result is fed back. Request side: litellm pulled the last user message into the top-level `message` and emitted the tool result as a TOOL entry in chatHistory. OCI rejects that ("cannot specify message if the last entry in chat history contains tool results"), and an empty message alone is rejected too ("message must be at least 1 token long or tool results must be specified"). OCI carries the current turn's results in a dedicated top-level `toolResults` field. The Cohere transform now sends an empty message, keeps the user turn in chatHistory, and puts the results in `toolResults`, matching the langchain-oracle reference. Tool results are no longer represented as chatHistory entries. Response side: tool-grounded answers come back with citations carrying `documentIds` (camelCase) and no `document_ids`, which made the required `CohereCitation.document_ids` field fail validation and sink the whole response parse. Those citations are never surfaced, so the field (and CohereSearchQuery's generation_id) is now optional. Verified live against us-chicago-1 (cohere.command-a-03-2025 and command-latest), single and multi-round tool loops. Adds unit regressions on the transformed request shape and on citation parsing, plus gated integration tests for the continuation. * feat: integrate Repelloai Argus guardrail (#30673) * feat(guardrails): add RepelloAI Argus guardrail integration (#1) * feat(guardrails): add RepelloAI Argus guardrail integration Add a new guardrail hook backed by RepelloAI Argus, with dashboard-managed asset policies enforced via an asset_id and X-API-Key auth. * fix(guardrails): harden RepelloAI Argus guardrail - scan streaming responses on output (was bypassing the guardrail) - log blocked verdicts as guardrail_intervened instead of success - treat auth/config errors (401/403/404/422) as misconfiguration that always blocks, not a fail-open-able unreachable error - default unreachable_fallback to fail_closed and read it directly; block on unknown/malformed verdicts so an API change can't silently disable enforcement - type unreachable_fallback as a Literal, drop the duplicate config model, expose unreachable_fallback in the config schema, and stop leaking the raw provider response / exception strings to the client * fix(guardrails): address RepelloAI Argus review feedback - support ARGUS_API_KEY (with REPELLOAI_API_KEY fallback) - make asset_id required in the config model - normalize unreachable_fallback so only fail_open opens; block on 400 misconfig - correct the shared unreachable_fallback field description * docs(guardrails): add RepelloAI Argus docs page and dashboard listing - add docs page covering config, env vars, modes, verdicts, failure semantics - list RepelloAI Argus in the Guardrail Garden with provider/logo mappings - add a regression test for the provider logo and display-name resolution * fix(guardrails): keep RepelloAI asset_id optional in config model A required asset_id leaked onto the shared LitellmParams (which inherits RepelloAIGuardrailConfigModel), breaking validation for every other guardrail. Keep it optional like sibling models; the guardrail __init__ still raises when asset_id is missing, which is the real enforcement. * Add comment for last user turn scanning * feat(guardrails): harden repelloai scanning * feat(guardrails): expand repelloai scanning to include tool definitions Add extraction of tool definitions and tool call arguments to the RepelloAI guardrail scanning. Improves detection coverage by including function schemas and parameters in the prompt sent to the guardrail service. Also captures detailed error responses in logs and adds guardrail header to streaming responses. * refactor(guardrails): fix and harden repelloai schema text extraction - Fix duplicate text in _iter_schema_text: previously all dict values were re-queued onto the stack even after scalar/list keys were already extracted explicitly, causing names/descriptions to appear twice in the scanned prompt - Extract schema key frozensets to module-level constants so they are not reconstructed on every call - Change _iter_schema_text from @classmethod to @staticmethod (cls unused) - Narrow _call_analyze stage param from str to Literal["prompt", "response"] - Add HttpxResponse type annotation to _raise_for_config_error - Add LLMResponseTypes annotation to async_post_call_success_hook response param * fix(guardrails): resolve pyright type errors in repelloai guardrail - Narrow async_handler.post return from Response|None to Response with explicit None guard before calling raise_for_status/json - Fix list comprehension returning str|None by switching to explicit loop with isinstance guard so pyright tracks the narrowing - Cast model_dump() result to Dict since hasattr does not narrow object type in pyright * fix(guardrails/repello): include Responses API instructions field in prompt scan The /v1/responses top-level `instructions` field was not included in _extract_prompt_text, allowing a caller to bypass guardrail policy checks by putting blocked content in `instructions` while keeping `input` benign. * feat: add api_key to config model and read prompt from data dict * fix(guardrails/repello): plug input_text and tool-call response bypass gaps Responses API input content parts with type 'input_text' were silently dropped by build_inspection_messages (which only handles type='text'), allowing callers to send blocked content via that path without triggering the pre-call scan. Fix: add _extract_input_text_parts to RepelloAIGuardrail and call it when walking the Responses API input messages. Post-call scanning skipped responses whose choices contained only tool_calls or function_call (message.content=None), letting models put blocked output in function arguments undetected. Fix: _extract_chat_completion_text now calls _extract_tool_call_args_from_message on each choice message. Also replace typing.Dict/List with builtin dict/list to clear TID251 strict ruff violations introduced by this file. * fix(guardrails/repello): scan Responses API function_call output arguments Output items with type 'function_call' in a /v1/responses response were skipped by _extract_responses_api_text; only 'message' items were walked. A model could return blocked content in function_call.arguments undetected. Now extract arguments from function_call output items before scanning. * refactor(guardrails/repello): clean up typing and remove lint-any workarounds - Replace Optional[X]/Union[X,Y] with X|None/X|Y union syntax throughout - Use dict[str, object] instead of bare dict in all signatures - Remove **kwargs from __init__; declare guardrail_name, event_hook, default_on explicitly - Replace getattr(litellm_params, ...) with direct attribute access now that LitellmParams inherits RepelloAIGuardrailConfigModel - Add _event_hook_from_mode() to convert str|list[str]|Mode to typed GuardrailEventHooks - Use TypeAdapter.validate_json() instead of response.json() + manual dict construction - Add _is_object_dict/_is_object_list TypeGuard helpers to narrow object types without Any - Remove cast() workarounds and typed intermediate variables that existed only for the now-removed lint-any CI check - Drop _AddLiteLLMCallback Protocol; budget has sufficient slack for the one reportUnknownMemberType - Fix GuardrailConfigModel missing type arg: GuardrailConfigModel[BaseModel] * fix(guardrails/repello): suppress LIT007 on TypeGuard helpers and add streaming scan-skip warning - Add guard-ok suppressions to _is_object_dict and _is_object_list to satisfy the LIT007 hard-zero budget gate - Emit verbose_proxy_logger.warning when the streaming hook finds no inspectable text after assembly, matching observability of pre/post hooks * refactor: modifications for lint check * feat: add Pinstripes as an OpenAI-compatible provider (#30567) * feat: add Pinstripes as an OpenAI-compatible provider Pinstripes (https://pinstripes.io) is an OpenAI-compatible inference provider serving open-source models (GLM-4.5-Air, Qwen3, DeepSeek, etc.) with per-token pricing and no subscriptions. Changes: - `litellm/llms/openai_like/providers.json`: register pinstripes with base_url, api_key_env, and max_completion_tokens→max_tokens mapping - `litellm/types/utils.py`: add `PINSTRIPES = "pinstripes"` to LlmProviders - `litellm/constants.py`: add to openai_compatible_providers and openai_compatible_endpoints lists - `litellm/litellm_core_utils/get_llm_provider_logic.py`: auto-detect provider when api_base is "https://pinstripes.io/v1" - `provider_endpoints_support.json`: document supported endpoints - `tests/`: 7 unit tests covering provider registration, resolution, URL auto-detection, api_base override, and Router config Usage: import litellm response = litellm.completion( model="pinstripes/ps/glm-4.5-air", messages=[{"role": "user", "content": "Hello"}], api_key=os.environ["PINSTRIPES_API_KEY"], ) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(pinstripes): resolve Greptile P1 review comments - Add api_base_env: PINSTRIPES_API_BASE to providers.json so env var override works - Set responses: false in provider_endpoints_support.json — not actually wired up - Remove docs/my-website/docs/providers/pinstripes.md — belongs in litellm-docs repo Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(pinstripes): add api_base_env and correct responses capability - Add api_base_env: PINSTRIPES_API_BASE to providers.json - Set responses: false in provider_endpoints_support.json Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(pinstripes): wire up Responses API — add supported_endpoints Adds supported_endpoints: ["/v1/chat/completions", "/v1/responses"] so JSONProviderRegistry.supports_responses_api returns true correctly, matching what provider_endpoints_support.json advertises. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * feat(pinstripes): enable embeddings endpoint Pinstripes serves nomic-embed-text-v1.5 and bge-m3 via /v1/embeddings. Add /v1/embeddings to supported_endpoints and set embeddings: true. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(pinstripes): use 4-space indentation in model_prices_and_context_window.json Matches the file's existing convention. Flagged by Greptile review. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(pinstripes): set a2a: false — A2A protocol not implemented All comparable JSON-configured providers (tensormesh, parasail, empiriolabs, libertai, neosantara) have a2a: false. Pinstripes does not implement the Google A2A protocol, so this should be false to match. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> --------- Co-authored-by: inference_provider <max@redactedlab.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(rag): attach existing OpenAI file ids (#30628) * fix(rag): attach existing OpenAI file ids * chore: use modern typing in rag ingest fix * chore: retrigger ci * fix(anthropic-messages): apply cache_control_injection_points on /v1/messages path (#30341) cache_control_injection_points was only consumed by the chat/completions prompt-management hook; on the native Anthropic /v1/messages path it was forwarded unused, so deployment-level cache injection was silently dropped (cache_creation_input_tokens stayed 0 for Anthropic-native clients). Add AnthropicCacheControlHook.apply_to_anthropic_messages_request to inject cache_control at block level for system / tools / message locations (the only forms /v1/messages accepts), wire it into the native anthropic_messages handler, and pop the param so it does not leak upstream as an unknown field. A {location: message, role: system} config is redirected to the top-level system prompt so the same YAML works on both endpoints. Injection respects Anthropic's 4-block cache_control limit shared across system, tools, and messages: client-supplied markers count toward the cap and are never overwritten, a slot is reserved per Bedrock tool_config point, and injection stops once the budget is exhausted. Locations this path cannot represent (tool_config) are forwarded downstream instead of being silently consumed, mirroring get_chat_completion_prompt's remaining_points pass-through. Built on litellm_internal_staging. Refs #30293 * fix(proxy): release budget reservation when a request is cancelled mid-flight (#30522) * fix(proxy): release budget reservation on cancel when no chunk was delivered The pre-call budget reservation increments the cross-pod spend counter by a request's worst-case cost, then reconciles it on success (cost callback) or error (failure hook). A client disconnect or timeout cancels the request and surfaces as CancelledError / GeneratorExit, which neither path catches, so the reservation leaks. Under a retry storm the leaked holds accumulate, pin the counter above real spend, and return spurious 429 "Budget has been exceeded" to keys whose spend is far below budget; the counter only recovers when its TTL lapses, so the failure is intermittent and self-healing. Release the reservation in async_streaming_data_generator (which the Anthropic and Google SSE generators delegate to) on the (CancelledError, GeneratorExit) path, alongside the existing max_parallel_requests release. release_budget_ reservation_on_cancel runs under asyncio.shield so it completes despite the in-progress cancellation, is guarded by the reservation's finalized flag, and swallows a failing release so it cannot replace the in-flight cancellation. The refund is gated on whether a chunk reached the client. The flag is set immediately before the yield, after the slow-path hook await: an async generator suspends at the yield, so a GeneratorExit on disconnect after a delivered chunk sees it True (keep the hold), while a cancellation during the slow-path await leaves it False (refund, nothing sent). A non-streaming cancellation delivers nothing and a completed non-streaming response is reconciled by the success callback, so neither needs a release here. Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix(proxy): reconcile a cancelled reservation to input cost, not zero A streaming request cancelled before the first chunk previously reconciled its reservation to zero and finalized it. But by the time the generator is consuming the response the provider call was already dispatched, so the input tokens were billed even though no chunk reached the client, and the success/failure cost callbacks are skipped on cancellation. Refunding to zero let a caller send an expensive request and abort pre-token to dodge the input charge. Compute the request's input-token cost at reservation time and reconcile the cancelled reservation to it instead of zero. The worst-case output portion of the reservation is still released (so a legitimate mid-flight cancellation no longer pins the counter and 429s the key), while the input the provider already processed is charged. --------- Co-authored-by: Bytechoreographer <Bytechoreographer@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix(caching): encode object name in GCS cache GET path (#30378) GCS cache reads always missed when gcs_path was set. The GET methods interpolated the object name directly into the URL path, while the GCS JSON API requires it to be URL-encoded (a "/" must be sent as %2F). With gcs_path configured the object name is "<prefix>/<sha256>", so the raw slash produced a malformed object path and GCS returned 404. httpx does not raise on 4xx, so the status_code == 200 check fell through and get/async_get returned None, silently missing on every read. Without gcs_path the key has no slash, which is why this went unnoticed. Wrap the object name with urllib.parse.quote(..., safe="") in get_cache and async_get_cache. Apply the same encoding to the name= query parameter in set_cache and async_set_cache so the key written matches the key read back. Adds regression tests asserting the GET path and SET query are encoded (%2F) when gcs_path is set, for both sync and async paths; these fail on the unpatched code. Fixes #30377 * chore: add soniox stt-async-v5 model (#30672) * fix(proxy): include model group aliases in v1 model info (#30626) * Include model group aliases in v1 model info * Fix model info alias implementation * removed extra blank line * chore: rerun CI * fix(lint): remove redundant noqa directive in proxy_cli.py * fix: address greptile review - restore bedrock_mantle auth symbols, guard OCI empty message list, validate DIRECT_URL scheme * Revert "fix: address greptile review - restore bedrock_mantle auth symbols, guard OCI empty message list, validate DIRECT_URL scheme" This reverts commit 52c7a07. * Revert "fix(anthropic-messages): apply cache_control_injection_points on /v1/messages path (#30341)" This reverts commit c9e8a17. * Revert "fix(proxy): stop IAM-refresh engine restart from cascading reconnects (#29176) (#30183)" This reverts commit 85828da. * fix(proxy): stop IAM-refresh engine restart from cascading reconnects (#29176) (#30183) An RDS IAM token refresh recreates the Prisma client, which SIGKILLs the running query-engine and spawns a new one. That planned kill was indistinguishable from a crash, and three reconnect paths used two uncoordinated locks, so a single refresh triggered a cascade of engine kill/respawn cycles: 1. `_safe_refresh_token` (holds `_reconnection_lock`) -> recreate -> kill old engine, spawn new one. 2. The engine-death watcher sees that kill, assumes a crash, and calls `attempt_db_reconnect(force=True)` (a different lock, `_db_reconnect_lock`) -> recreate again -> kills the fresh engine. 3. In-flight queries failing during the swap are classified as transport errors and trigger their own `attempt_db_reconnect` -> recreate again. Fix coordinates planned restarts across the wrapper and the watcher: - PrismaWrapper records the old engine PID in `_expected_engine_deaths` before killing it; all four watcher death-detectors (waitpid thread, pidfd, already-dead probe, os.kill poll) consume that PID and skip the reconnect instead of treating it as a crash. - `recreate_prisma_client` now serializes through `_reconnection_lock` and bumps a monotonic `_engine_generation`. Callers pass `expected_generation` as an optimistic-lock token, so racing/cascading recreates collapse into a single restart (losers no-op). This closes the two-lock gap. - The direct reconnect path probes the writer with SELECT 1 before recreating; a healthy connection (e.g. engine already replaced by a refresh) skips the recreate entirely. - `_safe_refresh_token` coalesces: it skips when the current token still has more than the refresh buffer of runway, so stacked triggers (proactive loop + __getattr__ fallback) don't each restart the engine. An `on_engine_replaced` hook re-arms the watcher on the new PID. RoutingPrismaWrapper forwards `expected_generation` and skips recreating the reader when the writer recreate was skipped. * fix(lint): modernize type annotations in IAM-refresh prisma client files (UP006/UP045) * Revert "feat(proxy): show session-aggregate cost and duration in request logs (#25708) (#30507)" This reverts commit f530b22. * Revert "fix(dashscope): treat an explicit 0.0 tier cost as a real price, not missing (#30653)" This reverts commit 4f58bd0. * Revert "fix(oci): make Cohere {{trace}} judges work (tool param types + agentic tool-calling continuation) (#30646)" This reverts commit 50f34e0. * Revert "fix(proxy): fail fast on non-PostgreSQL DATABASE_URL instead of hanging on startup (#30366)" This reverts commit 0544eed. * fix(bedrock_mantle): restore BedrockMantleAuthMixin and constants removed by routing rewrite * fix(key management): restore exact /key/list user_id & key_alias matching by default (#30593) Before substring search was added (commit 33bd570), /key/list matched user_id and key_alias exactly. That change made admin-authenticated calls substring-match by default, breaking the prior contract: a caller passing an exact user_id as an access filter (e.g. an integration scoping to one user with an admin key) then received other users' keys -- user_id="alice" also returned "alice2", "alice-test", etc. This is a cross-user key disclosure. Make substring matching opt-in via a new admin-only substring_matching=true query param; default to exact, restoring the prior behavior. The dashboard search box (keyListCall) passes the flag so partial search still works. Non-admins remain exact and scoped to their own keys. Updates the proxy-behavior key_alias test to opt in and adds an exact-by-default guard; adds list_keys unit coverage for the opt-in gate. --------- Co-authored-by: perseus <51974392+tcconnally@users.noreply.github.com> Co-authored-by: Hannah Smith <64043506+hannahmadison@users.noreply.github.com> Co-authored-by: Charlie Patterson <Pattersoncharlesl@gmail.com> Co-authored-by: Matthew Lapointe <mlapointe@alpha-sense.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: KRISH SONI <67964054+krishvsoni@users.noreply.github.com> Co-authored-by: Yash Raj Pandey <55940078+devYRPauli@users.noreply.github.com> Co-authored-by: Nitish Agarwal <1592163+nitishagar@users.noreply.github.com> Co-authored-by: hcl <chenglunhu@gmail.com> Co-authored-by: tushar8408 <32977767+tushar8408@users.noreply.github.com> Co-authored-by: AD Mohanraj <admohanraj@gmail.com> Co-authored-by: Fede Kamelhar <federico.kamelhar@oracle.com> Co-authored-by: Lavish Bansal <lavish.bansal619@gmail.com> Co-authored-by: max-amos <gruffulom@gmail.com> Co-authored-by: inference_provider <max@redactedlab.com> Co-authored-by: NK <93352237+Nithish-Yenaganti@users.noreply.github.com> Co-authored-by: 安妮的心动录 <74543653+anneheartrecord@users.noreply.github.com> Co-authored-by: Rick <26716961+Bytechoreographer@users.noreply.github.com> Co-authored-by: Bytechoreographer <Bytechoreographer@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Burak Ömür <burak.omur.1998@gmail.com> Co-authored-by: Dan Lemon <daniel.lemon@amazee.io> Co-authored-by: Vanika Dangi <166420943+vanika02@users.noreply.github.com> Co-authored-by: Jay Gowdy <130084966+jgowdy-godaddy@users.noreply.github.com>
* fix(proxy): bump health-check max_tokens default to 16 for GPT-5 compatibility (BerriAI#30708) OpenAI GPT-5 models require max_completion_tokens >= 16. Health checks were using 5 (proxy/health_check.py) and 10 (health_check_helpers.py), causing failures on GPT-5 models. Fixes BerriAI#23836 * fix: increase health check max_tokens from 5 to 16 (BerriAI#23836) (BerriAI#26610) GPT-5 models enforce a minimum of 16 for max_output_tokens. The current default of 5 still causes health checks to fail for these models. Bump the non-wildcard default to 16 — the smallest value that satisfies all known provider minimums while keeping health checks lightweight. Also tightens the wildcard test assertion from a weak disjunctive check to strict key-absence. Co-authored-by: Sameer Kankute <sameer@berri.ai> * fix: ensure checks show gemini-3-flash-preview supports responseJsonS… (BerriAI#30696) * fix: ensure checks show gemini-3-flash-preview supports responseJsonSchema. * fix: remove async keyword from test. * fix: make Bedrock Mantle Responses routing data-driven per model (BerriAI#30700) * Make Bedrock Mantle Responses routing data-driven per model Route Bedrock Mantle models to the native Responses API based on each model's price-map capability signal instead of a hardcoded model-name heuristic, and derive the OpenAI-compatible base path segment per model. Responses dispatch now selects the native config when the model advertises responses support (/v1/responses in supported_endpoints, or mode=responses), both overridable via register_model and proxy model_info. This enables native Responses for gpt-oss-120b/20b and the gemma-4 family while keeping chat-only models (gpt-oss safeguard, nvidia, mistral, ...) on the existing chat-completions emulation. Capability is per-model, so gpt-oss-120b routes natively while gpt-oss-safeguard-120b does not despite sharing the gpt-oss substring. The wire path is a separate concern, driven by the existing use_openai_responses_path flag rather than a model-name match: gpt-5.x and gemma-4-* on /openai/v1, everything else (incl. gpt-oss) on /v1. The chat config now derives its base from the same flag, fixing gemma-4 chat-completions requests that previously went to /v1 instead of /openai/v1. Cost maps: add supported_endpoints to the gpt-oss entries (responses for the non-safeguard variants, chat-only for safeguard) and supported_endpoints + use_openai_responses_path to all three gemma-4 entries. * Address review: move capability helper into bedrock_mantle package Move the Responses capability check out of utils.py into litellm/llms/bedrock_mantle/common_utils.py as mantle_supports_responses, alongside its companion wire-path helper mantle_base_segment. Both are now pure functions of (model, model_cost): the price-map mode/supported_endpoints read replaces the get_model_info call, so the rules are unit-testable without patching global state and the Bedrock Mantle package is self-contained. Use str | None instead of Optional[str] on the new signatures to satisfy the ruff UP045 strict-rule gate. Add direct unit tests for both helpers. Fix test_register_model_restore_undoes_existing_key_overwrite: gpt-oss-120b now legitimately supports Responses, so it can no longer be the "None after restore" vehicle; use the chat-only safeguard variant, which isolates the register/restore effect from the model's own capability. --------- Co-authored-by: Sameer Kankute <sameer@berri.ai> * fix(proxy): fail fast on non-PostgreSQL DATABASE_URL instead of hanging on startup (BerriAI#30366) * fix(proxy): fail fast on non-PostgreSQL DATABASE_URL instead of hanging on startup LiteLLM's Prisma datasource is pinned to provider = 'postgresql', so a sqlite:// or mysql:// DATABASE_URL can never connect. Today that surfaces as an opaque startup stall where the port never binds, and a separate 'DB not connected' 500 on /key/generate when no DATABASE_URL is set at all leaves operators guessing what to configure. Validate the DATABASE_URL / DIRECT_URL scheme in run_server before any Prisma call and exit with an actionable message naming the unsupported scheme. Also reword CommonProxyErrors.db_not_connected_error to tell the operator to set DATABASE_URL to a postgresql:// connection string. Add regression tests covering postgres acceptance and sqlite/mysql/mssql rejection. * fix: resolve CI failures and proxy DB URL typing issue * fix(dashscope): treat an explicit 0.0 tier cost as a real price, not missing (BerriAI#30653) The tiered cost calculator resolved a tier's per-token cost with `tier.get(cost_key) or tier.get(fallback_cost_key, 0)`. Because `or` short-circuits on any falsy value, a tier that legitimately prices a component at 0.0 (e.g. a free-cache-read tier with cache_read_input_token_cost: 0.0, or a free-reasoning tier) is treated as missing and silently billed at the full fallback rate (input_cost_per_token / output_cost_per_token). The flat-pricing path in the same module already handles this correctly with an `is None` guard. Resolve tier costs through a small helper that mirrors it, so 0.0 is honored at both the in-range and overflow sites. No shipped model currently has a 0.0 tier cost, so this is a latent defect; the fix makes the tiered path consistent with the flat path and prevents over-charging the first time such a tier appears. Adds unit tests covering the in-range and overflow paths, and drops an unused import flagged by ruff in the touched test file. * feat(proxy): show session-aggregate cost and duration in request logs (BerriAI#25708) (BerriAI#30507) * fix(anthropic): don't leak tool 'type' into OpenAI function parameters schema (BerriAI#30618) In the messages->chat/completions bridge, translate_anthropic_tools_to_openai merged every non-mapped tool key into the function parameters dict. The Anthropic tool 'type' (e.g. 'custom') thus overwrote parameters.type ('object' -> 'custom'), and providers reject it ('custom' is not a valid JSON-Schema type). Exclude 'type' from the passthrough. Fixes BerriAI#30557. * fix(proxy): stop IAM-refresh engine restart from cascading reconnects (BerriAI#29176) (BerriAI#30183) An RDS IAM token refresh recreates the Prisma client, which SIGKILLs the running query-engine and spawns a new one. That planned kill was indistinguishable from a crash, and three reconnect paths used two uncoordinated locks, so a single refresh triggered a cascade of engine kill/respawn cycles: 1. `_safe_refresh_token` (holds `_reconnection_lock`) -> recreate -> kill old engine, spawn new one. 2. The engine-death watcher sees that kill, assumes a crash, and calls `attempt_db_reconnect(force=True)` (a different lock, `_db_reconnect_lock`) -> recreate again -> kills the fresh engine. 3. In-flight queries failing during the swap are classified as transport errors and trigger their own `attempt_db_reconnect` -> recreate again. Fix coordinates planned restarts across the wrapper and the watcher: - PrismaWrapper records the old engine PID in `_expected_engine_deaths` before killing it; all four watcher death-detectors (waitpid thread, pidfd, already-dead probe, os.kill poll) consume that PID and skip the reconnect instead of treating it as a crash. - `recreate_prisma_client` now serializes through `_reconnection_lock` and bumps a monotonic `_engine_generation`. Callers pass `expected_generation` as an optimistic-lock token, so racing/cascading recreates collapse into a single restart (losers no-op). This closes the two-lock gap. - The direct reconnect path probes the writer with SELECT 1 before recreating; a healthy connection (e.g. engine already replaced by a refresh) skips the recreate entirely. - `_safe_refresh_token` coalesces: it skips when the current token still has more than the refresh buffer of runway, so stacked triggers (proactive loop + __getattr__ fallback) don't each restart the engine. An `on_engine_replaced` hook re-arms the watcher on the new PID. RoutingPrismaWrapper forwards `expected_generation` and skips recreating the reader when the writer recreate was skipped. * feat(bedrock): support file content retrieval for batch output files (BerriAI#30595) Implements transform_file_content_request and transform_file_content_response in BedrockFilesConfig so GET /v1/files/{id}/content works for Bedrock batch files. The request transform resolves the file id (direct s3:// URI or base64 unified id) to its S3 object, validates bucket and key prefix against the server-configured bucket, and SigV4-signs an S3 GetObject using the same credential and region resolution as the existing upload path. The credential and region params are validated into a typed model at the boundary, so the only untyped values left are the botocore signing primitives. Also fixes the proxy managed-files path: CredentialLiteLLMParams now carries s3_bucket_name (previously dropped when building deployment credentials) and the managed-files hook passes the deployment credential snapshot when routing afile_content, so unified-id content retrieval works with per-model bucket config instead of only the AWS_S3_BUCKET_NAME env var. Preserves managed-file access control: the proxy file-content endpoint now rejects raw cloud-storage ids (s3://, gs://), which would otherwise skip the owner/team check that only runs for unified ids and let a caller read another tenant's batch output by its object key. Managed outputs are reachable only through their unified file id. The afile_content "not found" error now reports the caller's unified id rather than the resolved internal S3 URI. Fixes BerriAI#16186, BerriAI#15563 * fix(oci): make Cohere {{trace}} judges work (tool param types + agentic tool-calling continuation) (BerriAI#30646) * fix(oci): map Cohere tool array/object params to lowercase builtins OCI's Cohere backend returns HTTP 500 on a tool parameter typed as a bare "List", which is what OCI_JSON_TO_PYTHON_TYPES produced for JSON-schema arrays. MLflow {{trace}} judges trip this: their tools (get_root_span, get_span) take an attributes_to_fetch array. The lowercase builtins list/dict are accepted; only the bare "List" 500s ("Dict" happens to be tolerated, but both are lowercased for consistency). Verified live against us-chicago-1 (cohere.command-a-03-2025 and command-latest). Adds a unit regression on the transformed parameterDefinitions plus a gated integration test exercising an array-param tool end to end. * fix(oci): make Cohere agentic tool-calling continuation work Two bugs broke the OCI Cohere tool-calling loop that MLflow {{trace}} judges drive once a tool has been executed and its result is fed back. Request side: litellm pulled the last user message into the top-level `message` and emitted the tool result as a TOOL entry in chatHistory. OCI rejects that ("cannot specify message if the last entry in chat history contains tool results"), and an empty message alone is rejected too ("message must be at least 1 token long or tool results must be specified"). OCI carries the current turn's results in a dedicated top-level `toolResults` field. The Cohere transform now sends an empty message, keeps the user turn in chatHistory, and puts the results in `toolResults`, matching the langchain-oracle reference. Tool results are no longer represented as chatHistory entries. Response side: tool-grounded answers come back with citations carrying `documentIds` (camelCase) and no `document_ids`, which made the required `CohereCitation.document_ids` field fail validation and sink the whole response parse. Those citations are never surfaced, so the field (and CohereSearchQuery's generation_id) is now optional. Verified live against us-chicago-1 (cohere.command-a-03-2025 and command-latest), single and multi-round tool loops. Adds unit regressions on the transformed request shape and on citation parsing, plus gated integration tests for the continuation. * feat: integrate Repelloai Argus guardrail (BerriAI#30673) * feat(guardrails): add RepelloAI Argus guardrail integration (#1) * feat(guardrails): add RepelloAI Argus guardrail integration Add a new guardrail hook backed by RepelloAI Argus, with dashboard-managed asset policies enforced via an asset_id and X-API-Key auth. * fix(guardrails): harden RepelloAI Argus guardrail - scan streaming responses on output (was bypassing the guardrail) - log blocked verdicts as guardrail_intervened instead of success - treat auth/config errors (401/403/404/422) as misconfiguration that always blocks, not a fail-open-able unreachable error - default unreachable_fallback to fail_closed and read it directly; block on unknown/malformed verdicts so an API change can't silently disable enforcement - type unreachable_fallback as a Literal, drop the duplicate config model, expose unreachable_fallback in the config schema, and stop leaking the raw provider response / exception strings to the client * fix(guardrails): address RepelloAI Argus review feedback - support ARGUS_API_KEY (with REPELLOAI_API_KEY fallback) - make asset_id required in the config model - normalize unreachable_fallback so only fail_open opens; block on 400 misconfig - correct the shared unreachable_fallback field description * docs(guardrails): add RepelloAI Argus docs page and dashboard listing - add docs page covering config, env vars, modes, verdicts, failure semantics - list RepelloAI Argus in the Guardrail Garden with provider/logo mappings - add a regression test for the provider logo and display-name resolution * fix(guardrails): keep RepelloAI asset_id optional in config model A required asset_id leaked onto the shared LitellmParams (which inherits RepelloAIGuardrailConfigModel), breaking validation for every other guardrail. Keep it optional like sibling models; the guardrail __init__ still raises when asset_id is missing, which is the real enforcement. * Add comment for last user turn scanning * feat(guardrails): harden repelloai scanning * feat(guardrails): expand repelloai scanning to include tool definitions Add extraction of tool definitions and tool call arguments to the RepelloAI guardrail scanning. Improves detection coverage by including function schemas and parameters in the prompt sent to the guardrail service. Also captures detailed error responses in logs and adds guardrail header to streaming responses. * refactor(guardrails): fix and harden repelloai schema text extraction - Fix duplicate text in _iter_schema_text: previously all dict values were re-queued onto the stack even after scalar/list keys were already extracted explicitly, causing names/descriptions to appear twice in the scanned prompt - Extract schema key frozensets to module-level constants so they are not reconstructed on every call - Change _iter_schema_text from @classmethod to @staticmethod (cls unused) - Narrow _call_analyze stage param from str to Literal["prompt", "response"] - Add HttpxResponse type annotation to _raise_for_config_error - Add LLMResponseTypes annotation to async_post_call_success_hook response param * fix(guardrails): resolve pyright type errors in repelloai guardrail - Narrow async_handler.post return from Response|None to Response with explicit None guard before calling raise_for_status/json - Fix list comprehension returning str|None by switching to explicit loop with isinstance guard so pyright tracks the narrowing - Cast model_dump() result to Dict since hasattr does not narrow object type in pyright * fix(guardrails/repello): include Responses API instructions field in prompt scan The /v1/responses top-level `instructions` field was not included in _extract_prompt_text, allowing a caller to bypass guardrail policy checks by putting blocked content in `instructions` while keeping `input` benign. * feat: add api_key to config model and read prompt from data dict * fix(guardrails/repello): plug input_text and tool-call response bypass gaps Responses API input content parts with type 'input_text' were silently dropped by build_inspection_messages (which only handles type='text'), allowing callers to send blocked content via that path without triggering the pre-call scan. Fix: add _extract_input_text_parts to RepelloAIGuardrail and call it when walking the Responses API input messages. Post-call scanning skipped responses whose choices contained only tool_calls or function_call (message.content=None), letting models put blocked output in function arguments undetected. Fix: _extract_chat_completion_text now calls _extract_tool_call_args_from_message on each choice message. Also replace typing.Dict/List with builtin dict/list to clear TID251 strict ruff violations introduced by this file. * fix(guardrails/repello): scan Responses API function_call output arguments Output items with type 'function_call' in a /v1/responses response were skipped by _extract_responses_api_text; only 'message' items were walked. A model could return blocked content in function_call.arguments undetected. Now extract arguments from function_call output items before scanning. * refactor(guardrails/repello): clean up typing and remove lint-any workarounds - Replace Optional[X]/Union[X,Y] with X|None/X|Y union syntax throughout - Use dict[str, object] instead of bare dict in all signatures - Remove **kwargs from __init__; declare guardrail_name, event_hook, default_on explicitly - Replace getattr(litellm_params, ...) with direct attribute access now that LitellmParams inherits RepelloAIGuardrailConfigModel - Add _event_hook_from_mode() to convert str|list[str]|Mode to typed GuardrailEventHooks - Use TypeAdapter.validate_json() instead of response.json() + manual dict construction - Add _is_object_dict/_is_object_list TypeGuard helpers to narrow object types without Any - Remove cast() workarounds and typed intermediate variables that existed only for the now-removed lint-any CI check - Drop _AddLiteLLMCallback Protocol; budget has sufficient slack for the one reportUnknownMemberType - Fix GuardrailConfigModel missing type arg: GuardrailConfigModel[BaseModel] * fix(guardrails/repello): suppress LIT007 on TypeGuard helpers and add streaming scan-skip warning - Add guard-ok suppressions to _is_object_dict and _is_object_list to satisfy the LIT007 hard-zero budget gate - Emit verbose_proxy_logger.warning when the streaming hook finds no inspectable text after assembly, matching observability of pre/post hooks * refactor: modifications for lint check * feat: add Pinstripes as an OpenAI-compatible provider (BerriAI#30567) * feat: add Pinstripes as an OpenAI-compatible provider Pinstripes (https://pinstripes.io) is an OpenAI-compatible inference provider serving open-source models (GLM-4.5-Air, Qwen3, DeepSeek, etc.) with per-token pricing and no subscriptions. Changes: - `litellm/llms/openai_like/providers.json`: register pinstripes with base_url, api_key_env, and max_completion_tokens→max_tokens mapping - `litellm/types/utils.py`: add `PINSTRIPES = "pinstripes"` to LlmProviders - `litellm/constants.py`: add to openai_compatible_providers and openai_compatible_endpoints lists - `litellm/litellm_core_utils/get_llm_provider_logic.py`: auto-detect provider when api_base is "https://pinstripes.io/v1" - `provider_endpoints_support.json`: document supported endpoints - `tests/`: 7 unit tests covering provider registration, resolution, URL auto-detection, api_base override, and Router config Usage: import litellm response = litellm.completion( model="pinstripes/ps/glm-4.5-air", messages=[{"role": "user", "content": "Hello"}], api_key=os.environ["PINSTRIPES_API_KEY"], ) * fix(pinstripes): resolve Greptile P1 review comments - Add api_base_env: PINSTRIPES_API_BASE to providers.json so env var override works - Set responses: false in provider_endpoints_support.json — not actually wired up - Remove docs/my-website/docs/providers/pinstripes.md — belongs in litellm-docs repo * fix(pinstripes): add api_base_env and correct responses capability - Add api_base_env: PINSTRIPES_API_BASE to providers.json - Set responses: false in provider_endpoints_support.json * fix(pinstripes): wire up Responses API — add supported_endpoints Adds supported_endpoints: ["/v1/chat/completions", "/v1/responses"] so JSONProviderRegistry.supports_responses_api returns true correctly, matching what provider_endpoints_support.json advertises. * feat(pinstripes): enable embeddings endpoint Pinstripes serves nomic-embed-text-v1.5 and bge-m3 via /v1/embeddings. Add /v1/embeddings to supported_endpoints and set embeddings: true. * fix(pinstripes): use 4-space indentation in model_prices_and_context_window.json Matches the file's existing convention. Flagged by Greptile review. * fix(pinstripes): set a2a: false — A2A protocol not implemented All comparable JSON-configured providers (tensormesh, parasail, empiriolabs, libertai, neosantara) have a2a: false. Pinstripes does not implement the Google A2A protocol, so this should be false to match. --------- Co-authored-by: inference_provider <max@redactedlab.com> * fix(rag): attach existing OpenAI file ids (BerriAI#30628) * fix(rag): attach existing OpenAI file ids * chore: use modern typing in rag ingest fix * chore: retrigger ci * fix(anthropic-messages): apply cache_control_injection_points on /v1/messages path (BerriAI#30341) cache_control_injection_points was only consumed by the chat/completions prompt-management hook; on the native Anthropic /v1/messages path it was forwarded unused, so deployment-level cache injection was silently dropped (cache_creation_input_tokens stayed 0 for Anthropic-native clients). Add AnthropicCacheControlHook.apply_to_anthropic_messages_request to inject cache_control at block level for system / tools / message locations (the only forms /v1/messages accepts), wire it into the native anthropic_messages handler, and pop the param so it does not leak upstream as an unknown field. A {location: message, role: system} config is redirected to the top-level system prompt so the same YAML works on both endpoints. Injection respects Anthropic's 4-block cache_control limit shared across system, tools, and messages: client-supplied markers count toward the cap and are never overwritten, a slot is reserved per Bedrock tool_config point, and injection stops once the budget is exhausted. Locations this path cannot represent (tool_config) are forwarded downstream instead of being silently consumed, mirroring get_chat_completion_prompt's remaining_points pass-through. Built on litellm_internal_staging. Refs BerriAI#30293 * fix(proxy): release budget reservation when a request is cancelled mid-flight (BerriAI#30522) * fix(proxy): release budget reservation on cancel when no chunk was delivered The pre-call budget reservation increments the cross-pod spend counter by a request's worst-case cost, then reconciles it on success (cost callback) or error (failure hook). A client disconnect or timeout cancels the request and surfaces as CancelledError / GeneratorExit, which neither path catches, so the reservation leaks. Under a retry storm the leaked holds accumulate, pin the counter above real spend, and return spurious 429 "Budget has been exceeded" to keys whose spend is far below budget; the counter only recovers when its TTL lapses, so the failure is intermittent and self-healing. Release the reservation in async_streaming_data_generator (which the Anthropic and Google SSE generators delegate to) on the (CancelledError, GeneratorExit) path, alongside the existing max_parallel_requests release. release_budget_ reservation_on_cancel runs under asyncio.shield so it completes despite the in-progress cancellation, is guarded by the reservation's finalized flag, and swallows a failing release so it cannot replace the in-flight cancellation. The refund is gated on whether a chunk reached the client. The flag is set immediately before the yield, after the slow-path hook await: an async generator suspends at the yield, so a GeneratorExit on disconnect after a delivered chunk sees it True (keep the hold), while a cancellation during the slow-path await leaves it False (refund, nothing sent). A non-streaming cancellation delivers nothing and a completed non-streaming response is reconciled by the success callback, so neither needs a release here. Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix(proxy): reconcile a cancelled reservation to input cost, not zero A streaming request cancelled before the first chunk previously reconciled its reservation to zero and finalized it. But by the time the generator is consuming the response the provider call was already dispatched, so the input tokens were billed even though no chunk reached the client, and the success/failure cost callbacks are skipped on cancellation. Refunding to zero let a caller send an expensive request and abort pre-token to dodge the input charge. Compute the request's input-token cost at reservation time and reconcile the cancelled reservation to it instead of zero. The worst-case output portion of the reservation is still released (so a legitimate mid-flight cancellation no longer pins the counter and 429s the key), while the input the provider already processed is charged. --------- Co-authored-by: Bytechoreographer <Bytechoreographer@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix(caching): encode object name in GCS cache GET path (BerriAI#30378) GCS cache reads always missed when gcs_path was set. The GET methods interpolated the object name directly into the URL path, while the GCS JSON API requires it to be URL-encoded (a "/" must be sent as %2F). With gcs_path configured the object name is "<prefix>/<sha256>", so the raw slash produced a malformed object path and GCS returned 404. httpx does not raise on 4xx, so the status_code == 200 check fell through and get/async_get returned None, silently missing on every read. Without gcs_path the key has no slash, which is why this went unnoticed. Wrap the object name with urllib.parse.quote(..., safe="") in get_cache and async_get_cache. Apply the same encoding to the name= query parameter in set_cache and async_set_cache so the key written matches the key read back. Adds regression tests asserting the GET path and SET query are encoded (%2F) when gcs_path is set, for both sync and async paths; these fail on the unpatched code. Fixes BerriAI#30377 * chore: add soniox stt-async-v5 model (BerriAI#30672) * fix(proxy): include model group aliases in v1 model info (BerriAI#30626) * Include model group aliases in v1 model info * Fix model info alias implementation * removed extra blank line * chore: rerun CI * fix(lint): remove redundant noqa directive in proxy_cli.py * fix: address greptile review - restore bedrock_mantle auth symbols, guard OCI empty message list, validate DIRECT_URL scheme * Revert "fix: address greptile review - restore bedrock_mantle auth symbols, guard OCI empty message list, validate DIRECT_URL scheme" This reverts commit 52c7a07. * Revert "fix(anthropic-messages): apply cache_control_injection_points on /v1/messages path (BerriAI#30341)" This reverts commit c9e8a17. * Revert "fix(proxy): stop IAM-refresh engine restart from cascading reconnects (BerriAI#29176) (BerriAI#30183)" This reverts commit 85828da. * fix(proxy): stop IAM-refresh engine restart from cascading reconnects (BerriAI#29176) (BerriAI#30183) An RDS IAM token refresh recreates the Prisma client, which SIGKILLs the running query-engine and spawns a new one. That planned kill was indistinguishable from a crash, and three reconnect paths used two uncoordinated locks, so a single refresh triggered a cascade of engine kill/respawn cycles: 1. `_safe_refresh_token` (holds `_reconnection_lock`) -> recreate -> kill old engine, spawn new one. 2. The engine-death watcher sees that kill, assumes a crash, and calls `attempt_db_reconnect(force=True)` (a different lock, `_db_reconnect_lock`) -> recreate again -> kills the fresh engine. 3. In-flight queries failing during the swap are classified as transport errors and trigger their own `attempt_db_reconnect` -> recreate again. Fix coordinates planned restarts across the wrapper and the watcher: - PrismaWrapper records the old engine PID in `_expected_engine_deaths` before killing it; all four watcher death-detectors (waitpid thread, pidfd, already-dead probe, os.kill poll) consume that PID and skip the reconnect instead of treating it as a crash. - `recreate_prisma_client` now serializes through `_reconnection_lock` and bumps a monotonic `_engine_generation`. Callers pass `expected_generation` as an optimistic-lock token, so racing/cascading recreates collapse into a single restart (losers no-op). This closes the two-lock gap. - The direct reconnect path probes the writer with SELECT 1 before recreating; a healthy connection (e.g. engine already replaced by a refresh) skips the recreate entirely. - `_safe_refresh_token` coalesces: it skips when the current token still has more than the refresh buffer of runway, so stacked triggers (proactive loop + __getattr__ fallback) don't each restart the engine. An `on_engine_replaced` hook re-arms the watcher on the new PID. RoutingPrismaWrapper forwards `expected_generation` and skips recreating the reader when the writer recreate was skipped. * fix(lint): modernize type annotations in IAM-refresh prisma client files (UP006/UP045) * Revert "feat(proxy): show session-aggregate cost and duration in request logs (BerriAI#25708) (BerriAI#30507)" This reverts commit f530b22. * Revert "fix(dashscope): treat an explicit 0.0 tier cost as a real price, not missing (BerriAI#30653)" This reverts commit 4f58bd0. * Revert "fix(oci): make Cohere {{trace}} judges work (tool param types + agentic tool-calling continuation) (BerriAI#30646)" This reverts commit 50f34e0. * Revert "fix(proxy): fail fast on non-PostgreSQL DATABASE_URL instead of hanging on startup (BerriAI#30366)" This reverts commit 0544eed. * fix(bedrock_mantle): restore BedrockMantleAuthMixin and constants removed by routing rewrite * fix(key management): restore exact /key/list user_id & key_alias matching by default (BerriAI#30593) Before substring search was added (commit 33bd570), /key/list matched user_id and key_alias exactly. That change made admin-authenticated calls substring-match by default, breaking the prior contract: a caller passing an exact user_id as an access filter (e.g. an integration scoping to one user with an admin key) then received other users' keys -- user_id="alice" also returned "alice2", "alice-test", etc. This is a cross-user key disclosure. Make substring matching opt-in via a new admin-only substring_matching=true query param; default to exact, restoring the prior behavior. The dashboard search box (keyListCall) passes the flag so partial search still works. Non-admins remain exact and scoped to their own keys. Updates the proxy-behavior key_alias test to opt in and adds an exact-by-default guard; adds list_keys unit coverage for the opt-in gate. --------- Co-authored-by: perseus <51974392+tcconnally@users.noreply.github.com> Co-authored-by: Hannah Smith <64043506+hannahmadison@users.noreply.github.com> Co-authored-by: Charlie Patterson <Pattersoncharlesl@gmail.com> Co-authored-by: Matthew Lapointe <mlapointe@alpha-sense.com> Co-authored-by: KRISH SONI <67964054+krishvsoni@users.noreply.github.com> Co-authored-by: Yash Raj Pandey <55940078+devYRPauli@users.noreply.github.com> Co-authored-by: Nitish Agarwal <1592163+nitishagar@users.noreply.github.com> Co-authored-by: hcl <chenglunhu@gmail.com> Co-authored-by: tushar8408 <32977767+tushar8408@users.noreply.github.com> Co-authored-by: AD Mohanraj <admohanraj@gmail.com> Co-authored-by: Fede Kamelhar <federico.kamelhar@oracle.com> Co-authored-by: Lavish Bansal <lavish.bansal619@gmail.com> Co-authored-by: max-amos <gruffulom@gmail.com> Co-authored-by: inference_provider <max@redactedlab.com> Co-authored-by: NK <93352237+Nithish-Yenaganti@users.noreply.github.com> Co-authored-by: 安妮的心动录 <74543653+anneheartrecord@users.noreply.github.com> Co-authored-by: Rick <26716961+Bytechoreographer@users.noreply.github.com> Co-authored-by: Bytechoreographer <Bytechoreographer@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Burak Ömür <burak.omur.1998@gmail.com> Co-authored-by: Dan Lemon <daniel.lemon@amazee.io> Co-authored-by: Vanika Dangi <166420943+vanika02@users.noreply.github.com> Co-authored-by: Jay Gowdy <130084966+jgowdy-godaddy@users.noreply.github.com>
`_query_first_with_cached_plan_fallback` recovers from Postgres's "cached plan must not change result type" by recreating the Prisma client, which drops both the server-side plans and the engine's client-side statement-name cache. Since #30183 the shared reconnect path probes the writer with `SELECT 1` first and skips the recreate when it answers, which is right for the IAM token refresh it was added for and wrong here: the connection is healthy, it is the session's prepared statements that are stale, so the probe always passes and always vetoes the recreate. Callers now pass `force_recreate` to skip that probe, and only the cached-plan fallback does. Getting past the probe is not enough on its own. Both cooldown checks would still skip the recreate for 15 seconds after any earlier reconnect, which outlives the 10 second auth retry window, so a migration landing in that window kept 503ing. `force=True` would fix that but would also let every concurrent caller of the same burst kill the engine the first one just built. The caller instead names the engine generation it observed before the query, and the cooldown is waived only while that generation is still live, so the first caller repairs the pool and the rest fall back to the normal cooldown once it has. Fixes #36418
`_query_first_with_cached_plan_fallback` recovers from Postgres's "cached plan must not change result type" by recreating the Prisma client, which drops both the server-side plans and the engine's client-side statement-name cache. Since #30183 the shared reconnect path probes the writer with `SELECT 1` first and skips the recreate when it answers, which is right for the IAM token refresh it was added for and wrong here: the connection is healthy, it is the session's prepared statements that are stale, so the probe always passes and always vetoes the recreate. Callers now pass `force_recreate` to skip that probe, and only the cached-plan fallback does. Getting past the probe is not enough on its own. Both cooldown checks would still skip the recreate for 15 seconds after any earlier reconnect, which outlives the 10 second auth retry window, so a migration landing in that window kept 503ing. `force=True` would fix that but would also let every concurrent caller of the same burst kill the engine the first one just built. The caller instead names the engine generation it observed before the query, and the cooldown is waived only while that generation is still live, so the first caller repairs the pool and the rest fall back to the normal cooldown once it has. That generation has to come from the engine the query actually ran on. `query_first` is a top-level read, so with a read replica configured it is dispatched to the reader and it is the reader's prepared statements that go stale, while `writer_db` names a different engine with its own counter. The snapshot and the cooldown comparison both read `read_db`, added alongside `writer_db` and backed by a `read_target` property on the routing wrapper that `__getattr__` now dispatches through so the two cannot drift. Two smaller holes on the way out. The waiver is withdrawn while recreates are failing: a failed recreate leaves the generation where it was, so every queued caller would see its own generation still live and run its own full recreate serially rather than collapsing onto one attempt, which is what the cooldown exists for. And a forced recreate that the optimistic-lock guard declines is no longer reported as a success, since the routing wrapper leaves the reader untouched in that case and the caller would otherwise retry the query that just failed. A decline is not counted as a reconnect failure, so the caller's own backoff still gets its waiver on the next attempt. Fixes #36418
`_query_first_with_cached_plan_fallback` recovers from Postgres's "cached plan must not change result type" by recreating the Prisma client, which drops both the server-side plans and the engine's client-side statement-name cache. Since #30183 the shared reconnect path probes the writer with `SELECT 1` first and skips the recreate when it answers, which is right for the IAM token refresh it was added for and wrong here: the connection is healthy, it is the session's prepared statements that are stale, so the probe always passes and always vetoes the recreate. Callers now pass `force_recreate` to skip that probe, and only the cached-plan fallback does. Getting past the probe is not enough on its own. Both cooldown checks would still skip the recreate for 15 seconds after any earlier reconnect, which outlives the 10 second auth retry window, so a migration landing in that window kept 503ing. `force=True` would fix that but would also let every concurrent caller of the same burst kill the engine the first one just built. The caller instead names the engine it observed before the query, and the cooldown is waived only while that engine is still the live one, so the first caller repairs the pool and the rest fall back to the normal cooldown. That engine has to be the one the query actually ran on. `query_first` is a top-level read, so with a read replica configured it is dispatched to the reader and it is the reader's prepared statements that go stale, while `writer_db` names a different engine with its own counter. The observation and the cooldown comparison both go through `read_db`, added alongside `writer_db` and backed by a `read_target` property on the routing wrapper that `__getattr__` now dispatches through so the two cannot drift. The observation carries the wrapper, not just its generation. `read_db` resolves to the reader while it is available and to the writer once it is not, and those counters are independent and both start at zero, so comparing a bare number across that switch pits one engine's counter against another's. Equal by coincidence waives the cooldown for an engine already replaced; unequal gates a caller that needs the recreate. Identity settles it, and is sound because the engine object is never re-pointed without the generation also moving. Three smaller holes on the way out. The waiver is withdrawn once a repair of that same engine has been tried and failed, so a burst collapses onto one attempt instead of each caller running its own recreate serially; the record is keyed per engine rather than counted globally, so an unrelated reconnect failure cannot suppress a stale reader's recovery and a writer failure cannot evict the reader's record. And a forced recreate that the optimistic-lock guard declines is no longer reported as a success on either the direct or the heavy path, since the routing wrapper leaves the reader untouched in that case; a decline is deliberately not counted as a failure, so the caller's own backoff still gets its waiver on the next attempt. Fixes #36418
`_query_first_with_cached_plan_fallback` recovers from Postgres's "cached plan must not change result type" by recreating the Prisma client, which drops both the server-side plans and the engine's client-side statement-name cache. Since #30183 the shared reconnect path probes the writer with `SELECT 1` first and skips the recreate when it answers, which is right for the IAM token refresh it was added for and wrong here: the connection is healthy, it is the session's prepared statements that are stale, so the probe always passes and always vetoes the recreate. Callers now pass `force_recreate` to skip that probe, and only the cached-plan fallback does. Getting past the probe is not enough on its own. Both cooldown checks would still skip the recreate for 15 seconds after any earlier reconnect, which outlives the 10 second auth retry window, so a migration landing in that window kept 503ing. `force=True` would fix that but would also let every concurrent caller of the same burst kill the engine the first one just built. The caller instead names the engine it observed before the query, and the cooldown is waived only while that engine is still the live one, so the first caller repairs the pool and the rest fall back to the normal cooldown. That engine has to be the one the query actually ran on. `query_first` is a top-level read, so with a read replica configured it is dispatched to the reader and it is the reader's prepared statements that go stale, while `writer_db` names a different engine with its own counter. The observation and the cooldown comparison both go through `read_db`, added alongside `writer_db` and backed by a `read_target` property on the routing wrapper that `__getattr__` now dispatches through so the two cannot drift. The observation carries the wrapper, not just its generation. `read_db` resolves to the reader while it is available and to the writer once it is not, and those counters are independent and both start at zero, so comparing a bare number across that switch pits one engine's counter against another's. Equal by coincidence waives the cooldown for an engine already replaced; unequal gates a caller that needs the recreate. Identity settles it, and is sound because the engine object is never re-pointed without the generation also moving. Three smaller holes on the way out. The waiver is withdrawn once a repair of that same engine has been tried and failed, so a burst collapses onto one attempt instead of each caller running its own recreate serially; the record is keyed per engine rather than counted globally, so an unrelated reconnect failure cannot suppress a stale reader's recovery and a writer failure cannot evict the reader's record. And a forced recreate that the optimistic-lock guard declines is no longer reported as a success on either the direct or the heavy path, since the routing wrapper leaves the reader untouched in that case; a decline is deliberately not counted as a failure, so the caller's own backoff still gets its waiver on the next attempt. A decline on the heavy path clears the dead-engine flag before raising. The clear after the cycle is skipped by any raise, which is right for a failure and wrong here, and the non-forced path already clears it on a decline, so this restores that policy rather than inventing one. Stranding the flag would route the next cycle back down the probe-free heavy branch, where the refreshed generation matches and the recreate kills the healthy engine a refresh just spawned, which is #29176. Fixes #36418
`_query_first_with_cached_plan_fallback` recovers from Postgres's "cached plan must not change result type" by recreating the Prisma client, which drops both the server-side plans and the engine's client-side statement-name cache. Since #30183 the shared reconnect path probes the writer with `SELECT 1` first and skips the recreate when it answers, which is right for the IAM token refresh it was added for and wrong here: the connection is healthy, it is the session's prepared statements that are stale, so the probe always passes and always vetoes the recreate. Callers now pass `force_recreate` to skip that probe, and only the cached-plan fallback does. Getting past the probe is not enough on its own. Both cooldown checks would still skip the recreate for 15 seconds after any earlier reconnect, which outlives the 10 second auth retry window, so a migration landing in that window kept 503ing. `force=True` would fix that but would also let every concurrent caller of the same burst kill the engine the first one just built. The caller instead names the engine it observed before the query, and the cooldown is waived only while that engine is still the live one, so the first caller repairs the pool and the rest fall back to the normal cooldown. That engine has to be the one the query actually ran on. `query_first` is a top-level read, so with a read replica configured it is dispatched to the reader and it is the reader's prepared statements that go stale, while `writer_db` names a different engine with its own counter. The observation and the cooldown comparison both go through `read_db`, added alongside `writer_db` and backed by a `read_target` property on the routing wrapper that `__getattr__` now dispatches through so the two cannot drift. The observation carries the wrapper, not just its generation. `read_db` resolves to the reader while it is available and to the writer once it is not, and those counters are independent and both start at zero, so comparing a bare number across that switch pits one engine's counter against another's. Equal by coincidence waives the cooldown for an engine already replaced; unequal gates a caller that needs the recreate. Identity settles it, and is sound because the engine object is never re-pointed without the generation also moving. Three smaller holes on the way out. The waiver is withdrawn once a repair of that same engine has been tried and failed, so a burst collapses onto one attempt instead of each caller running its own recreate serially; the record is keyed per engine rather than counted globally, so an unrelated reconnect failure cannot suppress a stale reader's recovery and a writer failure cannot evict the reader's record. And a forced recreate that the optimistic-lock guard declines is no longer reported as a success on either the direct or the heavy path, since the routing wrapper leaves the reader untouched in that case; a decline is deliberately not counted as a failure, so the caller's own backoff still gets its waiver on the next attempt. A decline on the heavy path clears the dead-engine flag before raising. The clear after the cycle is skipped by any raise, which is right for a failure and wrong here, and the non-forced path already clears it on a decline, so this restores that policy rather than inventing one. Stranding the flag would route the next cycle back down the probe-free heavy branch, where the refreshed generation matches and the recreate kills the healthy engine a refresh just spawned, which is #29176. Clearing that flag is necessary and not sufficient. The escalation check re-arms it whenever the consecutive-failure count sits at the threshold, so a decline that left the count alone sent the very next attempt back down the same path. A decline is raised only at the generation guard, and the generation moves only after a replacement connects, so a decline is proof that a replacement succeeded and the count is reset on it. Fixes #36418
`_query_first_with_cached_plan_fallback` recovers from Postgres's "cached plan must not change result type" by recreating the Prisma client, which drops both the server-side plans and the engine's client-side statement-name cache. Since #30183 the shared reconnect path probes the writer with `SELECT 1` first and skips the recreate when it answers, which is right for the IAM token refresh it was added for and wrong here: the connection is healthy, it is the session's prepared statements that are stale, so the probe always passes and always vetoes the recreate. Callers now pass `force_recreate` to skip that probe, and only the cached-plan fallback does. Getting past the probe is not enough on its own. Both cooldown checks would still skip the recreate for 15 seconds after any earlier reconnect, which outlives the 10 second auth retry window, so a migration landing in that window kept 503ing. `force=True` would fix that but would also let every concurrent caller of the same burst kill the engine the first one just built. The caller instead names the engine it observed before the query, and the cooldown is waived only while that engine is still the live one, so the first caller repairs the pool and the rest fall back to the normal cooldown. That engine has to be the one the query actually ran on. `query_first` is a top-level read, so with a read replica configured it is dispatched to the reader and it is the reader's prepared statements that go stale, while `writer_db` names a different engine with its own counter. The observation and the cooldown comparison both go through `read_db`, added alongside `writer_db` and backed by a `read_target` property on the routing wrapper that `__getattr__` now dispatches through so the two cannot drift. The observation carries the wrapper, not just its generation. `read_db` resolves to the reader while it is available and to the writer once it is not, and those counters are independent and both start at zero, so comparing a bare number across that switch pits one engine's counter against another's. Equal by coincidence waives the cooldown for an engine already replaced; unequal gates a caller that needs the recreate. Identity settles it, and is sound because the engine object is never re-pointed without the generation also moving. Three smaller holes on the way out. The waiver is withdrawn once a repair of that same engine has been tried and failed, so a burst collapses onto one attempt instead of each caller running its own recreate serially; the record is keyed per engine rather than counted globally, so an unrelated reconnect failure cannot suppress a stale reader's recovery and a writer failure cannot evict the reader's record. And a forced recreate that the optimistic-lock guard declines is no longer reported as a success on either the direct or the heavy path, since the routing wrapper leaves the reader untouched in that case; a decline is deliberately not counted as a failure, so the caller's own backoff still gets its waiver on the next attempt. A decline on the heavy path clears the dead-engine flag before raising. The clear after the cycle is skipped by any raise, which is right for a failure and wrong here, and the non-forced path already clears it on a decline, so this restores that policy rather than inventing one. Stranding the flag would route the next cycle back down the probe-free heavy branch, where the refreshed generation matches and the recreate kills the healthy engine a refresh just spawned, which is #29176. Clearing that flag is necessary and not sufficient. The escalation check re-arms it whenever the consecutive-failure count sits at the threshold, so a decline that left the count alone sent the very next attempt back down the same path. A decline is raised only at the generation guard, and the generation moves only after a replacement connects, so a decline is proof that a replacement succeeded and the count is reset on it. Fixes #36418 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Relevant issues
Fixes #29176
Linear ticket
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
make test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewDelays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
CI (LiteLLM team)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Screenshots / Proof of Fix
This is a concurrency race in the IAM-refresh path (no UI, no steady-state log line to screenshot), so the proof is the before/after of the end-to-end repro test, test_iam_refresh_racing_reconnect_recreates_engine_only_once.
It reproduces the exact reported scenario: an RDS IAM token refresh is mid-recreate — it has already spawned a fresh query-engine (PID 222) and is holding _reconnection_lock — when an in-flight transport-error reconnect (attempt_db_reconnect(force=True)) fires. The test asserts the engine is recreated exactly once (prisma_constructed == 1) and that the refresh's freshly-spawned engine (PID 222) is never killed (fresh_engine_killed == False).
❌ Before the fix — source reverted to the base branch, the repro test kept. The racing reconnect undoes the refresh, reproducing the kill/respawn storm from the issue:
$ pytest test_iam_refresh_racing_reconnect_recreates_engine_only_once
FAILED
{'prisma_constructed': 2} != {'prisma_constructed': 1} # engine spawned twice
{'fresh_engine_killed': True} != {'fresh_engine_killed': False} # refresh's new engine (PID 222) killed by the reconnect
1 failed
✅ After the fix — the racing reconnect probes the writer / checks the engine generation, recognizes the planned restart, and no-ops. The engine is spawned once and the fresh engine is never killed:
$ pytest test_iam_refresh_racing_reconnect_recreates_engine_only_once
tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py::test_iam_refresh_racing_reconnect_recreates_engine_only_once PASSED [100%]
1 passed
Full affected-suite run (engine watcher, reconnect cycle, IAM token expiry, routing wrapper, planned-restart, self-heal, backoff-retry): 311 passed, mypy clean, and codecov reports all modified lines covered.
Type
🐛 Bug Fix
✅ Test
Changes
Problem
When IAM_TOKEN_DB_AUTH is enabled, the proxy refreshes its RDS IAM token on a ~12-minute cadence. Each refresh recreates the Prisma client, which SIGKILLs the running query-engine and spawns a new one. The problem: this planned engine restart was indistinguishable from an engine crash, and three independent reconnect paths guarded by two uncoordinated locks all reacted to it — recreating the client again and again, each time killing the engine a prior path had just spawned.
The three colliding paths:
IAM refresh — PrismaWrapper._safe_refresh_token (holds _reconnection_lock) → recreate_prisma_client kills the old engine and spawns a new one.
Engine-death watcher — sees that kill and assumes a crash, calling attempt_db_reconnect(reason="engine_process_death", force=True), which takes a different lock (PrismaClient._db_reconnect_lock) and recreates the client again, killing the engine the refresh just spawned.
In-flight queries — requests hitting the brief restart window fail with httpx.ConnectError/ReadError, get classified as transport errors, and trigger their own attempt_db_reconnect via call_with_db_reconnect_retry → yet another recreate.
The net effect is a storm of engine kill/respawn cycles and a burst of DB exception errors on every token refresh.
Fix
Make planned restarts first-class and idempotent, and close the two-lock gap:
Planned-death tracking — recreate_prisma_client records the old engine PID in _expected_engine_deaths before killing it. All four watcher death detectors (waitpid thread, pidfd, already-dead probe, os.kill poll) consume that PID and skip the reconnect instead of treating it as a crash.
Engine-generation guard — every successful recreate increments a monotonic _engine_generation. Callers pass expected_generation as an optimistic-lock token; if it no longer matches once the lock is held, another path already replaced the engine and the recreate becomes a no-op. recreate_prisma_client now serializes all recreations through _reconnection_lock, so the refresh path and the reconnect paths can no longer recreate concurrently.
Probe-first reconnect — the direct (engine-alive) reconnect path runs SELECT 1 before recreating. A healthy connection — e.g. the engine was already replaced by a refresh — skips the recreate entirely instead of needlessly killing a working engine.
Refresh coalescing + watcher re-arm — _safe_refresh_token skips when the current token still has more than the refresh buffer of runway, so stacked triggers (proactive loop racing the getattr fallback) don't each restart the engine. A new on_engine_replaced hook lets the owning PrismaClient re-arm its watcher on the new PID after a planned restart.
RoutingPrismaWrapper forwards expected_generation to the writer and skips recreating the reader when the writer recreate was skipped.
Testing
New tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py, including an end-to-end repro of the race: an IAM refresh held mid-recreate while an in-flight transport-error reconnect fires — asserts the engine is recreated exactly once and the freshly-spawned engine is never killed.
New planned-death coverage across all four watcher death detectors, plus probe-first and generation-guard cases for the reconnect cycle.
Red-green verified: with the source fix reverted, the key regression tests fail (prisma_constructed: 2, fresh engine killed); restored, they pass.
336 passing across the 7 affected test files plus the exception-handler reconnect paths. Existing tests that pinned the old "always recreate" behavior were updated to the probe-first contract; the #26191 guarantee (never call the blocking disconnect()) is preserved.