scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

Merge remote-tracking branch 'origin/main' into litellm_mistral_magistral_reasoning_effort
Some checks failed
Cost map pins recheck / rerun-cost-map-guard (push) Has been cancelled
Details
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
Details
LiteLLM Rust DB / db-queries (push) Has been cancelled
Details
LiteLLM Rust / rust-lint (push) Has been cancelled
Details
LiteLLM Rust / rust-test (push) Has been cancelled
Details
LiteLLM Rust / rust-wheel (push) Has been cancelled
Details
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Details

Browse source 
# Conflicts:
#	litellm/llms/mistral/chat/transformation.py
This commit is contained in:
mateo-berri 2026-11-07 03:22:08 -08:00
commit e59bf7a477
1618 changed files with 115334 additions and 69933 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
.github/PULL_REQUEST_TEMPLATE/general.md vendored
View file

@ -6,19 +6,12 @@
## Handwritten 1-3 sentence description
<!-- The author writes a plain summary in their own words before requesting review.
Wrap newly AI-written summary text in the make-me-human-readable tags below, including backticks.
The author reviews or rewrites the draft and removes the tags.
When updating an existing PR, preserve the author's wording and removed tags.
Changes elsewhere in the body do not make an unchanged summary an AI draft again.
Add tags again only around newly AI-written replacement summary text, never merely to match this template.
This preserves author review across later updates (2026-11-02, #49779).
<!-- The author writes this in their own words, and it can be left empty. AI agents never write here and leave the
placeholder line below as is, and when updating an existing PR they keep whatever the author wrote. Keep this heading
Say what the PR does for users and why, then anything that keeps it working over time (a sync job, a config someone owns)
No bullets, tables, or code names here. Everything else goes inside the Details block below -->
`<make-me-human-readable>`
<blah>
`</make-me-human-readable>`
[insert handwritten description here]
<details>
<summary>Details</summary>
@ -100,6 +93,7 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
Otherwise proof must be fully e2e, no mocks, with real provider calls where applicable. Pytest alone is insufficient. The following format applies to live proof
Show only the latest proof. Before uses the merge base and After the tested tip. Rerun affected QA after behavior changes. Verified neutral diffs may carry proof with the tested SHA, current tip, equivalence reason and focused checks, following AGENTS.md and /qa
Use matching Before/After case headings with numbered commands and observed output. Put shared setup above them. For one case, omit case headings
A reviewer must be able to rerun every case from the body alone, so the shared setup gives the proxy start command, the config.yaml, env var names (never values) and the client code or curl used, all exactly as run. Wrap any block longer than about 20 lines in a collapsed <details> with a one-line summary
### Before (<hash>)

2
.github/workflows/image-scan.yml vendored
View file

@ -53,7 +53,7 @@ jobs:
if: >-
github.event_name != 'pull_request' ||
github.event.pull_request.head.repo.full_name == github.repository
timeout-minutes: 30
timeout-minutes: 60
permissions:
contents: read
steps:

4
.github/workflows/interactions-spec-canary.yml vendored
View file

@ -49,9 +49,9 @@ jobs:
id: compliance
run: |
uv run --no-sync pytest tests/unit/interactions/test_openapi_compliance.py \
--interactions-spec-url=https://ai.google.dev/static/api/interactions.openapi.json \
--interactions-spec-url=https://aistudio.google.com/api-reference/interactions.openapi.json \
--tb=short -vv -p no:rerunfailures -p no:pytest-retry \
--allow-hosts=127.0.0.1,::1,localhost,ai.google.dev
--allow-hosts=127.0.0.1,::1,localhost,ai.google.dev,aistudio.google.com
tracking-issue:
name: tracking-issue

2
.github/workflows/test-linting.yml vendored
View file

@ -21,7 +21,7 @@ jobs:
pull-requests: read
actions: read
runs-on: ubuntu-latest
timeout-minutes: 20
timeout-minutes: 30
services:
redis:
image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f

25
.github/workflows/test-unit.yml vendored
View file

@ -45,7 +45,7 @@ jobs:
id: rust-bridge-key
if: steps.changes.outputs.decision != 'skip'
env:
RUST_BRIDGE_CACHE_KEY: ${{ runner.os }}-rust-bridge-editable-${{ hashFiles('litellm-rust/Cargo.lock', 'litellm-rust/Cargo.toml', 'litellm-rust/crates/**', 'model_prices_and_context_window.json', 'pyproject.toml') }}
RUST_BRIDGE_CACHE_KEY: ${{ runner.os }}-rust-bridge-editable-${{ hashFiles('litellm-rust/Cargo.lock', 'litellm-rust/Cargo.toml', 'litellm-rust/crates/**', 'model_prices_and_context_window.json', 'pyproject.toml', 'packaging/litellm-proxy/pyproject.toml') }}
run: echo "key=$RUST_BRIDGE_CACHE_KEY" >> "$GITHUB_OUTPUT"
- name: Set up Python
@ -98,6 +98,18 @@ jobs:
mkdir -p rust-bridge-dist/litellm/rust_bridge
cp litellm/rust_bridge/_native.abi3.so rust-bridge-dist/litellm/rust_bridge/_native.abi3.so
- name: Build the gateway wheel once for test shards
if: steps.changes.outputs.decision != 'skip' && steps.rust-bridge-cache.outputs.cache-hit != 'true'
timeout-minutes: 15
working-directory: packaging/litellm-proxy
env:
RUSTC_WRAPPER: sccache
SCCACHE_GHA_ENABLED: "true"
SCCACHE_GHA_VERSION: litellm-ci-v1
SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }}
CARGO_INCREMENTAL: "0"
run: uv run --no-project --with maturin==1.15.0 maturin build --profile dev --strip --out ../../rust-bridge-dist
- name: Upload Rust bridge artifact
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 10
@ -170,6 +182,7 @@ jobs:
tests/unit/google_genai
tests/unit/router_strategy
tests/unit/router_utils
tests/unit/redis
tests/unit/proxy/common_utils/test_cache_aware_routing.py
tests/unit/enterprise/enterprise_callbacks/send_emails
tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py
@ -1022,9 +1035,9 @@ jobs:
run: |
diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json
if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
else
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl
cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so
uv run --no-sync python -c "import importlib.metadata; import litellm.rust_bridge._native; print(importlib.metadata.version('litellm'))"
@ -1281,7 +1294,7 @@ jobs:
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard
- name: Cache Prisma binaries
if: steps.changes.outputs.decision != 'skip'
@ -1484,9 +1497,9 @@ jobs:
RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }}
run: |
if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
else
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl
cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so
fi

3
AGENTS.md
View file

@ -45,14 +45,13 @@ Run QA proxies with `python litellm/proxy/proxy_cli.py --config litellm/proxy/de
Simple PRs may skip before/after e2e when the full diff is mechanical and focused checks completely prove it, such as #49030 restoring an existing helper name. Trace affected callers and run relevant tests/static checks at the tip. For a bug fix, establish that the regression catches the original defect. Size alone never qualifies. New behavior, API/provider semantics, auth, billing, persistence, concurrency or uncertain integration risk need live e2e. In Screenshots / Proof of Fix, record `E2E exception at <tip SHA>`, scope, rationale, commands and results. Tests count as proof only for this exception. Reassess changed diffs, retain CI/bot gates, and honor explicit e2e requests (2026-10-28, unnecessary live QA for #49030)
For other PRs, proof must show the end user's real flow before and after, with commands and output against a live proxy and real provider calls where applicable. Pytest alone is insufficient. For UI or interactive coding-tool flows, drive that surface and embed before/after screenshots, with URLs, clicks and input fields a reviewer can reproduce
For other PRs, proof must show the end user's real flow before and after, with commands and output against a live proxy and real provider calls where applicable. The body alone must let a reviewer rerun it, with the proxy command, config, env var names and client code exactly as run, and long blocks collapsed in <details>. Pytest alone is insufficient. For UI or interactive coding-tool flows, drive that surface and embed before/after screenshots, with URLs, clicks and input fields a reviewer can reproduce
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
- don't use emojis
- don't use "—". Don't use ";" or ":" unless it would be nonsensical not to. Instead, reach for ",", ".", conjunction words, etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them also feels AI-y. A word cap does not penalize you for adding more sentences
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a ".". Exception is that docs (i.e., litellm-docs) should have trailing "."
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure

6
Dockerfile
View file

@ -76,6 +76,7 @@ ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
COPY pyproject.toml uv.lock ./
COPY enterprise/pyproject.toml enterprise/
COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/
COPY packaging/litellm-proxy/pyproject.toml packaging/litellm-proxy/
# Install third-party dependencies (cached unless pyproject.toml/uv.lock change)
RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \
@ -101,7 +102,8 @@ COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/
RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
# Install project and workspace packages (fast - deps already cached)
RUN uv sync --frozen --no-default-groups --no-editable \
RUN uv build --package litellm-proxy --wheel --out-dir /wheels
RUN uv sync --frozen --no-install-package litellm-proxy --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--group admin-mcp \
@ -111,6 +113,8 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra bedrock-realtime \
--python python3.13
RUN uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm_proxy-*.whl
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
npm_config_cache=/root/.npm \
prisma generate --schema=./schema.prisma

40
cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md
View file

@ -46,7 +46,43 @@ For a managed workspace rollout, use OpenAI's [gateway configuration guide](http
## Native compatibility and Pi
For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials but does not need a model deployment or resolve model aliases. Restricted keys must allow the upstream model name.
For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials. A ChatGPT model keeps its upstream name and needs no deployment of its own.
### Mix ChatGPT and LiteLLM models
On the native route, Codex's model picker shows OpenAI's catalog first and then the other chat and responses models the virtual key can use. Pick one of those and Codex's requests go through LiteLLM's `/v1/responses` with the key's usual routing, budgets, guardrails, and logging. The proxy removes the ChatGPT bearer and account header from that request, so the other provider and your callbacks never receive them. A name is routed this way when another provider serves it as an exact `model_name`, through a router `model_group_alias`, or as a team's public model name. A ChatGPT model, a name that only a wildcard deployment matches, or a name OpenAI's catalog lists for the caller's ChatGPT account keeps the passthrough, so a model OpenAI adds later still uses the subscription even when a deployment of another provider has the same name. The proxy learns the account's list from the catalog Codex fetches, or asks OpenAI with the caller's login. When it cannot get the list, for example because OpenAI is down, the request times out, or the client is not Codex, every request keeps the passthrough for the next minute, as it did before LiteLLM models were added to the picker. When the key can use many models, the picker lists the ones that fit in Codex's catalog size limit, and the proxy logs a warning with how many it left out.
Codex's picker sends a `none` reasoning effort for a LiteLLM model whose catalog entry lists no effort levels. With `litellm_settings: drop_params: true`, the proxy drops it for a model that does not support `none`, and the provider applies its own default. A model whose cost map entry or `model_info` sets `supports_none_reasoning_effort: true` keeps it. Without `drop_params`, the effort reaches the provider unchanged.
Codex also sends its ChatGPT app tools with every turn, which on a LiteLLM model adds a few hundred thousand prompt tokens to each turn. Set `apps = false` under `[features]` in Codex's `config.toml` to turn them off.
Only a key that may call `/v1/responses` gets these models. A key whose `allowed_routes` covers only the native routes sees OpenAI's catalog unchanged, and every request it sends uses the subscription.
Codex sends earlier reasoning items back with each turn. The proxy tags the reasoning that a LiteLLM model returns on this route, OpenAI and Azure deployments included. When a session switches to a ChatGPT model, the proxy drops those items before forwarding the request to OpenAI, which would reject them. Messages, tool calls, and ChatGPT's own reasoning are kept, and a LiteLLM model still gets its own reasoning back.
When the catalog adds LiteLLM models for a key, its `ETag` is a digest of OpenAI's ETag and those models, and that key's responses carry the same value in `X-Models-Etag`. Codex keeps its cached catalog while the two match and fetches it again when OpenAI's catalog or the key's models change. A key whose catalog adds no LiteLLM models keeps OpenAI's ETags.
### Restrict native models with an access group
A key with a model allowlist needs a rule for ChatGPT names that OpenAI adds over time. Put the subscription wildcards in an access group under `chatgpt/` names:
```yaml
model_list:
- model_name: chatgpt/gpt-*
litellm_params:
model: chatgpt/gpt-*
model_info:
access_groups: ["chatgpt-native"]
- model_name: chatgpt/codex-*
litellm_params:
model: chatgpt/codex-*
model_info:
access_groups: ["chatgpt-native"]
```
Then add `chatgpt-native` to the key's models, next to any other model names it should reach. On the native route, a key that holds only the `chatgpt/` form of a name is checked as that form, so a request for `gpt-5.6-sol` is allowed when the key may call `chatgpt/gpt-5.6-sol`. The key can call every ChatGPT model that matches the wildcards, including ones the proxy configuration does not name.
On `/v1` these wildcards match only a request that asks for a `chatgpt/` name, so a request for a plain `gpt-` or `codex-` name still goes where it went before, such as a `*` catch-all deployment. A `/v1` request for a `chatgpt/` name without a ChatGPT login falls back to the proxy's own saved login if it has one, so keep that login unset as described above.
### Configure Pi
@ -71,7 +107,7 @@ Keep Pi's subscription login as the provider credential. The header value names
Send a prompt and check the completed request in LiteLLM's Logs page. Its gateway user/key and token counts should match the caller. Give each person their own LiteLLM virtual key for separate attribution. The gateway credential and ChatGPT login are independently authenticated; LiteLLM does not establish that both identities belong to the same person.
The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key must allow that name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions.
The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key needs the `chatgpt-native` access group above or that exact name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions.
The proxy forwards the ChatGPT bearer and account header to the fixed OpenAI Codex endpoint. It removes gateway credentials and cookies from the upstream request and redacts credential headers in logging callbacks. The client remains responsible for refresh and reauthentication. Native passthrough returns authentication and quota errors without switching subscriptions. Standard model routes retain the caller login across ChatGPT retries; explicitly configured cross-provider fallbacks use their own credentials. An explicit workspace header is preserved for OpenAI to validate.

60
cookbook/skill_scanner/README.md Normal file
View file

@ -0,0 +1,60 @@
# Manual skill security scans
This standalone adapter runs NVIDIA SkillSpector on a separate machine. It is not installed in the LiteLLM proxy and does not require a shared Helm or Compose deployment
Start it with Python 3.12 or newer and uv. The script pins SkillSpector and its HTTP dependencies
```sh
export SKILL_SCANNER_API_KEY="$(openssl rand -hex 32)"
uv run server.py
```
For another host, place the loopback service behind an HTTPS reverse proxy that forwards the Authorization header and enforces a 65 MiB request limit. Run it under an unprivileged account with a writable temporary directory and no production files or credentials. The scanner host needs Git and network access to requested sources and optional model providers. Apply network egress restrictions suitable for those sources
Configure the same scanner secret in the gateway's environment, then add this to LiteLLM's configuration
```yaml
general_settings:
skill_scanner:
endpoint: https://scanner.example.com/scan
api_key_env: SKILL_SCANNER_API_KEY
use_llm: false
timeout_seconds: 180
```
The endpoint must be reachable from the gateway. A customer's laptop works while it is reachable through their network, VPN or authenticated HTTPS tunnel. `localhost` refers to the gateway machine, not the admin's browser. HTTP is accepted only for loopback development
Open Skills, Submissions, and Review as a proxy admin. Click **Scan for threats**, inspect the report, then independently approve or reject the proposal. Scanning is optional and never changes approval permissions or blocks publication. The report shows findings, locations, remediation, risk score, scanner version and whether semantic analysis actually ran. Incomplete reports do not establish safety
Uploaded ZIPs are sent privately to the scanner with their exact stored bytes. Metadata-only edits reuse the matching stored archive. If that archive is missing or changed, submit a new ZIP before scanning. The displayed SHA-256 identifies that archive, and the scan examines its files rather than proposed catalog metadata. Remote URL reports describe the source at scan time and do not freeze its contents. GitHub branch/tag refs and explicit-ref subdirectories are supported. Other ref/subdirectory combinations require a ZIP so the adapter does not silently scan the wrong revision
The latest scan result is retained with the submission. Running scans have a deadline. If a gateway restarts or the scanner becomes unavailable, the attempt expires and an admin can retry. Retries replace the latest report. Approving or rejecting during a scan does not cancel it or change the decision
Set `use_llm: true` to request semantic analysis and configure the provider on the scanner machine using [SkillSpector's provider instructions](https://github.com/NVIDIA/skillspector#llm-analysis). Provider credentials stay on that machine. Static analysis can also call OSV for dependency checks
## HTTP contract
The gateway sends an authenticated `POST` to the configured endpoint. The body contains form fields `source` (JSON source descriptor), `use_llm` (boolean), and `expires_at` (timezone-aware ISO timestamp), plus `file` (ZIP) when bytes are available. The service returns HTTP 200 with this JSON shape
```json
{
"scanner_version": "2.12.0",
"complete": true,
"scan_mode": "static-only",
"risk_score": 70,
"findings": [{
"rule_id": "example-rule",
"severity": "HIGH",
"message": "Potential data exfiltration",
"file": "SKILL.md",
"start_line": 8,
"explanation": "Instructions request sending local files to an external host",
"remediation": "Remove the external upload instruction"
}],
"limitations": []
}
```
Reports are limited to 1 MiB and 500 findings. Scanner uploads allow 64 MiB because the gateway can expand compressed skill uploads into stored compatibility ZIPs. The gateway rejects malformed reports, refuses redirects and verifies TLS. Only the configured scanner credential is sent. A timeout, HTTP error or incomplete report remains visible to the reviewer
The bundled adapter runs at most two child scans concurrently. Each child and its descendants are terminated at the gateway deadline, capped at 570 seconds. Keep the gateway and scanner clocks synchronized. Direct adapter clients may omit `expires_at` to use the cap. Child processes never execute submitted skill scripts. Raise the gateway deadline up to 600 seconds when needed for semantic analysis and configure the reverse proxy timeout to match

172
cookbook/skill_scanner/server.py Normal file
View file

@ -0,0 +1,172 @@
# /// script
# requires-python = ">=3.12"
# dependencies = [
# "fastapi==0.162.0", "uvicorn==0.54.0", "python-multipart==0.0.33",
# "skillspector @ git+https://github.com/NVIDIA/skillspector.git@62c2fe8622f580ac292d4078c3b71f790194dbbd",
# ]
# ///
import asyncio
import hmac
import json
import os
import signal
import sys
import tempfile
from contextlib import suppress
from datetime import datetime, timezone
from pathlib import Path
from typing import Annotated, Final
from urllib.parse import quote, urlsplit
from fastapi import Depends, FastAPI, File, Form, HTTPException, UploadFile
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from pydantic import AwareDatetime, TypeAdapter, ValidationError
MAX_UPLOAD: Final = 64 * 1024 * 1024
MAX_REPORT: Final = 1024 * 1024
SCAN_TIMEOUT: Final = 570
SOURCE: Final = TypeAdapter(dict[str, str])
SLOTS: Final = asyncio.Semaphore(2)
app: Final = FastAPI(title="SkillSpector adapter", docs_url=None, redoc_url=None, openapi_url=None)
def authenticate(credentials: Annotated[HTTPAuthorizationCredentials, Depends(HTTPBearer())]) -> None:
token: Final = os.environ.get("SKILL_SCANNER_API_KEY", "")
if not token or not hmac.compare_digest(credentials.credentials, token):
raise HTTPException(401, "Invalid scanner credentials")
def source_target(source: dict[str, str]) -> str:
target: Final = (
"https://github.com/" + source.get("repo", "") if source.get("source") == "github" else source.get("url", "")
)
parsed: Final = urlsplit(target)
if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password or parsed.fragment:
raise HTTPException(400, "Remote sources require an HTTPS URL without credentials or a fragment")
ref: Final = source.get("ref")
subdir: Final = source.get("path")
if ref or subdir:
if parsed.hostname != "github.com" or len(parsed.path.strip("/").split("/")) != 2 or not ref:
raise HTTPException(400, "Upload a ZIP to scan this source's exact ref and subdirectory")
suffix: Final = "/" + quote(subdir, safe="/") if subdir else ""
return target.removesuffix(".git").rstrip("/") + "/tree/" + quote(ref, safe="/") + suffix
return target
async def worker(target: str, output: Path, use_llm: bool) -> None:
from skillspector.mcp_server import run_scan
result: Final = await run_scan(target, use_llm=use_llm, allow_local_targets=not target.startswith("https://"))
coverage: Final = result.get("analysis_completeness", {})
findings: Final = result["findings"]
complete: Final = (
result.get("execution_successful") is True
and coverage.get("is_complete") is True
and (not use_llm or (result.get("llm_available") is True and result.get("llm_used") is True))
and len(findings) <= 500
)
limitations: Final = tuple(str(value)[:2000] for value in coverage.get("limitations", ()))
report: Final = {
"scanner_version": result["version"],
"complete": complete,
"scan_mode": result["scan_mode"],
"risk_score": result["risk_score"],
"findings": [
{
"rule_id": finding["id"],
"severity": finding["severity"],
"message": finding.get("pattern") or finding.get("finding") or "Security finding",
"file": finding["location"]["file"],
"start_line": finding["location"].get("start_line"),
"explanation": finding.get("explanation"),
"remediation": finding.get("remediation"),
}
for finding in findings[:500]
],
"limitations": (
*limitations,
*(("Only the first 500 findings are displayed.",) if len(findings) > 500 else ()),
),
}
encoded: Final = json.dumps(report).encode()
if len(encoded) > MAX_REPORT:
raise ValueError("Scanner report exceeded the response limit")
output.write_bytes(encoded)
async def run_worker(target: str, output: Path, use_llm: bool, expires_at: datetime | None = None) -> dict[str, object]:
timeout: Final = (
SCAN_TIMEOUT
if expires_at is None
else min(SCAN_TIMEOUT, (expires_at - datetime.now(timezone.utc)).total_seconds())
)
if timeout <= 0:
raise HTTPException(504, "SkillSpector scan deadline expired")
process: Final = await asyncio.create_subprocess_exec(
sys.executable,
str(Path(__file__).resolve()),
"--scan",
target,
str(output),
str(use_llm).lower(),
stdout=asyncio.subprocess.DEVNULL,
stderr=asyncio.subprocess.DEVNULL,
start_new_session=True,
env={
**os.environ,
"HOME": str(output.parent),
"TMPDIR": str(output.parent),
"GIT_CONFIG_GLOBAL": os.devnull,
"GIT_CONFIG_NOSYSTEM": "1",
"GIT_TERMINAL_PROMPT": "0",
"GIT_ALLOW_PROTOCOL": "https",
},
)
try:
code: Final = await asyncio.wait_for(process.wait(), timeout)
if code != 0 or not output.is_file() or output.stat().st_size > MAX_REPORT:
raise HTTPException(502, "SkillSpector could not produce a report for this source")
return TypeAdapter(dict[str, object]).validate_json(output.read_bytes())
except TimeoutError:
raise HTTPException(504, "SkillSpector scan timed out") from None
finally:
with suppress(ProcessLookupError):
os.killpg(process.pid, signal.SIGKILL)
await process.wait()
@app.post("/scan", dependencies=[Depends(authenticate)])
async def scan(
source: Annotated[str, Form(max_length=16384)],
use_llm: Annotated[bool, Form()] = False,
expires_at: Annotated[AwareDatetime | None, Form()] = None,
file: Annotated[UploadFile | None, File()] = None,
) -> dict[str, object]:
if SLOTS.locked():
raise HTTPException(429, "Scanner is busy; try again shortly")
try:
parsed: Final = SOURCE.validate_json(source)
except ValidationError:
raise HTTPException(400, "Invalid source metadata") from None
async with SLOTS:
with tempfile.TemporaryDirectory(prefix="skill-scan-") as directory:
root: Final = Path(directory)
if file is not None:
content: Final = await file.read(MAX_UPLOAD + 1)
if len(content) > MAX_UPLOAD:
raise HTTPException(413, "Upload exceeds 64 MiB")
archive: Final = root / "skill.zip"
archive.write_bytes(content)
return await run_worker(str(archive), root / "report.json", use_llm, expires_at)
return await run_worker(source_target(parsed), root / "report.json", use_llm, expires_at)
if __name__ == "__main__":
if len(sys.argv) > 1 and sys.argv[1] == "--scan":
asyncio.run(worker(sys.argv[2], Path(sys.argv[3]), sys.argv[4] == "true"))
else:
import uvicorn
if not os.environ.get("SKILL_SCANNER_API_KEY"):
raise SystemExit("Set SKILL_SCANNER_API_KEY before starting the scanner")
uvicorn.run(app, host="127.0.0.1", port=int(os.environ.get("PORT", "8099")))

6
docker/Dockerfile.database
View file

@ -74,6 +74,7 @@ ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
COPY pyproject.toml uv.lock ./
COPY enterprise/pyproject.toml enterprise/
COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/
COPY packaging/litellm-proxy/pyproject.toml packaging/litellm-proxy/
# Install third-party dependencies (cached unless pyproject.toml/uv.lock change)
RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \
@ -99,7 +100,8 @@ COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/
RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
# Install project and workspace packages (fast - deps already cached)
RUN uv sync --frozen --no-default-groups --no-editable \
RUN uv build --package litellm-proxy --wheel --out-dir /wheels
RUN uv sync --frozen --no-install-package litellm-proxy --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--group admin-mcp \
@ -109,6 +111,8 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra bedrock-realtime \
--python python3.13
RUN uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm_proxy-*.whl
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
npm_config_cache=/root/.npm \
prisma generate --schema=./schema.prisma

13
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -8,7 +8,6 @@ import re
from collections.abc import Callable, Mapping, Sequence
from dataclasses import replace as dataclasses_replace
from datetime import datetime, timedelta
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast
from pydantic import TypeAdapter, ValidationError
@ -766,7 +765,6 @@ class CheckBatchCost:
credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {}
_file_content = await afile_content(
file_id=cast(str, raw_output_file_id),
_litellm_internal_model_credentials=MappingProxyType(dict(credentials)),
**cast(FileContentCallOptions, credentials),
)
@ -867,10 +865,7 @@ class CheckBatchCost:
error_file_failed_requests: Final = await count_error_file_failed_requests(
response,
custom_llm_provider=batch_file_provider,
litellm_params={
**credentials,
"_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),
},
litellm_params=credentials,
)
batch_result: Final = (
output_file_result
@ -1029,11 +1024,7 @@ class CheckBatchCost:
custom_llm_provider=deployment_info.litellm_params.custom_llm_provider,
)
cost_provider: Final = cast(BatchFileProvider, llm_provider)
credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {}
file_access: Final = {
**credentials,
"_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),
}
file_access: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {}
output_content: Final = (
await fetch_batch_managed_file_content(
response.output_file_id, custom_llm_provider=cost_provider, litellm_params=file_access

30
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -127,7 +127,6 @@ class _ManagedFileRetrieve(Protocol):
self,
*,
file_id: str,
_litellm_internal_model_credentials: Mapping[str, object] | None = None,
**kwargs: Unpack[FileRetrieveCallOptions],
) -> OpenAIFileObject: ...
@ -400,7 +399,7 @@ class _RouterFileCallKwargs(TypedDict, total=False):
_MODEL_FILE_ID_MAPPING: Final = TypeAdapter(dict[str, str])
_FILE_UPLOAD_ONLY_FIELDS: Final = frozenset({"file", "purpose", "expires_after", "extra_body"})
_FILE_DELETE_RESERVED_FIELDS: Final = frozenset({"model", "file_id", "_litellm_internal_model_credentials"})
_FILE_DELETE_RESERVED_FIELDS: Final = frozenset({"model", "file_id"})
@dataclass(frozen=True, slots=True)
@ -459,15 +458,7 @@ async def _delete_provider_file(
provider_file_id: str,
request_data: Mapping[str, object],
) -> None:
credentials: Final = llm_router.registry().credentials(model_id, form="with_provider")
delete_data: Final = {
**{k: v for k, v in request_data.items() if k not in _FILE_DELETE_RESERVED_FIELDS},
**(
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
if credentials is not None
else {}
),
}
delete_data: Final = {k: v for k, v in request_data.items() if k not in _FILE_DELETE_RESERVED_FIELDS}
await llm_router.afile_delete(
model=model_id,
file_id=provider_file_id,
@ -1160,7 +1151,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
*,
provider_file_id: str,
call_options: Mapping[str, object],
internal_model_credentials: Mapping[str, object] | None = None,
) -> OpenAIFileObject:
"""Retrieve provider file details with transient retries and SDK retries disabled."""
retrieve_file: Final = cast(_ManagedFileRetrieve, litellm.afile_retrieve)
@ -1172,14 +1162,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if attempt > 0:
await self._sleep(_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS[attempt - 1])
try:
if internal_model_credentials is None:
return await retrieve_file(
file_id=provider_file_id,
**retrieve_options,
)
return await retrieve_file(
file_id=provider_file_id,
_litellm_internal_model_credentials=MappingProxyType(dict(internal_model_credentials)),
**retrieve_options,
)
except Exception as error:
@ -1221,7 +1205,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
self._afile_retrieve_with_retries(
provider_file_id=provider_file_id,
call_options=router_credentials,
internal_model_credentials=router_credentials,
),
timeout=BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS,
)
@ -2866,15 +2849,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
exception_dict = {}
for model_id, provider_file_id in specific_model_file_id_mapping.items():
try:
# Cloud-storage providers (e.g. Bedrock S3) validate file ids
# against the deployment's configured bucket, which they only
# trust from this immutable server-side snapshot, never from
# request params.
credentials = llm_router.registry().credentials(model_id, form="with_provider")
if credentials is not None:
data["_litellm_internal_model_credentials"] = cast(Dict, MappingProxyType(dict(credentials)))
else:
data.pop("_litellm_internal_model_credentials", None)
return cast(
HttpxBinaryResponseContent,
await llm_router.afile_content(

28
enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py
View file

@ -1,9 +1,9 @@
import enum
from typing import Dict, List, Optional
from pydantic import Field
from litellm.proxy._types import WebhookEvent
from litellm.types.integrations.model_deprecation_email import MODEL_DEPRECATION_EMAIL_EVENT
from litellm.types.llms.base import LiteLLMBaseModel
@ -26,7 +26,7 @@ class SendKeyCreatedEmailEvent(WebhookEvent):
class SendKeyRotatedEmailEvent(WebhookEvent):
virtual_key: str
key_alias: Optional[str] = None
key_alias: str | None = None
"""
The virtual key that was rotated
this will be sk-123xxx, since we will be emailing this to the user to start using the new key
@ -34,23 +34,33 @@ class SendKeyRotatedEmailEvent(WebhookEvent):
class EmailEvent(str, enum.Enum):
model_deprecations = MODEL_DEPRECATION_EMAIL_EVENT
virtual_key_created = "Virtual Key Created"
new_user_invitation = "New User Invitation"
virtual_key_rotated = "Virtual Key Rotated"
soft_budget_crossed = "Soft Budget Crossed"
max_budget_alert = "Max Budget Alert"
class EmailEventSettings(LiteLLMBaseModel):
event: EmailEvent
enabled: bool
class EmailEventSettingsUpdateRequest(LiteLLMBaseModel):
settings: List[EmailEventSettings]
settings: list[EmailEventSettings]
class EmailEventSettingsResponse(LiteLLMBaseModel):
settings: List[EmailEventSettings]
settings: list[EmailEventSettings]
class DefaultEmailSettings(LiteLLMBaseModel):
"""Default settings for email events"""
settings: Dict[EmailEvent, bool] = Field(
settings: dict[EmailEvent, bool] = Field(
default_factory=lambda: {
EmailEvent.model_deprecations: False, # Explicit opt-in for owner retirement reminders.
EmailEvent.virtual_key_created: True, # On by default
EmailEvent.new_user_invitation: True, # On by default
EmailEvent.virtual_key_rotated: True, # On by default
@ -58,10 +68,12 @@ class DefaultEmailSettings(LiteLLMBaseModel):
EmailEvent.max_budget_alert: True, # On by default
}
)
def to_dict(self) -> Dict[str, bool]:
def to_dict(self) -> dict[str, bool]:
"""Convert to dictionary with string keys for storage"""
return {event.value: enabled for event, enabled in self.settings.items()}
@classmethod
def get_defaults(cls) -> Dict[str, bool]:
def get_defaults(cls) -> dict[str, bool]:
"""Get the default settings as a dictionary with string keys"""
return cls().to_dict()
return cls().to_dict()

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.81"
version = "0.1.83"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.81"
version = "0.1.83"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

9
gateway/Dockerfile
View file

@ -41,10 +41,14 @@ COPY pyproject.toml README.md rust-toolchain.toml ./
COPY .cargo ./.cargo
COPY litellm-rust ./litellm-rust
COPY model_prices_and_context_window.json ./
COPY packaging/litellm-proxy/pyproject.toml packaging/litellm-proxy/
RUN --mount=type=cache,target=/root/.cache/uv \
mkdir -p litellm/rust_bridge /wheels && \
touch litellm/__init__.py litellm/rust_bridge/__init__.py && \
uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \
python -c "import maturin; maturin.build_wheel('/wheels')" && \
cd packaging/litellm-proxy && \
uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \
python -c "import maturin; maturin.build_wheel('/wheels')"
@ -85,6 +89,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,source=uv.lock,target=uv.lock \
--mount=type=bind,source=enterprise/pyproject.toml,target=enterprise/pyproject.toml \
--mount=type=bind,source=litellm-proxy-extras/pyproject.toml,target=litellm-proxy-extras/pyproject.toml \
--mount=type=bind,source=packaging/litellm-proxy/pyproject.toml,target=packaging/litellm-proxy/pyproject.toml \
uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
@ -98,7 +103,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
COPY . .
RUN --mount=type=cache,target=/root/.cache/uv \
uv sync --frozen --no-install-project --no-default-groups --no-editable \
uv sync --frozen --no-install-project --no-install-package litellm-proxy --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--group admin-mcp \
@ -108,7 +113,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--python python3.13
RUN --mount=type=bind,from=rust-native,source=/wheels,target=/wheels \
uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl
uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl /wheels/litellm_proxy-*.whl
# PYTHONPATH=/app makes the source tree shadow the installed package, so the
# compiled Rust extension must live next to the source or it is never imported.

16
litellm-proxy-extras/litellm_proxy_extras/migrations/20261105000000_skill_version_history/migration.sql Normal file
View file

@ -0,0 +1,16 @@
SET LOCAL lock_timeout = '10s';
CREATE TABLE IF NOT EXISTS "LiteLLM_ClaudeCodePluginVersion" (
"id" TEXT NOT NULL,
"plugin_id" TEXT NOT NULL,
"version" TEXT,
"description" TEXT,
"manifest_json" TEXT,
"published_at" TIMESTAMP(3),
"published_by" TEXT,
"archive_available" BOOLEAN NOT NULL DEFAULT false,
"content" BYTEA,
CONSTRAINT "LiteLLM_ClaudeCodePluginVersion_pkey" PRIMARY KEY ("id"),
CONSTRAINT "LiteLLM_ClaudeCodePluginVersion_plugin_id_fkey" FOREIGN KEY ("plugin_id") REFERENCES "LiteLLM_ClaudeCodePluginTable"("id") ON DELETE CASCADE ON UPDATE CASCADE
);
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ClaudeCodePluginVersion_plugin_id_version_key" ON "LiteLLM_ClaudeCodePluginVersion"("plugin_id", "version");
CREATE INDEX IF NOT EXISTS "LiteLLM_ClaudeCodePluginVersion_plugin_id_published_at_idx" ON "LiteLLM_ClaudeCodePluginVersion"("plugin_id", "published_at");

27
litellm-proxy-extras/litellm_proxy_extras/migrations/20261105010000_skill_submissions/migration.sql Normal file
View file

@ -0,0 +1,27 @@
SET LOCAL lock_timeout = '10s';
CREATE TABLE IF NOT EXISTS "LiteLLM_SkillSubmission" (
"id" TEXT NOT NULL,
"plugin_name" TEXT NOT NULL,
"operation" TEXT NOT NULL,
"status" TEXT NOT NULL DEFAULT 'pending_review',
"proposal_json" TEXT NOT NULL,
"archive_content" BYTEA,
"submitted_by" TEXT NOT NULL,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"base_plugin_id" TEXT,
"base_updated_at" TIMESTAMP(3),
"reviewed_by" TEXT,
"reviewed_at" TIMESTAMP(3),
"rejection_reason" TEXT,
CONSTRAINT "LiteLLM_SkillSubmission_pkey" PRIMARY KEY ("id")
);
CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_submitted_by_created_at_id_idx"
ON "LiteLLM_SkillSubmission"("submitted_by", "created_at" DESC, "id" DESC);
CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_submitted_by_status_created_at_id_idx"
ON "LiteLLM_SkillSubmission"("submitted_by", "status", "created_at" DESC, "id" DESC);
CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_status_created_at_id_idx"
ON "LiteLLM_SkillSubmission"("status", "created_at" DESC, "id" DESC);
CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_created_at_id_idx"
ON "LiteLLM_SkillSubmission"("created_at" DESC, "id" DESC);

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20261105020000_skill_scanning/migration.sql Normal file
View file

@ -0,0 +1,3 @@
SET LOCAL lock_timeout = '10s';
ALTER TABLE "LiteLLM_SkillSubmission" ADD COLUMN IF NOT EXISTS "scan_json" TEXT;

37
litellm-proxy-extras/litellm_proxy_extras/migrations/20261105094500_add_skill_download_history/migration.sql Normal file
View file

@ -0,0 +1,37 @@
SET LOCAL lock_timeout = '10s';
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1 FROM pg_attribute
WHERE attrelid = to_regclass('"LiteLLM_ClaudeCodePluginTable"')
AND attname = 'anonymous_download_count' AND NOT attisdropped
) THEN
ALTER TABLE "LiteLLM_ClaudeCodePluginTable"
ADD COLUMN IF NOT EXISTS "anonymous_download_count" BIGINT NOT NULL DEFAULT 0;
END IF;
IF to_regclass('"LiteLLM_ClaudeCodePluginDownload"') IS NULL THEN
CREATE TABLE IF NOT EXISTS "LiteLLM_ClaudeCodePluginDownload" (
"id" BIGSERIAL NOT NULL,
"plugin_id" TEXT NOT NULL,
"downloaded_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"version" TEXT,
"digest" TEXT NOT NULL,
"authenticated" BOOLEAN NOT NULL,
"user_id" TEXT,
"user_email" TEXT,
"key_alias" TEXT,
"team_id" TEXT,
"team_alias" TEXT,
CONSTRAINT "LiteLLM_ClaudeCodePluginDownload_pkey" PRIMARY KEY ("id"),
CONSTRAINT "LiteLLM_ClaudeCodePluginDownload_plugin_id_fkey" FOREIGN KEY ("plugin_id")
REFERENCES "LiteLLM_ClaudeCodePluginTable"("id") ON DELETE CASCADE ON UPDATE CASCADE
);
END IF;
IF to_regclass('"LiteLLM_ClaudeCodePluginDownload_plugin_id_id_idx"') IS NULL THEN
CREATE INDEX IF NOT EXISTS "LiteLLM_ClaudeCodePluginDownload_plugin_id_id_idx"
ON "LiteLLM_ClaudeCodePluginDownload"("plugin_id", "id");
END IF;
END $$;

59
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -1616,10 +1616,69 @@ model LiteLLM_ClaudeCodePluginTable {
created_by String?
uploaded_archive LiteLLM_ClaudeCodePluginUpload?
versions LiteLLM_ClaudeCodePluginVersion[]
downloads LiteLLM_ClaudeCodePluginDownload[]
anonymous_download_count BigInt @default(0)
@@map("LiteLLM_ClaudeCodePluginTable")
}
model LiteLLM_ClaudeCodePluginVersion {
id String @id @default(uuid())
plugin_id String
version String?
description String?
manifest_json String?
published_at DateTime?
published_by String?
archive_available Boolean @default(false)
content Bytes?
plugin LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_id], references: [id], onDelete: Cascade, onUpdate: Cascade)
@@unique([plugin_id, version])
@@index([plugin_id, published_at])
}
model LiteLLM_SkillSubmission {
id String @id @default(uuid())
plugin_name String
operation String
status String @default("pending_review")
proposal_json String
archive_content Bytes?
submitted_by String
created_at DateTime @default(now())
base_plugin_id String?
base_updated_at DateTime?
reviewed_by String?
reviewed_at DateTime?
rejection_reason String?
scan_json String?
@@index([submitted_by, created_at(sort: Desc), id(sort: Desc)])
@@index([submitted_by, status, created_at(sort: Desc), id(sort: Desc)])
@@index([status, created_at(sort: Desc), id(sort: Desc)])
@@index([created_at(sort: Desc), id(sort: Desc)])
}
model LiteLLM_ClaudeCodePluginDownload {
id BigInt @id @default(autoincrement())
plugin_id String
downloaded_at DateTime @default(now())
version String?
digest String
authenticated Boolean
user_id String?
user_email String?
key_alias String?
team_id String?
team_alias String?
plugin LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_id], references: [id], onDelete: Cascade, onUpdate: Cascade)
@@index([plugin_id, id])
}
model LiteLLM_ClaudeCodePluginUpload {
plugin_name String @id
content Bytes

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.112"
version = "0.4.114"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -30,7 +30,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.112"
version = "0.4.114"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

2
litellm-rust/.agents/skills/rust-gateway/SKILL.md
View file

@ -7,7 +7,7 @@ metadata:
# Rust gateway
`crates/gateway` mounts endpoint routers and serves the listener. `crates/gateway-inference`, `gateway-mcp` and `gateway-ui` own their endpoints, `gateway-auth` is the inbound authentication middleware and `gateway-management` holds management logic that has no router yet. A crate's `AGENTS.md`, where it has one, says what belongs in it. Read the reference for the part you are touching
`crates/gateway` mounts endpoint routers and serves the listener. `crates/gateway-inference`, `gateway-mcp` and `gateway-ui` own their endpoints, `gateway-auth` is the inbound authentication middleware and `gateway-management` holds management logic that has no router yet. `gateway-forward` is the reverse proxy that relays every request to the Python proxy and has no routes of its own. A crate's `AGENTS.md`, where it has one, says what belongs in it. Read the reference for the part you are touching
- [surfaces](references/surfaces.md)
- Each API surface mirrors the envelope the Python proxy returns there, so the shapes differ on purpose

53
litellm-rust/.agents/skills/rust-python-boundary/SKILL.md Normal file
View file

@ -0,0 +1,53 @@
---
name: rust-python-boundary
description: Decide where Python-shaped input becomes typed Rust in litellm-rust. Use when reading config.yaml values, os.environ/ references, REDIS_* or other environment layering, get_secret fallbacks, or kwargs and objects handed over from Python, when coercing text to a bool, int or float the way Python does, or when adding a dependency on litellm-python-compat
metadata:
internal: true
---
# Python boundary
Python-shaped input becomes typed Rust in exactly two crates. Every crate past them handles typed values, so no crate re-guesses how Python would have read a value
## The two boundaries
- `crates/config` reads config.yaml, which is Python-shaped: `os.environ/` references, scalars that Helm or the environment turned into text, raw kwargs maps. It also owns the environment layering and `get_secret` fallbacks Python applies while it loads that config. It has no Python runtime, so it reproduces Python value semantics with `litellm-python-compat`
- `crates/python-bridge` receives live Python objects, such as `RedisCache.redis_kwargs`. It reads them through PyO3 with the real protocol (`is_truthy()`, `str()`, `extract`), so it never imitates Python
- Both produce the same typed input, defined in the data crate the consumer already depends on (`redis-types` for Redis), so a downstream crate cannot tell which host it runs under
## Python value semantics
Reading a value the way Python itself reads it. This lives only at the two boundaries
- Truthiness and `x or fallback`, `str(x)`, `int(text)` and `float(text)`
- Each surface's true and false text: `_str_to_bool`, pydantic's lax bool, `get_secret`'s `true` and `false`
- `os.environ/` references and `REDIS_<KWARG>` style environment layering
A boundary knows which surface a value came from. `general_settings.coordination_redis` is a pydantic model and `cache_params` goes through `_str_to_bool`, so the boundary applies that surface's rule exactly instead of refusing the text two surfaces disagree on
Each typed field keeps where its value came from (the key, `REDIS_<KWARG>` or the fallback name), so an error or warning downstream still names it
## What stays downstream
A crate that stands in for a Python library keeps that library's semantics, because they are its compatibility contract and not value coercion. `redis-resolve` stands in for redis-py, so it keeps
- Which target wins, and the defaults redis-py applies
- How `parse_url` reads a url, including `int()` on `?db=` and redis-py's `to_bool` on the TLS query parameters
- Which kwargs redis-py accepts, and whether the Rust client reads, warns about or refuses each one
Its input is typed. A rule it applies reads typed fields, never a `serde_json::Value` it has to coerce first. A value it only carries back to Python untouched, like `passthrough`, may stay untyped
## litellm-python-compat
The rule is about where config and kwargs get coerced, not about who may depend on the crate
- Its data formats (`pickle`, `json.dumps`, `repr`) are fine anywhere Rust reads or writes state Python also reads or writes, such as cache values in a shared Redis
- Its coercions (`truthy`, `number`) belong at the two boundaries, and in a crate that stands in for a Python library when that library calls them itself
## Existing code
Some code still coerces Python-shaped values past the boundary, for example the `REDIS_*` merge in `gateway/src/redis/layers.rs` and the flag and number readers in `redis-resolve/src/options.rs`. Do not add more. When you change one, move the coercion to its boundary instead of extending it in place
## Tests
Pin a boundary's coercion with named cases next to the Python function it mirrors, in the Rust test and in its Python twin, like the Redis contract cases in `redis-resolve/tests/interpret.rs` and `tests/unit/test_redis.py`. A downstream crate's tests build its typed input directly and never feed it raw text to coerce

4
litellm-rust/AGENTS.md
View file

@ -8,6 +8,10 @@ For string-valued enums and their Serde conversions, follow [.agents/skills/rust
For error types, error variants and conversions between crate errors, follow [.agents/skills/rust-error/SKILL.md](.agents/skills/rust-error/SKILL.md)
For where config.yaml values, environment layering and kwargs from Python become typed Rust, follow [.agents/skills/rust-python-boundary/SKILL.md](.agents/skills/rust-python-boundary/SKILL.md)
The process-wide rustls `CryptoProvider` (ring) is installed once by the host. rust-proxy does it at the top of `main` through `litellm_gateway::install_crypto_provider`, and python-bridge must do the same in its module init before it builds any rustls config. Library crates never install it. The AWS SDK links aws-lc-rs next to ring, so without an installed default the first rustls config a library builds, such as redis-rs TLS, panics
For fieldless enums, derive `strum::VariantArray` and use `VARIANTS` instead of a hand-listed `ALL` array. Derive Strum string conversions instead of hand-written variant-to-string matches
Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `str::parse`) with no `as_str`/`parse` wrapper that only delegates, and spell each variant with explicit `#[strum(serialize = "...")]` instead of `serialize_all`

270
litellm-rust/Cargo.lock generated
View file

@ -8,6 +8,17 @@ version = "2.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa"
[[package]]
name = "aes"
version = "0.8.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0"
dependencies = [
"cfg-if",
"cipher",
"cpufeatures 0.2.17",
]
[[package]]
name = "ahash"
version = "0.8.12"
@ -547,7 +558,7 @@ dependencies = [
"http 1.4.2",
"http-body 0.4.6",
"hyper 0.14.32",
"hyper 1.10.1",
"hyper 1.12.0",
"hyper-rustls 0.24.2",
"hyper-rustls 0.27.9",
"hyper-util",
@ -736,7 +747,7 @@ dependencies = [
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"hyper 1.10.1",
"hyper 1.12.0",
"hyper-util",
"itoa",
"matchit",
@ -982,6 +993,15 @@ dependencies = [
"hybrid-array",
]
[[package]]
name = "block-padding"
version = "0.3.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a8894febbff9f758034a5b8e12d87918f56dfc64a8e1fe757d65e29041538d93"
dependencies = [
"generic-array",
]
[[package]]
name = "bollard"
version = "0.20.2"
@ -1000,7 +1020,7 @@ dependencies = [
"home",
"http 1.4.2",
"http-body-util",
"hyper 1.10.1",
"hyper 1.12.0",
"hyper-named-pipe",
"hyper-rustls 0.27.9",
"hyper-util",
@ -1134,6 +1154,15 @@ dependencies = [
"rustversion",
]
[[package]]
name = "cbc"
version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "26b52a9543ae338f279b96b0b9fed9c8093744685043739079ce85cd58f289a6"
dependencies = [
"cipher",
]
[[package]]
name = "cc"
version = "1.3.0"
@ -1210,6 +1239,16 @@ dependencies = [
"half",
]
[[package]]
name = "cipher"
version = "0.4.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad"
dependencies = [
"crypto-common 0.1.7",
"inout",
]
[[package]]
name = "clap"
version = "4.6.7"
@ -1774,6 +1813,15 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "des"
version = "0.8.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ffdd80ce8ce993de27e9f063a444a4d53ce8e8db4c1f00cc03af5ad5a9867a1e"
dependencies = [
"cipher",
]
[[package]]
name = "digest"
version = "0.10.7"
@ -2255,7 +2303,7 @@ dependencies = [
"chrono",
"http 1.4.2",
"http-body-util",
"hyper 1.10.1",
"hyper 1.12.0",
"hyper-rustls 0.27.9",
"hyper-util",
"ring",
@ -2391,7 +2439,7 @@ dependencies = [
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"hyper 1.10.1",
"hyper 1.12.0",
"lazy_static",
"opentelemetry",
"opentelemetry-semantic-conventions",
@ -2833,9 +2881,9 @@ dependencies = [
[[package]]
name = "hyper"
version = "1.10.1"
version = "1.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498"
checksum = "2c3e324da4c95177d6291d4c8730197c0d1822f8a9766814a4a44fa5ab797c9c"
dependencies = [
"atomic-waker",
"bytes",
@ -2860,7 +2908,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fab3637d6b04a8037af8a266fdf6cf92ea957e8c53981a2bf6136572531025bf"
dependencies = [
"hex",
"hyper 1.10.1",
"hyper 1.12.0",
"hyper-util",
"pin-project-lite",
"tokio",
@ -2889,7 +2937,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f"
dependencies = [
"http 1.4.2",
"hyper 1.10.1",
"hyper 1.12.0",
"hyper-util",
"rustls 0.23.42",
"rustls-native-certs",
@ -2904,7 +2952,7 @@ version = "0.5.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0"
dependencies = [
"hyper 1.10.1",
"hyper 1.12.0",
"hyper-util",
"pin-project-lite",
"tokio",
@ -2923,7 +2971,7 @@ dependencies = [
"futures-util",
"http 1.4.2",
"http-body 1.1.0",
"hyper 1.10.1",
"hyper 1.12.0",
"ipnet",
"libc",
"percent-encoding",
@ -2942,7 +2990,7 @@ checksum = "986c5ce3b994526b3cd75578e62554abd09f0899d6206de48b3e96ab34ccc8c7"
dependencies = [
"hex",
"http-body-util",
"hyper 1.10.1",
"hyper 1.12.0",
"hyper-util",
"pin-project-lite",
"tokio",
@ -3105,6 +3153,16 @@ dependencies = [
"serde_core",
]
[[package]]
name = "inout"
version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01"
dependencies = [
"block-padding",
"generic-array",
]
[[package]]
name = "ipnet"
version = "2.12.0"
@ -3663,8 +3721,7 @@ dependencies = [
"strum",
"tempfile",
"thiserror 2.0.19",
"tracing",
"tracing-subscriber",
"veil",
]
[[package]]
@ -3767,8 +3824,10 @@ dependencies = [
"http-body-util",
"litellm-circuit-breaker",
"litellm-config",
"litellm-core-utils",
"litellm-error",
"litellm-gateway-auth",
"litellm-gateway-forward",
"litellm-gateway-inference",
"litellm-gateway-mcp",
"litellm-gateway-ui",
@ -3776,11 +3835,15 @@ dependencies = [
"litellm-inference",
"litellm-llms",
"litellm-redis-client",
"litellm-redis-connect",
"litellm-redis-resolve",
"litellm-redis-testing",
"litellm-redis-types",
"litellm-secret-value",
"litellm-secrets",
"litellm-tracing",
"rstest",
"rustls 0.23.42",
"serde",
"serde_json",
"tempfile",
@ -3791,6 +3854,7 @@ dependencies = [
"tower-sessions-moka-store",
"tracing",
"uuid",
"veil",
]
[[package]]
@ -3815,6 +3879,29 @@ dependencies = [
"veil",
]
[[package]]
name = "litellm-gateway-forward"
version = "0.1.0"
dependencies = [
"axum",
"bytes",
"envy",
"futures-util",
"http-body-util",
"httparse",
"hyper 1.12.0",
"hyper-util",
"proptest",
"rstest",
"serde",
"thiserror 2.0.19",
"tokio",
"tokio-tungstenite",
"tokio-util",
"tower-service",
"tracing",
]
[[package]]
name = "litellm-gateway-inference"
version = "0.1.0"
@ -4294,11 +4381,13 @@ dependencies = [
"criterion",
"fancy-regex 0.19.2",
"futures-util",
"indexmap 2.14.0",
"litellm-auth",
"litellm-auth-aws",
"litellm-cache",
"litellm-cache-response",
"litellm-callbacks-legacy-python",
"litellm-circuit-breaker",
"litellm-core-utils",
"litellm-error",
"litellm-host",
@ -4313,6 +4402,12 @@ dependencies = [
"litellm-llms",
"litellm-llms-types",
"litellm-model-catalog",
"litellm-python-compat",
"litellm-redis-client",
"litellm-redis-connect",
"litellm-redis-resolve",
"litellm-redis-store",
"litellm-redis-types",
"litellm-router-core",
"litellm-router-types",
"litellm-secret-value",
@ -4363,23 +4458,110 @@ version = "0.1.0"
dependencies = [
"futures-util",
"litellm-circuit-breaker",
"litellm-python-compat",
"litellm-redis-testing",
"litellm-redis-types",
"litellm-secret-value",
"pem",
"pkcs8",
"redis",
"rstest",
"rustls 0.23.42",
"rustls-native-certs",
"socket2 0.6.5",
"thiserror 2.0.19",
"tokio",
"tracing",
]
[[package]]
name = "litellm-redis-connect"
version = "0.1.0"
dependencies = [
"futures-util",
"litellm-auth-aws",
"litellm-auth-azure",
"litellm-auth-gcp",
"litellm-auth-types",
"litellm-circuit-breaker",
"litellm-core-utils",
"litellm-error",
"litellm-http",
"litellm-redis-client",
"litellm-redis-resolve",
"litellm-redis-testing",
"litellm-redis-types",
"litellm-secret-value",
"rstest",
"rustls 0.23.42",
"serde",
"serde_json",
"strum",
"thiserror 2.0.19",
"time",
"tokio",
"tracing",
"url",
"veil",
]
[[package]]
name = "litellm-redis-resolve"
version = "0.1.0"
dependencies = [
"litellm-python-compat",
"litellm-redis-types",
"litellm-secret-value",
"rstest",
"serde",
"serde_json",
"strum",
"thiserror 2.0.19",
"tracing",
"url",
]
[[package]]
name = "litellm-redis-store"
version = "0.1.0"
dependencies = [
"futures-util",
"litellm-cache",
"litellm-cache-memory",
"litellm-circuit-breaker",
"litellm-python-compat",
"litellm-redis-client",
"litellm-redis-testing",
"litellm-secret-value",
"litellm-tracing",
"rstest",
"strum",
"thiserror 2.0.19",
"tokio",
"tracing",
]
[[package]]
name = "litellm-redis-testing"
version = "0.1.0"
dependencies = [
"pkcs8",
"rcgen",
"tempfile",
"testcontainers-modules",
"thiserror 2.0.19",
"tokio",
]
[[package]]
name = "litellm-redis-types"
version = "0.1.0"
dependencies = [
"litellm-secret-value",
"rstest",
"url",
]
[[package]]
name = "litellm-router"
version = "0.1.0"
@ -5262,6 +5444,16 @@ version = "0.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4"
[[package]]
name = "pbkdf2"
version = "0.12.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f8ed6a7761f76e3b9f92dfb0a60a6a6477c61024b775147ff0973a02653abaf2"
dependencies = [
"digest 0.10.7",
"hmac 0.12.1",
]
[[package]]
name = "pem"
version = "4.0.0"
@ -5372,6 +5564,22 @@ dependencies = [
"spki",
]
[[package]]
name = "pkcs5"
version = "0.7.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e847e2c91a18bfa887dd028ec33f2fe6f25db77db3619024764914affe8b69a6"
dependencies = [
"aes",
"cbc",
"der",
"des",
"pbkdf2",
"scrypt",
"sha2 0.10.9",
"spki",
]
[[package]]
name = "pkcs8"
version = "0.10.2"
@ -5379,6 +5587,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7"
dependencies = [
"der",
"pkcs5",
"rand_core 0.6.4",
"spki",
]
@ -5566,6 +5776,7 @@ version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b"
dependencies = [
"indexmap 2.14.0",
"libc",
"once_cell",
"portable-atomic",
@ -6031,7 +6242,7 @@ dependencies = [
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"hyper 1.10.1",
"hyper 1.12.0",
"hyper-rustls 0.27.9",
"hyper-util",
"js-sys",
@ -6387,6 +6598,15 @@ version = "1.0.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f"
[[package]]
name = "salsa20"
version = "0.10.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "97a22f5af31f73a954c10289c93e8a50cc23d971e80ee446f1f6f7137a088213"
dependencies = [
"cipher",
]
[[package]]
name = "same-file"
version = "1.0.6"
@ -6450,6 +6670,17 @@ version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
[[package]]
name = "scrypt"
version = "0.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0516a385866c09368f0b5bcd1caff3366aace790fcd46e2bb032697bb172fd1f"
dependencies = [
"pbkdf2",
"salsa20",
"sha2 0.10.9",
]
[[package]]
name = "sct"
version = "0.7.1"
@ -7512,6 +7743,7 @@ dependencies = [
"bytes",
"futures-core",
"futures-sink",
"futures-util",
"pin-project-lite",
"tokio",
]
@ -7591,7 +7823,7 @@ dependencies = [
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"hyper 1.10.1",
"hyper 1.12.0",
"hyper-timeout",
"hyper-util",
"percent-encoding",
@ -8167,9 +8399,9 @@ dependencies = [
[[package]]
name = "want"
version = "0.3.1"
version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e"
checksum = "ec4cdd0dd910afe868b7ef477227d8d538b46b3075031afee8a9f2acb0a2ed0b"
dependencies = [
"try-lock",
]
@ -8551,7 +8783,7 @@ dependencies = [
"futures",
"http 1.4.2",
"http-body-util",
"hyper 1.10.1",
"hyper 1.12.0",
"hyper-util",
"log",
"once_cell",

5
litellm-rust/Cargo.toml
View file

@ -26,6 +26,7 @@ litellm-inference-ocr = { path = "crates/inference-ocr" }
litellm-inference-testing = { path = "crates/inference-testing" }
litellm-gateway-mcp = { path = "crates/gateway-mcp" }
litellm-gateway = { path = "crates/gateway" }
litellm-gateway-forward = { path = "crates/gateway-forward" }
litellm-gateway-inference = { path = "crates/gateway-inference" }
litellm-gateway-auth = { path = "crates/gateway-auth" }
litellm-gateway-management = { path = "crates/gateway-management" }
@ -72,7 +73,11 @@ litellm-host-python = { path = "crates/host-python" }
litellm-python-compat = { path = "crates/python-compat" }
litellm-circuit-breaker = { path = "crates/circuit-breaker" }
litellm-redis-client = { path = "crates/redis-client" }
litellm-redis-connect = { path = "crates/redis-connect" }
litellm-redis-store = { path = "crates/redis-store" }
litellm-redis-testing = { path = "crates/redis-testing" }
litellm-redis-resolve = { path = "crates/redis-resolve" }
litellm-redis-types = { path = "crates/redis-types" }
litellm-cost = { path = "crates/cost" }
litellm-model-catalog = { path = "crates/model-catalog" }
litellm-model-catalog-refresh = { path = "crates/model-catalog-refresh" }

53
litellm-rust/crates/cache-memory/src/cache.rs
View file

@ -2,7 +2,10 @@ use std::{
cmp::Reverse,
collections::{BinaryHeap, HashMap, HashSet},
hash::Hash,
sync::{Arc, Mutex, MutexGuard},
sync::{
Arc, Mutex, MutexGuard,
atomic::{AtomicUsize, Ordering},
},
time::{Duration, SystemTime, UNIX_EPOCH},
};
@ -31,7 +34,7 @@ struct CacheState<V> {
pub struct InMemoryCache<V: Clone> {
state: Mutex<CacheState<V>>,
max_size_in_memory: usize,
max_size_in_memory: AtomicUsize,
default_ttl: Duration,
max_entry_bytes: Option<usize>,
measure_value: Option<ValueMeasure<V>>,
@ -74,7 +77,9 @@ impl<V: Clone> InMemoryCache<V> {
expirations: HashMap::new(),
expiration_heap: BinaryHeap::new(),
}),
max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
max_size_in_memory: AtomicUsize::new(
max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
),
default_ttl: default_ttl
.filter(|ttl| !ttl.is_zero())
.unwrap_or(DEFAULT_TTL),
@ -94,7 +99,7 @@ impl<V: Clone> InMemoryCache<V> {
value: V,
ttl: Option<Duration>,
) -> Result<CacheWrite, Error> {
if self.max_size_in_memory == 0 {
if self.max_size_in_memory() == 0 {
return Ok(CacheWrite::Disabled);
}
let now = (self.now)();
@ -129,7 +134,7 @@ impl<V: Clone> InMemoryCache<V> {
pub fn evict_cache(&self) -> Result<(), Error> {
let now = (self.now)();
let mut state = self.lock()?;
Self::evict(&mut state, self.max_size_in_memory, now, None);
Self::evict(&mut state, self.max_size_in_memory(), now, None);
Ok(())
}
@ -170,7 +175,16 @@ impl<V: Clone> InMemoryCache<V> {
}
pub fn max_size_in_memory(&self) -> usize {
self.max_size_in_memory
self.max_size_in_memory.load(Ordering::Relaxed)
}
/// `DualCache.update_in_memory_max_size`: `None` restores the default. Entries already held stay
/// until the next write evicts down to the new size.
pub fn update_max_size_in_memory(&self, max_size_in_memory: Option<usize>) {
self.max_size_in_memory.store(
max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
Ordering::Relaxed,
);
}
pub fn max_entry_bytes(&self) -> Option<usize> {
@ -200,6 +214,23 @@ impl<V: Clone> InMemoryCache<V> {
.collect())
}
/// Every live entry, for callers that inspect the whole tier.
pub fn entries(&self) -> Result<Vec<(String, V)>, Error> {
let now = (self.now)();
let state = self.lock()?;
Ok(state
.values
.iter()
.filter(|(key, _)| {
state
.expirations
.get(*key)
.is_none_or(|expiration| *expiration >= now)
})
.map(|(key, value)| (key.clone(), value.clone()))
.collect())
}
pub fn delete_cache(&self, key: &str) -> Result<(), Error> {
let mut state = self.lock()?;
Self::remove(&mut state, key);
@ -276,7 +307,7 @@ impl<V: Clone> InMemoryCache<V> {
ttl: Option<Duration>,
now: Duration,
) -> Result<CacheWrite, Error> {
Self::evict(state, self.max_size_in_memory, now, None);
Self::evict(state, self.max_size_in_memory(), now, None);
if !self.check_value_size(&value)? {
return Ok(CacheWrite::TooLarge);
}
@ -300,12 +331,12 @@ where
eligible: &[V],
context: ExactCacheContext,
) -> Result<V, Error> {
if self.max_size_in_memory == 0 {
if self.max_size_in_memory() == 0 {
return Ok(candidate);
}
let now = (self.now)();
let mut state = self.lock()?;
Self::evict(&mut state, self.max_size_in_memory, now, Some(key));
Self::evict(&mut state, self.max_size_in_memory(), now, Some(key));
let existing = state
.values
.get(key)
@ -335,7 +366,7 @@ impl CounterCache for InMemoryCache<f64> {
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
if self.max_size_in_memory == 0 {
if self.max_size_in_memory() == 0 {
return Ok(amount);
}
let now = (self.now)();
@ -408,7 +439,7 @@ where
values: Vec<Self::SetValue>,
ttl: Option<Duration>,
) -> Result<Self::SetResult, Error> {
if self.max_size_in_memory == 0 {
if self.max_size_in_memory() == 0 {
return Ok(values);
}
let now = (self.now)();

9
litellm-rust/crates/circuit-breaker/AGENTS.md
View file

@ -1,13 +1,14 @@
- `fsm` is a generic state machine runtime and knows nothing about breakers. `breaker` is a generic circuit breaker and knows nothing about Redis. Deciding which errors count as failures belongs to the consumer
- `Machine::step` stays pure. `&self` is config and time comes in on the event, never read inside the machine. The only clock read is `Admission`'s drop, at the call-site boundary
- Transition rules are one `match (state, event)` per machine, with no macro DSL and no `_` arm that would hide a new state from the compiler
- `Shared` counts a transition, bumps the generation and calls the observer only when the state's enum variant changes, matching Python bumping `generation` only when `_state` changes. Data inside a variant (failure counters, `since`) changes without a transition
- `Shared` counts a transition, bumps the generation and calls the observer only when the state's enum variant changes. Data inside a variant (failure counters, `since`) changes without a transition
- The observer runs under the lock so transitions are seen in generation order. Keep it to logging or a metric
- A disabled breaker is `Option<Breaker>::None` at the consumer, not a flag inside the machine. `Config` comes in already resolved, reading `REDIS_CIRCUIT_BREAKER_*` env is the consumer's job
## Parity with Python `RedisCircuitBreaker` (`litellm/caching/redis_cache.py`)
- The rstest table `matches_the_python_redis_circuit_breaker` was checked case by case against the Python breaker. A behavior change here needs the same check, by driving the Python class with the same steps
- Divergences, each on purpose
## The only Redis breaker
- Python's `RedisCache` admits through this breaker via `NativeRedisBreaker` (`python-bridge/src/redis/breaker.rs`), sharing one per Redis server with the native client through `redis-connect`'s `RedisHosts`, so the Python `RedisCircuitBreaker` it replaced is gone
- The rstest table `matches_the_python_redis_circuit_breaker` was checked case by case against that Python class and now pins its behavior. A behavior change here changes what every Redis caller in the proxy sees
- Divergences from the old Python class, each on purpose
- monotonic `Instant` instead of `time.time()`, so a wall clock jump can't open or close the breaker
- a probe `Admission` dropped without an outcome records a hard failure and reopens. Python's probe that dies on `CancelledError` (a `BaseException`) records nothing and leaves the breaker HALF_OPEN for good, and a dropped Rust future is common
- `_failure_count` is dropped, it only fed a log line

15
litellm-rust/crates/circuit-breaker/src/breaker.rs
View file

@ -1,5 +1,5 @@
use std::{
sync::Arc,
sync::{Arc, Weak},
time::{Duration, Instant},
};
@ -128,6 +128,9 @@ impl Machine for Config {
#[derive(Clone)]
pub struct Breaker(Arc<Shared<Config>>);
#[derive(Clone)]
pub struct WeakBreaker(Weak<Shared<Config>>);
#[derive(Debug, PartialEq, Eq)]
pub struct Rejected;
@ -165,6 +168,16 @@ impl Breaker {
pub fn state(&self) -> State {
self.0.state()
}
pub fn downgrade(&self) -> WeakBreaker {
WeakBreaker(Arc::downgrade(&self.0))
}
}
impl WeakBreaker {
pub fn upgrade(&self) -> Option<Breaker> {
self.0.upgrade().map(Breaker)
}
}
impl Admission {

4
litellm-rust/crates/circuit-breaker/src/lib.rs
View file

@ -1,5 +1,7 @@
mod breaker;
mod fsm;
pub use breaker::{Admission, Breaker, Config, Event, Failure, Output, Rejected, State};
pub use breaker::{
Admission, Breaker, Config, Event, Failure, Output, Rejected, State, WeakBreaker,
};
pub use fsm::{Machine, Observer, Shared, Transition};

3
litellm-rust/crates/config/AGENTS.md Normal file
View file

@ -0,0 +1,3 @@
## Redis settings
- This crate only deserializes YAML. `RedisSettings` is the raw `coordination_redis` or `cache_params` map and `CacheSettings` tells a `type: redis` cache apart from the rest. Reading those maps the way `litellm/_redis.py` does belongs to `gateway/src/redis`, and turning merged kwargs into a `Resolution` belongs to `redis-resolve`

7
litellm-rust/crates/config/Cargo.toml
View file

@ -6,17 +6,16 @@ license.workspace = true
repository.workspace = true
[dependencies]
litellm-secret-value.workspace = true
litellm-router-types.workspace = true
litellm-secret-value.workspace = true
serde.workspace = true
serde_json.workspace = true
serde_yaml_ng = "0.10.0"
strum.workspace = true
thiserror.workspace = true
tracing.workspace = true
veil.workspace = true
[dev-dependencies]
litellm-error.workspace = true
rstest.workspace = true
serde_json.workspace = true
tempfile.workspace = true
tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }

7
litellm-rust/crates/config/src/error.rs
View file

@ -11,11 +11,4 @@ pub enum Error {
path: PathBuf,
source: serde_yaml_ng::Error,
},
#[error("Redis host does not yet support configured setting {0}")]
RedisSetting(String),
#[error("invalid Redis setting {setting}: {reason}")]
RedisInvalid {
setting: &'static str,
reason: &'static str,
},
}

5
litellm-rust/crates/config/src/lib.rs
View file

@ -14,10 +14,7 @@ pub use error::Error;
pub use litellm_router_types::{LitellmParams, Spelled};
pub use mcp::{McpAuth, McpServer, McpTransport};
pub use model::Model;
pub use redis::{
CacheSettings, Lookup, RedisCluster, RedisEndpoint, RedisHost, RedisNode, RedisSettings,
ResolvedRedis, environment_has_target,
};
pub use redis::{CacheSettings, RedisSettings};
pub use settings::{
ClickHouseStoreSettings, GeneralSettings, LiteLlmSettings, RouterSettings, TracingSettings,
TracingStoreSettings,

658
litellm-rust/crates/config/src/redis.rs
View file

@ -1,195 +1,21 @@
use std::{collections::BTreeMap, fmt, time::Duration};
use std::collections::BTreeMap;
use litellm_router_types::Spelled;
use litellm_secret_value::SecretValue;
use serde::{Deserialize, Deserializer, de::Error as _};
use veil::Redact;
use crate::{AdditionalFields, Error, Object, Value};
use crate::{Object, Value};
pub type Lookup<'a> = &'a dyn Fn(&str) -> Option<String>;
const SUPPORTED: &[&str] = &[
"url",
"host",
"port",
"username",
"password",
"ssl",
"db",
"socket_timeout",
"startup_nodes",
"namespace",
];
const IGNORED: &[&str] = &[
"max_connections",
"single_connection_client",
"health_check_interval",
"socket_keepalive",
"socket_keepalive_options",
"decode_responses",
"encoding",
"encoding_errors",
"charset",
"errors",
"client_name",
"lib_name",
"lib_version",
"sentinel_password",
"gcp_ssl_ca_certs",
"azure_client_id",
"azure_tenant_id",
"azure_client_secret",
"aws_iam_user_name",
"aws_iam_cache_name",
"aws_iam_region",
"aws_iam_serverless",
];
const WARN_IGNORED: &[&str] = &[
"max_connections",
"health_check_interval",
"socket_keepalive",
"socket_keepalive_options",
];
const REDIS_PY_KWARGS: &[&str] = &[
"aws_iam_auth",
"aws_iam_cache_name",
"aws_iam_region",
"aws_iam_serverless",
"aws_iam_user_name",
"azure_client_id",
"azure_client_secret",
"azure_redis_ad_token",
"azure_tenant_id",
"cache",
"cache_config",
"charset",
"client_name",
"db",
"decode_responses",
"encoding",
"encoding_errors",
"errors",
"event_dispatcher",
"gcp_service_account",
"gcp_ssl_ca_certs",
"health_check_interval",
"host",
"lib_name",
"lib_version",
"max_connections",
"password",
"port",
"protocol",
"redis_connect_func",
"retry_on_error",
"retry_on_timeout",
"single_connection_client",
"socket_connect_timeout",
"socket_keepalive",
"socket_keepalive_options",
"socket_timeout",
"ssl",
"ssl_ca_certs",
"ssl_ca_data",
"ssl_ca_path",
"ssl_cert_reqs",
"ssl_certfile",
"ssl_check_hostname",
"ssl_ciphers",
"ssl_keyfile",
"ssl_min_version",
"ssl_ocsp_context",
"ssl_ocsp_expected_cert",
"ssl_password",
"ssl_validate_ocsp",
"ssl_validate_ocsp_stapled",
"unix_socket_path",
"url",
"username",
];
const UNSUPPORTED_EXTRA: &[&str] = &[
"credential_provider",
"sentinel_nodes",
"retry",
"connection_pool",
];
const CACHE_ONLY: &[&str] = &[
"type",
"mode",
"ttl",
"default_in_memory_ttl",
"default_in_redis_ttl",
"supported_call_types",
"similarity_threshold",
"max_messages",
"_backend",
];
#[derive(Clone, Default)]
/// A Redis section as `config.yaml` spells it, before any surface reads it.
#[derive(Clone, Default, Deserialize, Redact)]
#[serde(transparent)]
pub struct RedisSettings {
pub url: Option<SecretValue>,
pub host: Option<Spelled<String>>,
pub port: Option<Spelled<u16>>,
pub username: Option<SecretValue>,
pub password: Option<SecretValue>,
pub ssl: Option<Spelled<bool>>,
pub db: Option<Spelled<i64>>,
pub socket_timeout: Option<Spelled<f64>>,
pub startup_nodes: Option<Value>,
pub namespace: Option<Spelled<String>>,
pub additional_fields: AdditionalFields,
unparsed: DeferredInputs,
#[redact(with = "[REDACTED]")]
fields: BTreeMap<String, Value>,
}
impl<'de> Deserialize<'de> for RedisSettings {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let fields = BTreeMap::<String, Value>::deserialize(deserializer)?;
let scalar = |key: &str| fields.get(key).and_then(as_text);
Ok(Self {
url: scalar("url").map(SecretValue::new),
host: scalar("host").map(Spelled::Value),
port: spelled(scalar("port")),
username: scalar("username").map(SecretValue::new),
password: scalar("password").map(SecretValue::new),
ssl: spelled(scalar("ssl")),
db: spelled(scalar("db")),
socket_timeout: spelled(scalar("socket_timeout")),
startup_nodes: fields
.get("startup_nodes")
.filter(|value| !value.is_null())
.cloned(),
namespace: scalar("namespace").map(Spelled::Value),
additional_fields: cloned(&fields, |key, _| !SUPPORTED.contains(&key)),
unparsed: deferred_inputs(&fields),
})
}
}
impl fmt::Debug for RedisSettings {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("RedisSettings")
.field("url", &self.url)
.field("host", &self.host)
.field("port", &self.port)
.field("username", &self.username)
.field("password", &self.password)
.field("ssl", &self.ssl)
.field("db", &self.db)
.field("socket_timeout", &self.socket_timeout)
.field(
"startup_nodes",
&self.startup_nodes.as_ref().map(|_| "<configured>"),
)
.field("namespace", &self.namespace)
.field("additional_fields", &self.additional_fields.keys())
.field("unparsed", &self.unparsed.keys())
.finish()
impl RedisSettings {
pub fn fields(&self) -> &BTreeMap<String, Value> {
&self.fields
}
}
@ -216,452 +42,22 @@ impl<'de> Deserialize<'de> for CacheSettings {
}
}
#[derive(Clone, PartialEq)]
pub struct ResolvedRedis {
pub endpoint: RedisEndpoint,
pub namespace: Option<String>,
pub timeout: Option<Duration>,
}
#[cfg(test)]
mod tests {
use super::RedisSettings;
use rstest::rstest;
impl fmt::Debug for ResolvedRedis {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("ResolvedRedis")
.field("endpoint", &self.endpoint)
.field("namespace", &self.namespace)
.field("timeout", &self.timeout)
.finish()
#[rstest]
#[case::scalar(serde_json::json!("credential-value"))]
#[case::object(serde_json::json!({"secret": "credential-value"}))]
#[case::array(serde_json::json!(["credential-value"]))]
fn credentials_are_redacted(#[case] password: serde_json::Value) {
let settings: RedisSettings = serde_json::from_value(serde_json::json!({
"host": "redis", "password": password,
}))
.unwrap();
let debug = format!("{settings:?}");
assert!(!debug.contains("credential-value"));
assert!(debug.contains("[REDACTED]"));
}
}
#[derive(Clone, PartialEq)]
pub enum RedisEndpoint {
Url(SecretValue),
Host(RedisHost),
Cluster(RedisCluster),
}
impl fmt::Debug for RedisEndpoint {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Url(_) => formatter.debug_tuple("Url").field(&"<redacted>").finish(),
Self::Host(host) => host.fmt(formatter),
Self::Cluster(cluster) => cluster.fmt(formatter),
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct RedisHost {
pub host: String,
pub port: u16,
pub username: Option<SecretValue>,
pub password: Option<SecretValue>,
pub tls: bool,
pub db: i64,
}
#[derive(Clone, Debug, PartialEq)]
pub struct RedisCluster {
pub nodes: Vec<RedisNode>,
pub username: Option<SecretValue>,
pub password: Option<SecretValue>,
pub tls: bool,
}
#[derive(Clone, Debug, PartialEq)]
pub struct RedisNode {
pub host: String,
pub port: u16,
}
#[derive(Clone)]
struct Entry {
value: Value,
source: String,
}
type Kwargs = BTreeMap<String, Entry>;
impl RedisSettings {
pub fn resolve(&self, lookup: Lookup<'_>) -> Result<ResolvedRedis, Error> {
self.resolve_with(lookup, false)
}
pub fn resolve_cache(&self, lookup: Lookup<'_>) -> Result<ResolvedRedis, Error> {
self.resolve_with(lookup, true)
}
fn resolve_with(&self, lookup: Lookup<'_>, cache: bool) -> Result<ResolvedRedis, Error> {
let block = self.kwargs(cache);
let block: Kwargs = block
.into_iter()
.filter_map(|(key, entry)| {
let value = environment_reference(&entry.value, lookup)?;
Some((key, Entry { value, ..entry }))
})
.collect();
let environment = environment_kwargs(lookup);
let names_a_target = ["host", "startup_nodes", "sentinel_nodes"]
.iter()
.any(|key| block.contains_key(*key));
let drop_environment_url = names_a_target && !block.contains_key("url");
let merged: Kwargs = environment
.into_iter()
.filter(|(key, _)| !(key == "url" && drop_environment_url))
.chain(block)
.collect();
reject_unsupported(&merged, lookup)?;
warn_ignored(&merged);
let startup_nodes = match merged.get("startup_nodes") {
Some(entry) => Some(cluster_nodes(&entry.value)?),
None => lookup("REDIS_CLUSTER_NODES")
.map(|nodes| cluster_nodes(&Value::String(nodes)))
.transpose()?,
};
let username = secret_value(&merged, "username");
let password = secret_value(&merged, "password");
let tls = bool_value(&merged, "ssl");
let endpoint = match (startup_nodes, text_value(&merged, "url")) {
(Some(nodes), _) if !nodes.is_empty() => RedisEndpoint::Cluster(RedisCluster {
nodes,
username,
password,
tls,
}),
(_, Some(url)) => RedisEndpoint::Url(SecretValue::new(url)),
_ => RedisEndpoint::Host(RedisHost {
host: text_value(&merged, "host").ok_or(Error::RedisInvalid {
setting: "host",
reason: "either host or url must be specified for redis",
})?,
port: integer(&merged, "port")?.unwrap_or(6379),
username,
password,
tls,
db: integer(&merged, "db")?.unwrap_or(0),
}),
};
let timeout = merged
.get("socket_timeout")
.and_then(|entry| as_text(&entry.value))
.map(|seconds| {
seconds
.parse::<f64>()
.ok()
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
.ok_or(Error::RedisInvalid {
setting: "socket_timeout",
reason: "must be a non-negative number of seconds",
})
})
.transpose()?;
Ok(ResolvedRedis {
endpoint,
namespace: text_value(&merged, "namespace"),
timeout,
})
}
fn kwargs(&self, cache: bool) -> Kwargs {
let known = [
(
"url",
self.url.as_ref().map(|value| value.expose().to_owned()),
),
("host", spelled_text(self.host.as_ref())),
("port", spelled_text(self.port.as_ref())),
(
"username",
self.username
.as_ref()
.map(|value| value.expose().to_owned()),
),
(
"password",
self.password
.as_ref()
.map(|value| value.expose().to_owned()),
),
("ssl", spelled_text(self.ssl.as_ref())),
("db", spelled_text(self.db.as_ref())),
("socket_timeout", spelled_text(self.socket_timeout.as_ref())),
("namespace", spelled_text(self.namespace.as_ref())),
];
let known = known.into_iter().filter_map(|(key, value)| {
value.map(|value| {
(
key.to_owned(),
Entry {
value: Value::String(value),
source: key.to_owned(),
},
)
})
});
let startup_nodes = self.startup_nodes.clone().map(|value| {
(
"startup_nodes".to_owned(),
Entry {
value,
source: "startup_nodes".to_owned(),
},
)
});
let extras = self.additional_fields.iter().filter_map(|(key, value)| {
if cache && CACHE_ONLY.contains(&key.as_str()) {
return None;
}
let key = match (cache, key.as_str()) {
(true, "redis_startup_nodes") => "startup_nodes".to_owned(),
_ => key.clone(),
};
Some((
key.clone(),
Entry {
value: value.clone(),
source: key,
},
))
});
let unparsed = self.unparsed.iter().map(|(key, value)| {
(
key.clone(),
Entry {
value: value.0.clone(),
source: key.clone(),
},
)
});
extras
.chain(unparsed)
.chain(known)
.chain(startup_nodes)
.collect()
}
}
pub fn environment_has_target(lookup: Lookup<'_>) -> bool {
[
"REDIS_HOST",
"REDIS_URL",
"REDIS_CLUSTER_NODES",
"REDIS_SENTINEL_NODES",
]
.iter()
.any(|name| lookup(name).is_some())
}
fn reject_unsupported(merged: &Kwargs, lookup: Lookup<'_>) -> Result<(), Error> {
let active_auth = [
merged
.get("gcp_service_account")
.map(|entry| entry.source.as_str())
.or_else(|| lookup("REDIS_GCP_SERVICE_ACCOUNT").map(|_| "REDIS_GCP_SERVICE_ACCOUNT")),
text_value(merged, "azure_redis_ad_token")
.or_else(|| lookup("REDIS_AZURE_AD_TOKEN"))
.filter(|value| value.eq_ignore_ascii_case("true"))
.map(|_| {
merged
.get("azure_redis_ad_token")
.map_or("REDIS_AZURE_AD_TOKEN", |entry| entry.source.as_str())
}),
bool_value(merged, "aws_iam_auth").then(|| merged["aws_iam_auth"].source.as_str()),
merged
.get("sentinel_nodes")
.map(|entry| entry.source.as_str())
.or_else(|| lookup("REDIS_SENTINEL_NODES").map(|_| "REDIS_SENTINEL_NODES")),
];
if let Some(setting) = active_auth.into_iter().flatten().next() {
return Err(Error::RedisSetting(setting.to_owned()));
}
match merged
.iter()
.find(|(key, _)| is_unsupported(key))
.map(|(_, entry)| entry.source.clone())
{
Some(setting) => Err(Error::RedisSetting(setting)),
None => Ok(()),
}
}
fn warn_ignored(merged: &Kwargs) {
for setting in WARN_IGNORED {
if merged.contains_key(*setting) {
tracing::warn!(setting, "Redis setting is accepted but ignored");
}
}
}
fn is_unsupported(key: &str) -> bool {
let known = REDIS_PY_KWARGS.contains(&key) || UNSUPPORTED_EXTRA.contains(&key);
let handled = SUPPORTED.contains(&key)
|| IGNORED.contains(&key)
|| matches!(
key,
"gcp_service_account" | "azure_redis_ad_token" | "aws_iam_auth"
);
known && !handled
}
fn environment_kwargs(lookup: Lookup<'_>) -> Kwargs {
REDIS_PY_KWARGS
.iter()
.filter_map(|key| {
let source = format!("REDIS_{}", key.to_ascii_uppercase());
let value = lookup(&source)?;
Some((
(*key).to_owned(),
Entry {
value: Value::String(value),
source,
},
))
})
.collect()
}
fn environment_reference(value: &Value, lookup: Lookup<'_>) -> Option<Value> {
match value {
Value::String(text) => match text.strip_prefix("os.environ/") {
Some(name) => lookup(name).map(Value::String),
None => Some(value.clone()),
},
Value::Null => None,
_ => Some(value.clone()),
}
}
fn as_text(value: &Value) -> Option<String> {
match value {
Value::String(text) => Some(text.clone()),
Value::Number(number) => Some(number.to_string()),
Value::Bool(flag) => Some(flag.to_string()),
_ => None,
}
}
fn text_value(kwargs: &Kwargs, key: &str) -> Option<String> {
kwargs.get(key).and_then(|entry| as_text(&entry.value))
}
fn secret_value(kwargs: &Kwargs, key: &str) -> Option<SecretValue> {
text_value(kwargs, key).map(SecretValue::new)
}
fn bool_value(kwargs: &Kwargs, key: &str) -> bool {
match kwargs.get(key).map(|entry| &entry.value) {
Some(Value::Bool(flag)) => *flag,
Some(value) => as_text(value)
.is_some_and(|text| matches!(text.to_ascii_lowercase().as_str(), "true" | "1" | "yes")),
None => false,
}
}
fn integer<T: std::str::FromStr>(kwargs: &Kwargs, key: &'static str) -> Result<Option<T>, Error> {
kwargs
.get(key)
.and_then(|entry| as_text(&entry.value))
.map(|text| {
text.parse().map_err(|_| Error::RedisInvalid {
setting: key,
reason: "must be an integer",
})
})
.transpose()
}
const INVALID_NODES: Error = Error::RedisInvalid {
setting: "startup_nodes",
reason: "must be a list of {host, port} objects",
};
fn cluster_nodes(value: &Value) -> Result<Vec<RedisNode>, Error> {
match value {
Value::String(json) => {
let nodes: Vec<serde_json::Map<String, serde_json::Value>> =
serde_json::from_str(json).map_err(|_| INVALID_NODES)?;
nodes
.iter()
.map(|node| {
cluster_node(
node.get("host").and_then(json_text),
node.get("port").and_then(json_text),
)
})
.collect()
}
Value::Sequence(nodes) => nodes
.iter()
.map(|node| {
cluster_node(
node.get("host").and_then(as_text),
node.get("port").and_then(as_text),
)
})
.collect(),
_ => Err(INVALID_NODES),
}
}
fn cluster_node(host: Option<String>, port: Option<String>) -> Result<RedisNode, Error> {
match (host, port.and_then(|port| port.parse().ok())) {
(Some(host), Some(port)) => Ok(RedisNode { host, port }),
_ => Err(INVALID_NODES),
}
}
fn json_text(value: &serde_json::Value) -> Option<String> {
match value {
serde_json::Value::String(text) => Some(text.clone()),
serde_json::Value::Number(number) => Some(number.to_string()),
_ => None,
}
}
fn spelled_text<T: ToString>(value: Option<&Spelled<T>>) -> Option<String> {
value.map(|value| match value {
Spelled::Value(value) => value.to_string(),
Spelled::Text(value) => value.clone(),
})
}
fn spelled<T: std::str::FromStr>(text: Option<String>) -> Option<Spelled<T>> {
text.map(|text| match text.parse() {
Ok(value) => Spelled::Value(value),
Err(_) => Spelled::Text(text),
})
}
fn cloned(
fields: &BTreeMap<String, Value>,
keep: impl Fn(&str, &Value) -> bool,
) -> AdditionalFields {
fields
.iter()
.filter(|(key, value)| keep(key, value))
.map(|(key, value)| (key.clone(), value.clone()))
.collect()
}
#[derive(Clone)]
struct DeferredInput(Value);
type DeferredInputs = BTreeMap<String, DeferredInput>;
fn deferred_inputs(fields: &BTreeMap<String, Value>) -> DeferredInputs {
fields
.iter()
.filter(|(key, value)| is_unparsed(key, value))
.map(|(key, value)| (key.clone(), DeferredInput(value.clone())))
.collect()
}
fn is_unparsed(key: &str, value: &Value) -> bool {
SUPPORTED.contains(&key)
&& match key {
"startup_nodes" => value.is_null(),
_ => as_text(value).is_none(),
}
}

26
litellm-rust/crates/config/tests/config.rs
View file

@ -440,3 +440,29 @@ fn mcp_config_redacts_nested_credentials_and_preserves_policy_for_validation() {
fn rejects_invalid_typed_mcp_settings(#[case] setting: &str) {
assert!(Config::from_yaml(&format!("mcp_servers:\n docs:\n {setting}\n")).is_err());
}
#[rstest]
fn inactive_redis_cache_keeps_malformed_known_inputs_deferred() {
let config = Config::from_yaml(
"litellm_settings:\n cache: false\n cache_params:\n type: redis\n host: []\n port: invalid\n socket_timeout: invalid",
)
.unwrap();
assert!(matches!(
config.litellm_settings.cache_params,
Some(litellm_config::CacheSettings::Redis(_))
));
}
#[rstest]
fn non_redis_cache_settings_remain_opaque() {
let config = Config::from_yaml(
"litellm_settings:\n cache_params:\n type: local\n custom_option:\n nested: true",
)
.unwrap();
assert!(matches!(
config.litellm_settings.cache_params,
Some(litellm_config::CacheSettings::Other(_))
));
}

405
litellm-rust/crates/config/tests/redis.rs
View file

@ -1,405 +0,0 @@
use std::{
collections::HashMap,
fmt::Write,
sync::{Arc, Mutex},
};
use litellm_config::{CacheSettings, Config, Error, RedisEndpoint, RedisSettings, ResolvedRedis};
use rstest::rstest;
use tracing::{Event, Subscriber, field::Visit};
use tracing_subscriber::{Layer, layer::Context, prelude::*};
fn resolve(
yaml: &str,
environment: &[(&str, &str)],
) -> Result<litellm_config::ResolvedRedis, Error> {
let environment: HashMap<String, String> = environment
.iter()
.map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
.collect();
let config = Config::from_yaml(yaml).unwrap();
config
.general_settings
.coordination_redis
.as_ref()
.unwrap_or(&RedisSettings::default())
.resolve(&|name| environment.get(name).cloned())
}
fn resolve_cache(yaml: &str) -> Result<ResolvedRedis, Error> {
let config = Config::from_yaml(yaml).unwrap();
match config.litellm_settings.cache_params.unwrap() {
CacheSettings::Redis(settings) => settings.resolve_cache(&|_| None),
CacheSettings::Other(_) => panic!("expected a Redis cache"),
}
}
fn cluster_ports(resolved: &ResolvedRedis) -> Vec<u16> {
match &resolved.endpoint {
RedisEndpoint::Cluster(cluster) => cluster.nodes.iter().map(|node| node.port).collect(),
endpoint => panic!("expected a cluster, got {endpoint:?}"),
}
}
#[rstest]
#[case::host(
"general_settings:\n coordination_redis:\n host: redis\n port: 6380\n db: 2\n ssl: true\n namespace: app\n socket_timeout: 1.5",
&[],
"host:redis:6380:2:true:app:1500"
)]
#[case::environment_url(
"model_list: []",
&[("REDIS_URL", "redis://environment:6379")],
"url:redis://environment:6379"
)]
#[case::explicit_host_drops_environment_url(
"general_settings:\n coordination_redis:\n host: configured",
&[("REDIS_URL", "redis://environment:6379")],
"host:configured:6379:0:false::none"
)]
#[case::yaml_cluster(
"general_settings:\n coordination_redis:\n startup_nodes:\n - host: node-a\n port: 7000\n - host: node-b\n port: \"7001\"",
&[],
"cluster:node-a:7000,node-b:7001"
)]
#[case::json_environment_cluster(
"model_list: []",
&[("REDIS_CLUSTER_NODES", r#"[{"host":"node-a","port":"7000"},{"host":"node-b","port":7001}]"#)],
"cluster:node-a:7000,node-b:7001"
)]
#[case::missing_reference_keeps_environment_fallback(
"general_settings:\n coordination_redis:\n host: os.environ/MISSING_HOST",
&[("REDIS_HOST", "environment")],
"host:environment:6379:0:false::none"
)]
#[case::configured_values_beat_the_same_environment_values(
"general_settings:\n coordination_redis:\n host: configured\n port: 6380\n db: 3",
&[("REDIS_HOST", "environment"), ("REDIS_PORT", "6390"), ("REDIS_DB", "9")],
"host:configured:6380:3:false::none"
)]
#[case::environment_url_beats_the_environment_host(
"model_list: []",
&[("REDIS_URL", "redis://environment:6379"), ("REDIS_HOST", "ignored")],
"url:redis://environment:6379"
)]
#[case::startup_nodes_beat_a_url(
"general_settings:\n coordination_redis:\n url: redis://single:6379\n startup_nodes:\n - host: node-a\n port: 7000",
&[],
"cluster:node-a:7000"
)]
#[case::environment_references_resolve_the_host_and_port(
"general_settings:\n coordination_redis:\n host: os.environ/COORD_HOST\n port: os.environ/COORD_PORT",
&[("COORD_HOST", "referenced"), ("COORD_PORT", "6400")],
"host:referenced:6400:0:false::none"
)]
#[case::a_null_host_falls_back_to_the_environment(
"general_settings:\n coordination_redis:\n host: null",
&[("REDIS_HOST", "environment")],
"host:environment:6379:0:false::none"
)]
fn resolves_redis_through_the_parsed_config(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: &str,
) {
let resolved = resolve(yaml, environment).unwrap();
let description = match resolved.endpoint {
RedisEndpoint::Url(url) => format!("url:{}", url.expose()),
RedisEndpoint::Host(host) => format!(
"host:{}:{}:{}:{}:{}:{}",
host.host,
host.port,
host.db,
host.tls,
resolved.namespace.unwrap_or_default(),
resolved
.timeout
.map(|timeout| timeout.as_millis().to_string())
.unwrap_or_else(|| "none".to_owned())
),
RedisEndpoint::Cluster(cluster) => format!(
"cluster:{}",
cluster
.nodes
.iter()
.map(|node| format!("{}:{}", node.host, node.port))
.collect::<Vec<_>>()
.join(",")
),
};
assert_eq!(description, expected);
}
#[rstest]
#[case::custom_tls_file(
"general_settings:\n coordination_redis:\n host: redis\n ssl_ca_certs: /ca.pem",
&[],
"ssl_ca_certs"
)]
#[case::gcp_auth_from_config(
"general_settings:\n coordination_redis:\n host: redis\n gcp_service_account: account",
&[],
"gcp_service_account"
)]
#[case::gcp_auth_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_GCP_SERVICE_ACCOUNT", "account")],
"REDIS_GCP_SERVICE_ACCOUNT"
)]
#[case::azure_auth_from_config(
"general_settings:\n coordination_redis:\n host: redis\n azure_redis_ad_token: true",
&[],
"azure_redis_ad_token"
)]
#[case::azure_auth_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_AZURE_AD_TOKEN", "true")],
"REDIS_AZURE_AD_TOKEN"
)]
#[case::aws_auth_from_config(
"general_settings:\n coordination_redis:\n host: redis\n aws_iam_auth: true",
&[],
"aws_iam_auth"
)]
#[case::aws_auth_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_AWS_IAM_AUTH", "true")],
"REDIS_AWS_IAM_AUTH"
)]
#[case::sentinel_from_config(
"general_settings:\n coordination_redis:\n sentinel_nodes: [[redis, 26379]]",
&[],
"sentinel_nodes"
)]
#[case::sentinel_from_environment(
"model_list: []",
&[("REDIS_SENTINEL_NODES", "[[\"redis\",26379]]")],
"REDIS_SENTINEL_NODES"
)]
fn rejected_settings_name_their_source(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: &str,
) {
let error = resolve(yaml, environment).unwrap_err();
assert!(matches!(error, Error::RedisSetting(setting) if setting == expected));
}
#[rstest]
fn an_explicit_host_keeps_environment_credentials() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: configured",
&[
("REDIS_PASSWORD", "environment-secret"),
("REDIS_SSL", "True"),
],
)
.unwrap();
let RedisEndpoint::Host(host) = resolved.endpoint else {
panic!("expected a host endpoint");
};
assert_eq!(
host.password.as_ref().map(|password| password.expose()),
Some("environment-secret")
);
assert!(host.tls);
}
#[rstest]
#[case::only_a_port("general_settings:\n coordination_redis:\n port: 6379", &[])]
#[case::a_structured_host_with_an_environment_target(
"general_settings:\n coordination_redis:\n host: []",
&[("REDIS_HOST", "environment"), ("REDIS_URL", "redis://environment:6379")]
)]
fn a_block_without_a_usable_target_fails(#[case] yaml: &str, #[case] environment: &[(&str, &str)]) {
let error = resolve(yaml, environment).unwrap_err();
assert!(matches!(
error,
Error::RedisInvalid {
setting: "host",
..
}
));
}
#[rstest]
fn the_cache_alias_names_cluster_nodes() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n redis_startup_nodes:\n - host: node-a\n port: 7000",
)
.unwrap();
assert_eq!(cluster_ports(&resolved), vec![7000]);
}
#[rstest]
fn canonical_startup_nodes_beat_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n startup_nodes:\n - host: canonical\n port: 7000\n redis_startup_nodes:\n - host: alias\n port: 7100\n - host: alias\n port: 7101",
)
.unwrap();
assert_eq!(cluster_ports(&resolved), vec![7000]);
}
#[rstest]
fn a_null_canonical_startup_nodes_setting_shadows_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: null\n redis_startup_nodes:\n - host: alias\n port: 7100",
)
.unwrap();
assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_)));
}
#[rstest]
fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: os.environ/MISSING_NODES\n redis_startup_nodes:\n - host: alias\n port: 7100",
)
.unwrap();
assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_)));
}
#[rstest]
fn rejected_environment_settings_never_print_the_value() {
let error = resolve(
"model_list: []",
&[
("REDIS_HOST", "redis"),
("REDIS_SSL_CERT_REQS", "private-option-value"),
],
)
.unwrap_err();
let message = error.to_string();
assert!(message.contains("REDIS_SSL_CERT_REQS"));
assert!(!message.contains("private-option-value"));
}
#[rstest]
#[case::negative("-1")]
#[case::not_a_number("slow")]
fn rejects_bad_socket_timeout(#[case] timeout: &str) {
let yaml = format!(
"general_settings:\n coordination_redis:\n host: redis\n socket_timeout: {timeout}"
);
let error = resolve(&yaml, &[]).unwrap_err();
assert!(matches!(
error,
Error::RedisInvalid {
setting: "socket_timeout",
..
}
));
}
#[rstest]
fn defaults_socket_timeout_for_the_gateway() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: redis",
&[],
)
.unwrap();
assert_eq!(resolved.timeout, None);
}
#[rstest]
fn debug_redacts_redis_secrets_and_url() {
let config = Config::from_yaml(
"general_settings:\n coordination_redis:\n url: redis://user:password@redis:6379\n username: credential-user\n password: credential-password",
)
.unwrap();
let settings = config.general_settings.coordination_redis.unwrap();
let resolved = settings.resolve(&|_| None).unwrap();
let debug = format!("{settings:?} {resolved:?}");
assert!(!debug.contains("redis://user:password"));
assert!(!debug.contains("credential-user"));
assert!(!debug.contains("credential-password"));
assert!(debug.contains("redacted") || debug.contains("REDACTED"));
}
#[rstest]
fn inactive_redis_cache_keeps_malformed_known_inputs_deferred() {
let config = Config::from_yaml(
"litellm_settings:\n cache: false\n cache_params:\n type: redis\n host: []\n port: invalid\n socket_timeout: invalid",
)
.unwrap();
assert!(matches!(
config.litellm_settings.cache_params,
Some(litellm_config::CacheSettings::Redis(_))
));
}
#[rstest]
fn non_redis_cache_settings_remain_opaque() {
let config = Config::from_yaml(
"litellm_settings:\n cache_params:\n type: local\n custom_option:\n nested: true",
)
.unwrap();
assert!(matches!(
config.litellm_settings.cache_params,
Some(litellm_config::CacheSettings::Other(_))
));
}
#[rstest]
fn legacy_unknown_settings_do_not_change_resolution() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: redis\n redis_flush_size: 42\n service_name: primary",
&[],
)
.unwrap();
assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_)));
}
struct WarningCapture(Arc<Mutex<Vec<String>>>);
impl<S: Subscriber> Layer<S> for WarningCapture {
fn on_event(&self, event: &Event<'_>, _context: Context<'_, S>) {
let mut fields = Fields::default();
event.record(&mut fields);
self.0.lock().unwrap().push(fields.0);
}
}
#[derive(Default)]
struct Fields(String);
impl Visit for Fields {
fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
write!(self.0, "{}={value:?} ", field.name()).unwrap();
}
fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
write!(self.0, "{}={value} ", field.name()).unwrap();
}
}
#[rstest]
fn ignored_setting_warning_names_the_setting_without_its_value() {
let events = Arc::new(Mutex::new(Vec::new()));
let subscriber = tracing_subscriber::registry().with(WarningCapture(events.clone()));
tracing::subscriber::with_default(subscriber, || {
resolve(
"general_settings:\n coordination_redis:\n host: redis\n max_connections: 47",
&[],
)
.unwrap();
});
let warnings = events.lock().unwrap().join(" ");
assert!(warnings.contains("max_connections"));
assert!(!warnings.contains("47"));
}

37
litellm-rust/crates/gateway-forward/AGENTS.md Normal file
View file

@ -0,0 +1,37 @@
# Forwarding gateway
`Forward` is a reverse proxy in front of the Python proxy on loopback. Python must see each request as if the client had connected to it directly, because every Python route (inference, MCP, pass-through, UI, WebSocket) is served behind it without Rust knowing which one it is
- Serve HTTP/1.1 only, as uvicorn does. The supervisor never puts Rust in front of a Python server that speaks HTTP/2
- Never branch on a route or path. A behavior one endpoint needs is a protocol behavior every endpoint gets
- Forward the method, request target, headers, and body byte for byte within the [URL contract](FORWARDING.md#url-fragments-and-literal-hashes). An unescaped `#` starts a fragment and is omitted with everything after it, while `%23` is preserved. Only drop hop-by-hop headers (RFC 9110 section 7.6.1) and set `x-litellm-rust-peer`, which replaces any client copy
- Stream both bodies. Never buffer, cap, or time out a body, since uploads and provider streams have no upper bound
- A client disconnect closes the upstream connection at every stage, because that is how Python learns to stop a provider stream
- An upstream that dies mid-response fails the client body. Never end it cleanly, or SSE clients treat a truncated stream as complete
- Never invent a Host. Python builds absolute URLs from it, so a missing one must never become the private upstream address
- A request Python could not have received unchanged (HTTP/1.1 without Host, a transfer coding other than chunked, a target without a path) gets the status uvicorn would have given, from Rust
- Pooled upstream connections are retired well before Python's keep-alive timeout, so a request never lands on a connection Python is closing
- Idle client connections close after Python's keep-alive timeout, as uvicorn closed them, so a client that vanishes cannot hold a socket forever. Only the wait for the next request is timed, never an exchange in progress
- WebSocket is the one protocol switch Python serves. Upgrades are relayed as raw bytes after the 101 and drained on shutdown
## Layout
Every forwarding decision has one home. `rewrite.rs` owns what Python and the client receive, as pure functions over requests and responses. `Rejection` in `error.rs` is the only place that picks a status for a request Rust answers itself. `settings.rs` parses the environment into typed values, so an `Upstream` is a loopback origin by construction. `server.rs` accepts client connections and owns their timeouts, since `axum::serve` exposes none, and `connections.rs` tracks every connection and relay for shutdown. `forward.rs` and `websocket.rs` only move requests, bodies and upgraded connections and never touch a header themselves. `connector.rs` keeps hyper reading Python's answer after Python stops reading the request
## Environment
`Settings` reads only these variables, and only when `LITELLM_FORWARD_UPSTREAM` is set. The Python supervisor in `litellm/proxy/rust_forward/` sets every one of them from the values Python actually runs with, so all of them are required and a missing or invalid one fails startup naming it. The crate never reads Python's own settings (`HOST`, `PORT`, `KEEPALIVE_TIMEOUT`, `GRACEFUL_SHUTDOWN_TIMEOUT`) and keeps no copy of Python's defaults, since Python can get those values from CLI flags the environment never shows
| Variable | Meaning |
|---|---|
| `LITELLM_FORWARD_UPSTREAM` | Python's loopback origin, e.g. `http://127.0.0.1:53211` |
| `LITELLM_FORWARD_HOST` | public bind host |
| `LITELLM_FORWARD_PORT` | public bind port |
| `LITELLM_FORWARD_UPSTREAM_KEEPALIVE` | seconds uvicorn keeps an idle connection open. Pooled connections are retired at half of it, and idle client connections are closed after it |
| `LITELLM_FORWARD_SHUTDOWN_DRAIN` | seconds to drain on shutdown, already including Python's graceful timeout and teardown headroom |
## Tests
Test through real sockets with `tests/support`, never `oneshot`. `Gateway` runs the forwarder on a listener like production, `recording` is a well-behaved upstream that captures what Python would receive, and `raw_upstream` and `Raw` drive bytes by hand for framing, timing and disconnects a real server would not produce. Waits measured in minutes use paused tokio time instead of real sleeps
The forwarder sees HTTP/1.1 messages, not routes, so the tests are split into one file per traffic shape uvicorn can emit: `exchange` (plain request and response), `framing` (requests whose framing a parser could read two ways), `slow_head` (minutes before response headers), `stream` (incremental response bodies), `idle_stream` (long-lived streams that go quiet), `upload`, `download`, `websocket`, `client_identity`, `lifecycle` (disconnects, pooling, idle timeouts, shutdown) and `roundtrip` (the forwarding invariant for generated messages). A new route never gets its own test. Add its behavior to the shape it already has, and when it needs a wire behavior none of the files covers, add a new shape file for it

30
litellm-rust/crates/gateway-forward/Cargo.toml Normal file
View file

@ -0,0 +1,30 @@
[package]
name = "litellm-gateway-forward"
version = "0.1.0"
edition.workspace = true
rust-version.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
axum.workspace = true
envy = "0.4.2"
hyper = { version = "1.12.0", features = ["http1", "server"] }
hyper-util = { workspace = true, features = ["client-legacy", "http1", "tokio"] }
serde.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["io-util"] }
tokio-util = { version = "0.7", features = ["rt"] }
tower-service = "0.3.3"
tracing.workspace = true
[dev-dependencies]
bytes.workspace = true
futures-util.workspace = true
http-body-util = "0.1"
httparse = "1"
hyper = { version = "1.12.0", features = ["client"] }
proptest.workspace = true
rstest.workspace = true
tokio = { workspace = true, features = ["io-util", "sync", "test-util"] }
tokio-tungstenite.workspace = true

31
litellm-rust/crates/gateway-forward/FORWARDING.md Normal file
View file

@ -0,0 +1,31 @@
# Python forwarding compatibility
The opt-in forwarding gateway serves HTTP/1.1 and WebSockets in front of Python. Python still owns every route and policy. Its protocol contract and environment settings are in [AGENTS.md](AGENTS.md)
## URL fragments and literal hashes
The gateway follows Hyper's `http::Uri` parsing. If a raw client sends an unescaped `#` in the request target, only the path and query before the first `#` reach Python. For example, `/v1/files?purpose=batch#retry&limit=1` becomes `/v1/files?purpose=batch`, and `/gemini/v1beta/files/a#b?key=k` becomes `/gemini/v1beta/files/a`
This is accepted behavior because `#` starts a URL fragment, which is separated before dereferencing ([RFC 3986 section 3.5](https://www.rfc-editor.org/rfc/rfc3986.html#section-3.5)). HTTP origin-form contains only the path and optional query, not a fragment ([RFC 9112 section 3.2.1](https://www.rfc-editor.org/rfc/rfc9112.html#section-3.2.1)). Browsers and URL-based clients omit the fragment before sending the request, so UI hash routing is unaffected
A literal hash in path or query data must be encoded as `%23`, which the gateway forwards unchanged. Direct Python with h11 also accepts unescaped hashes as data, but that permissive raw-target behavior is outside this URL contract. Active socket tests cover both [fragment omission](tests/exchange.rs) and [encoded-hash preservation](tests/exchange.rs)
## Known differences from uvicorn with h11
These cases remain executable as ignored socket tests. They are not passing parity checks
| Input | Direct Python | Through Rust | Regression test |
|---|---|---|---|
| Header after a removed hop-by-hop field | Keeps the remaining header order | `HeaderMap::remove` can move a later header | [dropping_a_hop_by_hop_header_keeps_the_others_in_order](tests/exchange.rs) |
| Repeated header names interleaved with other names | Keeps the interleaving | Groups values by name, keeping their order within that name | [interleaved_repeated_headers_reach_python_interleaved](tests/exchange.rs) |
| Obsolete folded header, such as `x-a: b\r\n c` | Unfolds to `b c` | Rejects with 400 | [a_folded_header_reaches_python_unfolded](tests/exchange.rs) |
| Identical Content-Length list, such as `7, 7` | Accepts length 7 | Rejects with 400 | [a_content_length_list_of_one_value_reaches_python](tests/framing.rs) |
| HTTP/1.0 request with chunked transfer coding | Decodes the body | Rejects with 400 | [an_http10_chunked_upload_reaches_python](tests/framing.rs) |
Clients can avoid the parser rejections by sending unfolded headers, one Content-Length value, and HTTP/1.1 for chunked uploads. Integrations must not depend on header ordering across different field names
Run `cargo test -p litellm-gateway-forward --test exchange --test framing -- --ignored` from this directory to reproduce the differences
## Connection header handling
The gateway strips standard hop-by-hop fields and fields named by Connection, except routing, representation, credential, forwarding, and WebSocket handshake headers. A client cannot nominate those end-to-end fields for removal. The private `x-litellm-rust-peer` always replaces any client-supplied copy

58
litellm-rust/crates/gateway-forward/src/connections.rs Normal file
View file

@ -0,0 +1,58 @@
use tokio_util::{
sync::{CancellationToken, DropGuard},
task::TaskTracker,
};
/// Every client connection and WebSocket relay the gateway runs, so shutdown can wait for them
/// and then close whatever outlives the drain.
#[derive(Clone, Default)]
pub(crate) struct Connections {
tasks: TaskTracker,
draining: CancellationToken,
closed: CancellationToken,
}
impl Connections {
/// Runs `work` until it finishes or the gateway closes every connection.
pub fn spawn(&self, work: impl Future<Output = ()> + Send + 'static) {
let closed = self.closed.clone();
self.tasks.spawn(async move {
tokio::select! {
() = closed.cancelled() => {}
() = work => {}
}
});
}
pub fn begin_shutdown(&self) {
self.draining.cancel();
}
pub async fn draining(&self) {
self.draining.cancelled().await;
}
pub fn is_draining(&self) -> bool {
self.draining.is_cancelled()
}
pub async fn until_closed<T>(&self, future: impl Future<Output = T>) -> Option<T> {
tokio::select! {
() = self.closed.cancelled() => None,
output = future => Some(output),
}
}
pub fn close(&self) {
self.closed.cancel();
}
pub fn close_on_drop(&self) -> DropGuard {
self.closed.clone().drop_guard()
}
pub async fn wait(&self) {
self.tasks.close();
self.tasks.wait().await;
}
}

122
litellm-rust/crates/gateway-forward/src/connector.rs Normal file
View file

@ -0,0 +1,122 @@
use std::{
io::{self, IoSlice},
pin::Pin,
task::{Context, Poll, ready},
};
use axum::http::Uri;
use hyper::rt::{Read, ReadBufCursor, Write};
use hyper_util::{
client::legacy::connect::{Connected, Connection, HttpConnector},
rt::TokioIo,
};
use tokio::net::TcpStream;
/// Connects to Python like `HttpConnector`, but keeps hyper reading Python's answer after Python
/// stops reading the request. uvicorn rejects an upload (413, 401) as soon as it reads the head
/// and closes with the body unread, so writing the rest fails, and hyper would fail the request
/// and drop the answer it already received. Like curl and httpcore, the forwarder discards the
/// rest of the request instead and reads on. A connection that never answered still fails, on the
/// read, and no request is ever sent twice.
#[derive(Clone)]
pub(crate) struct Connector(pub HttpConnector);
impl tower_service::Service<Uri> for Connector {
type Response = AnswerAfterRefusal;
type Error = <HttpConnector as tower_service::Service<Uri>>::Error;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.0.poll_ready(context)
}
fn call(&mut self, uri: Uri) -> Self::Future {
let connecting = self.0.call(uri);
Box::pin(async move {
Ok(AnswerAfterRefusal {
stream: connecting.await?,
refused: false,
})
})
}
}
pub(crate) struct AnswerAfterRefusal {
stream: TokioIo<TcpStream>,
refused: bool,
}
impl AnswerAfterRefusal {
fn written(&mut self, written: io::Result<usize>, length: usize) -> io::Result<usize> {
match written {
Err(error)
if matches!(
error.kind(),
io::ErrorKind::BrokenPipe | io::ErrorKind::ConnectionReset
) =>
{
self.refused = true;
Ok(length)
}
written => written,
}
}
}
impl Read for AnswerAfterRefusal {
fn poll_read(
mut self: Pin<&mut Self>,
context: &mut Context<'_>,
buffer: ReadBufCursor<'_>,
) -> Poll<io::Result<()>> {
Pin::new(&mut self.stream).poll_read(context, buffer)
}
}
impl Write for AnswerAfterRefusal {
fn poll_write(
mut self: Pin<&mut Self>,
context: &mut Context<'_>,
buffer: &[u8],
) -> Poll<io::Result<usize>> {
if self.refused {
return Poll::Ready(Ok(buffer.len()));
}
let written = ready!(Pin::new(&mut self.stream).poll_write(context, buffer));
Poll::Ready(self.written(written, buffer.len()))
}
fn poll_write_vectored(
mut self: Pin<&mut Self>,
context: &mut Context<'_>,
buffers: &[IoSlice<'_>],
) -> Poll<io::Result<usize>> {
let length = buffers.iter().map(|buffer| buffer.len()).sum();
if self.refused {
return Poll::Ready(Ok(length));
}
let written = ready!(Pin::new(&mut self.stream).poll_write_vectored(context, buffers));
Poll::Ready(self.written(written, length))
}
fn is_write_vectored(&self) -> bool {
self.stream.is_write_vectored()
}
fn poll_flush(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<io::Result<()>> {
if self.refused {
return Poll::Ready(Ok(()));
}
Pin::new(&mut self.stream).poll_flush(context)
}
fn poll_shutdown(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.stream).poll_shutdown(context)
}
}
impl Connection for AnswerAfterRefusal {
fn connected(&self) -> Connected {
self.stream.connected()
}
}

91
litellm-rust/crates/gateway-forward/src/error.rs Normal file
View file

@ -0,0 +1,91 @@
use axum::{
http::StatusCode,
response::{IntoResponse, Response},
};
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("invalid forwarding upstream URI")]
UpstreamUri(#[source] axum::http::uri::InvalidUri),
#[error(
"LITELLM_FORWARD_UPSTREAM must be a loopback IPv4 http origin without a path, query or credentials"
)]
Upstream,
#[error("{0} must be set in forward mode")]
MissingSetting(String),
#[error("invalid forwarding setting: {0}")]
InvalidSetting(String),
#[error("forwarding gateway I/O failed: {0}")]
Io(#[from] std::io::Error),
}
/// Why a request got an answer from Rust instead of from Python.
#[derive(Debug, thiserror::Error)]
pub(crate) enum Rejection {
#[error("HTTP/1.1 request without a Host header")]
MissingHost,
#[error("request target has no path to forward")]
NoPath,
#[error("request uses a transfer coding other than chunked")]
UnsupportedTransferCoding,
#[error("client request body failed: {0}")]
ClientBody(#[source] hyper_util::client::legacy::Error),
#[error("python proxy unreachable: {0}")]
Unreachable(#[source] hyper_util::client::legacy::Error),
#[error("python proxy switched protocols without being asked to")]
UnrequestedUpgrade,
#[error("python proxy websocket upgrade failed: {0}")]
UpgradeFailed(#[source] hyper::Error),
#[error("forwarding gateway is shutting down")]
ShuttingDown,
#[error("could not build the request for the python proxy")]
Internal,
}
impl Rejection {
/// A client body that ends early or breaks its framing fails the upstream request too, but
/// the fault is the client's, not Python's.
pub fn from_upstream(error: hyper_util::client::legacy::Error) -> Self {
let client_body = std::error::Error::source(&error)
.and_then(|source| source.downcast_ref::<hyper::Error>())
.is_some_and(hyper::Error::is_user);
match client_body {
true => Self::ClientBody(error),
false => Self::Unreachable(error),
}
}
fn status(&self) -> StatusCode {
match self {
Self::MissingHost
| Self::NoPath
| Self::UnsupportedTransferCoding
| Self::ClientBody(_) => StatusCode::BAD_REQUEST,
Self::Unreachable(_) | Self::UnrequestedUpgrade | Self::UpgradeFailed(_) => {
StatusCode::BAD_GATEWAY
}
Self::ShuttingDown => StatusCode::SERVICE_UNAVAILABLE,
Self::Internal => StatusCode::INTERNAL_SERVER_ERROR,
}
}
fn is_gateway_fault(&self) -> bool {
matches!(
self,
Self::Unreachable(_)
| Self::UnrequestedUpgrade
| Self::UpgradeFailed(_)
| Self::Internal
)
}
}
impl IntoResponse for Rejection {
fn into_response(self) -> Response {
match self.is_gateway_fault() {
true => tracing::warn!(error = %self, "request not forwarded"),
false => tracing::debug!(error = %self, "request not forwarded"),
}
self.status().into_response()
}
}

169
litellm-rust/crates/gateway-forward/src/forward.rs Normal file
View file

@ -0,0 +1,169 @@
use std::{
pin::Pin,
task::{Context, Poll},
time::Duration,
};
use axum::{
body::Body,
extract::Request,
http::{StatusCode, Uri},
response::{IntoResponse, Response},
};
use hyper::body::{Frame, Incoming, SizeHint};
use hyper_util::{
client::legacy::{Client, connect::HttpConnector},
rt::{TokioExecutor, TokioTimer},
};
use tokio::net::TcpListener;
use crate::{
Rejection, Upstream, connections::Connections, connector::Connector, rewrite, server,
server::ClientSocket, websocket,
};
// h11 counts no headers and only caps a head at 16 KiB, so it accepts as many headers as fit in
// that many shortest lines (`a:\r\n`). Python copies every provider header into its response
// too, so hyper's default of 100 would turn ordinary traffic into 431s and 502s.
pub(crate) const MAX_HEADERS: usize = 16 * 1024 / b"a:\r\n".len();
#[derive(Clone)]
pub struct Forward {
upstream: Uri,
client: Client<Connector, Body>,
connections: Connections,
keepalive: Duration,
}
impl Forward {
/// `upstream_keepalive` is how long the upstream keeps an idle connection open. Pooled
/// connections are dropped well before that, so a request never lands on one it is closing,
/// and idle client connections are closed after it, as uvicorn did.
pub fn new(upstream: &Upstream, upstream_keepalive: Duration) -> Self {
let mut connector = HttpConnector::new();
connector.set_nodelay(true);
let client = Client::builder(TokioExecutor::new())
.http1_max_headers(MAX_HEADERS)
.pool_timer(TokioTimer::new())
.pool_idle_timeout(upstream_keepalive / 2)
.set_host(false)
.build(Connector(connector));
Self {
upstream: upstream.uri().clone(),
client,
connections: Connections::default(),
keepalive: upstream_keepalive,
}
}
/// Serves until `shutdown` resolves, then lets open requests and WebSockets finish for up to
/// `drain` before closing whatever is left. Python shuts down alongside and ends long streams
/// itself, so `drain` only backstops a Python that never does.
pub async fn serve(
self,
listener: TcpListener,
shutdown: impl Future<Output = ()>,
drain: Duration,
) {
server::serve(self, listener, shutdown, drain).await;
}
pub(crate) fn connections(&self) -> &Connections {
&self.connections
}
pub(crate) fn keepalive(&self) -> Duration {
self.keepalive
}
pub(crate) async fn handle(&self, client: ClientSocket, request: Request) -> Response {
self.forward(client, request)
.await
.unwrap_or_else(IntoResponse::into_response)
}
async fn forward(
&self,
client: ClientSocket,
mut request: Request,
) -> Result<Response, Rejection> {
if let Some(downstream) = websocket::take_handshake(&self.connections, &mut request)? {
return self.forward_websocket(client, request, downstream).await;
}
let request = rewrite::upstream_request(request, &self.upstream, client, false)?;
let response = self
.client
.request(request)
.await
.map_err(Rejection::from_upstream)?;
http_response(response)
}
async fn forward_websocket(
&self,
client: ClientSocket,
request: Request,
downstream: hyper::upgrade::OnUpgrade,
) -> Result<Response, Rejection> {
let request = rewrite::upstream_request(request, &self.upstream, client, true)?;
let response = self
.connections
.until_closed(self.client.request(request))
.await
.ok_or(Rejection::ShuttingDown)?
.map_err(Rejection::from_upstream)?;
if response.status() == StatusCode::SWITCHING_PROTOCOLS {
return websocket::switch(&self.connections, downstream, response).await;
}
http_response(response)
}
}
fn http_response(response: Response<Incoming>) -> Result<Response, Rejection> {
if response.status() == StatusCode::SWITCHING_PROTOCOLS {
return Err(Rejection::UnrequestedUpgrade);
}
Ok(
rewrite::downstream_response(response)
.map(|body| Body::new(Relayed { body, failed: None })),
)
}
/// Python's response body on its way to the client.
struct Relayed {
body: Incoming,
failed: Option<hyper::Error>,
}
impl hyper::body::Body for Relayed {
type Data = <Incoming as hyper::body::Body>::Data;
type Error = hyper::Error;
// hyper aborts a client connection whose response body fails without flushing the head it
// buffered, so a head Python sent right before dying would never arrive and SDKs would retry
// a request Python already ran. Holding the error for one poll lets hyper flush the head.
fn poll_frame(
mut self: Pin<&mut Self>,
context: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
if let Some(error) = self.failed.take() {
return Poll::Ready(Some(Err(error)));
}
match Pin::new(&mut self.body).poll_frame(context) {
Poll::Ready(Some(Err(error))) => {
self.failed = Some(error);
context.waker().wake_by_ref();
Poll::Pending
}
frame => frame,
}
}
fn is_end_stream(&self) -> bool {
self.failed.is_none() && self.body.is_end_stream()
}
fn size_hint(&self) -> SizeHint {
self.body.size_hint()
}
}

29
litellm-rust/crates/gateway-forward/src/lib.rs Normal file
View file

@ -0,0 +1,29 @@
mod connections;
mod connector;
mod error;
mod forward;
mod rewrite;
mod server;
mod settings;
mod websocket;
pub use error::Error;
pub(crate) use error::Rejection;
pub use forward::Forward;
pub use rewrite::PEER_HEADER;
pub use settings::{Settings, Upstream};
/// Binds the public listener and forwards to Python until `shutdown` resolves, then drains for
/// `settings.shutdown_drain`.
pub async fn run(
settings: Settings,
shutdown: impl Future<Output = ()> + Send + 'static,
) -> Result<(), Error> {
let forward = Forward::new(&settings.upstream, settings.upstream_keepalive);
let listener = tokio::net::TcpListener::bind((settings.host.as_str(), settings.port)).await?;
tracing::info!(address = %listener.local_addr()?, "forwarding gateway listening");
forward
.serve(listener, shutdown, settings.shutdown_drain)
.await;
Ok(())
}

161
litellm-rust/crates/gateway-forward/src/rewrite.rs Normal file
View file

@ -0,0 +1,161 @@
use axum::{
body::HttpBody,
extract::Request,
http::{HeaderMap, HeaderValue, Response, Uri, Version, header},
};
use crate::{Rejection, server::ClientSocket};
pub const PEER_HEADER: &str = "x-litellm-rust-peer";
/// The request Python receives: the client's method, target, headers and body unchanged, sent to
/// `upstream` over HTTP/1.1 without the client connection's own headers, plus the client's TCP
/// peer in `PEER_HEADER`, which replaces any copy the client sent.
pub(crate) fn upstream_request(
mut request: Request,
upstream: &Uri,
client: ClientSocket,
websocket: bool,
) -> Result<Request, Rejection> {
if !only_chunked(request.headers()) {
return Err(Rejection::UnsupportedTransferCoding);
}
let path_and_query = request
.uri()
.path_and_query()
.cloned()
.ok_or(Rejection::NoPath)?;
if let Some(host) = public_host(&request, client)? {
request.headers_mut().insert(header::HOST, host);
}
let mut uri = upstream.clone().into_parts();
uri.path_and_query = Some(path_and_query);
*request.uri_mut() = Uri::from_parts(uri).map_err(|_| Rejection::Internal)?;
*request.version_mut() = Version::HTTP_11;
let chunked = request.body().size_hint().exact().is_none();
let headers = request.headers_mut();
strip_hop_by_hop(headers);
if chunked {
// hyper sends no body for a GET with an unknown length unless told it is chunked.
headers.insert(
header::TRANSFER_ENCODING,
HeaderValue::from_static("chunked"),
);
}
if websocket {
restore_websocket_handshake(headers);
}
headers.insert(
PEER_HEADER,
HeaderValue::from_str(&client.peer.to_string()).map_err(|_| Rejection::Internal)?,
);
Ok(request)
}
// Python builds absolute URLs from Host and, without a usable one, from the address it listens
// on, which behind Rust is the private loopback port. A client that sent no usable Host gets the
// public address it connected to, which is what Python used before. h11 rejects HTTP/1.1 without
// Host.
fn public_host(request: &Request, client: ClientSocket) -> Result<Option<HeaderValue>, Rejection> {
let host = request.headers().get(header::HOST);
if host.is_some_and(|host| !host.is_empty()) {
return Ok(None);
}
if host.is_none() && request.version() == Version::HTTP_11 {
return Err(Rejection::MissingHost);
}
HeaderValue::from_str(&client.local.to_string())
.map(Some)
.map_err(|_| Rejection::Internal)
}
// hyper only decodes chunked framing, so anything else would reach Python still coded but framed
// as a plain body. uvicorn rejects it with the 400 it sends for every h11 parse error.
fn only_chunked(headers: &HeaderMap) -> bool {
headers
.get_all(header::TRANSFER_ENCODING)
.iter()
.flat_map(|value| value.as_bytes().split(|byte| *byte == b','))
.all(|coding| coding.trim_ascii().eq_ignore_ascii_case(b"chunked"))
}
/// The response the client receives: Python's, without the upstream connection's own headers.
pub(crate) fn downstream_response<B>(mut response: Response<B>) -> Response<B> {
strip_hop_by_hop(response.headers_mut());
response
}
/// Python's 101 as the client receives it, still naming the protocol both sides switch to.
pub(crate) fn switched_response<B>(mut response: Response<B>) -> Response<B> {
strip_hop_by_hop(response.headers_mut());
restore_websocket_handshake(response.headers_mut());
response
}
pub(crate) fn is_websocket_handshake(headers: &HeaderMap) -> bool {
headers
.get(header::UPGRADE)
.is_some_and(|value| value.as_bytes().eq_ignore_ascii_case(b"websocket"))
&& headers
.get_all(header::CONNECTION)
.iter()
.filter_map(|value| value.to_str().ok())
.flat_map(|value| value.split(','))
.any(|token| token.trim().eq_ignore_ascii_case("upgrade"))
}
const HOP_BY_HOP: [&str; 9] = [
"connection",
"keep-alive",
"proxy-authenticate",
"proxy-authorization",
"te",
"trailer",
"transfer-encoding",
"upgrade",
"proxy-connection",
];
// Per-connection headers (RFC 9110 section 7.6.1) describe the client's connection to Rust, not
// Rust's connection to Python, so a proxy must drop them instead of forwarding them.
fn strip_hop_by_hop(headers: &mut HeaderMap) {
let named: Vec<_> = headers
.get_all(header::CONNECTION)
.iter()
.filter_map(|value| value.to_str().ok())
.flat_map(|value| value.split(','))
.filter_map(|name| header::HeaderName::from_bytes(name.trim().as_bytes()).ok())
.filter(|name| !must_forward(name))
.collect();
for name in named {
headers.remove(name);
}
for name in HOP_BY_HOP {
headers.remove(name);
}
}
fn must_forward(name: &header::HeaderName) -> bool {
matches!(
name.as_str(),
"host"
| "authorization"
| "cookie"
| "set-cookie"
| "www-authenticate"
| "authentication-info"
| "api-key"
| "x-api-key"
| "litellm-api-key"
| "forwarded"
| PEER_HEADER
) || name.as_str().starts_with("content-")
|| name.as_str().starts_with("x-forwarded-")
|| name.as_str().starts_with("sec-websocket-")
}
// Stripping drops the handshake's Connection and Upgrade, so each hop gets them back.
fn restore_websocket_handshake(headers: &mut HeaderMap) {
headers.insert(header::CONNECTION, HeaderValue::from_static("upgrade"));
headers.insert(header::UPGRADE, HeaderValue::from_static("websocket"));
}

116
litellm-rust/crates/gateway-forward/src/server.rs Normal file
View file

@ -0,0 +1,116 @@
use std::{convert::Infallible, io, net::SocketAddr, pin::pin, time::Duration};
use axum::{body::Body, extract::Request};
use hyper::{body::Incoming, server::conn::http1, service::service_fn};
use hyper_util::rt::{TokioIo, TokioTimer};
use tokio::net::TcpListener;
use crate::{Forward, forward::MAX_HEADERS};
/// The client's TCP connection to the public listener.
#[derive(Clone, Copy, Debug)]
pub(crate) struct ClientSocket {
pub peer: SocketAddr,
pub local: SocketAddr,
}
/// Serves each public connection with hyper directly, because `axum::serve` cannot configure
/// the connection timeouts uvicorn applied to the same clients.
pub(crate) async fn serve(
forward: Forward,
listener: TcpListener,
shutdown: impl Future<Output = ()>,
drain: Duration,
) {
let connections = forward.connections().clone();
let _close_on_drop = connections.close_on_drop();
let keepalive = forward.keepalive();
let builder = builder(keepalive);
let mut shutdown = pin!(shutdown);
loop {
let (stream, peer) = tokio::select! {
() = &mut shutdown => break,
accepted = listener.accept() => match accepted {
Ok(accepted) => accepted,
Err(error) => {
back_off(error).await;
continue;
}
},
};
let local = match stream.local_addr() {
Ok(local) => local,
Err(error) => {
tracing::debug!(%error, "client connection closed before it was served");
continue;
}
};
if let Err(error) = stream.set_nodelay(true) {
tracing::warn!(%error, "could not disable Nagle on a client connection");
}
let socket = ClientSocket { peer, local };
let forward = forward.clone();
let service = service_fn(move |request: Request<Incoming>| {
let forward = forward.clone();
async move { Ok::<_, Infallible>(forward.handle(socket, request.map(Body::new)).await) }
});
let connection = builder
.serve_connection(TokioIo::new(stream), service)
.with_upgrades();
let draining = connections.clone();
connections.spawn(async move {
let mut connection = pin!(connection);
let ended = tokio::select! {
ended = connection.as_mut() => ended,
() = draining.draining() => {
connection.as_mut().graceful_shutdown();
connection.await
}
};
if let Err(error) = ended {
tracing::debug!(%error, "client connection ended with an error");
}
});
}
drop(listener);
connections.begin_shutdown();
if tokio::time::timeout(drain, connections.wait())
.await
.is_err()
{
tracing::warn!(
?drain,
"closing connections still open after the shutdown drain"
);
}
connections.close();
connections.wait().await;
}
// uvicorn closes a client connection that sits idle for `--keepalive_timeout` between requests.
// Without the same limit, a client that vanishes without closing (NAT expiry, a crashed host)
// would hold its socket for the life of the process until the gateway runs out of descriptors.
// Operators already tune that timeout against their load balancer, so the public side keeps
// Python's value instead of introducing its own. hyper only runs the header timer while a
// connection waits for its next request, so it never cuts a long stream, upload or WebSocket.
fn builder(keepalive: Duration) -> http1::Builder {
let mut builder = http1::Builder::new();
builder
.max_headers(MAX_HEADERS)
.timer(TokioTimer::new())
.header_read_timeout(keepalive);
builder
}
async fn back_off(error: io::Error) {
if matches!(
error.kind(),
io::ErrorKind::ConnectionRefused
| io::ErrorKind::ConnectionAborted
| io::ErrorKind::ConnectionReset
) {
return;
}
tracing::error!(%error, "could not accept a client connection");
tokio::time::sleep(Duration::from_secs(1)).await;
}

113
litellm-rust/crates/gateway-forward/src/settings.rs Normal file
View file

@ -0,0 +1,113 @@
use std::{ffi::OsString, net::SocketAddr, str::FromStr, time::Duration};
use axum::http::Uri;
use serde::Deserialize;
use crate::Error;
const UPSTREAM: &str = "LITELLM_FORWARD_UPSTREAM";
/// The Python proxy's private origin. Only a loopback IPv4 address parses, so nothing can point
/// the forwarder at another host.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Upstream(Uri);
impl Upstream {
pub(crate) fn uri(&self) -> &Uri {
&self.0
}
}
impl FromStr for Upstream {
type Err = Error;
fn from_str(value: &str) -> Result<Self, Error> {
let uri: Uri = value.parse().map_err(Error::UpstreamUri)?;
let address = uri
.authority()
.filter(|_| uri.scheme_str() == Some("http"))
.filter(|authority| !authority.as_str().contains('@'))
.filter(|_| uri.path() == "/" && uri.query().is_none())
.and_then(|authority| authority.as_str().parse::<SocketAddr>().ok());
match address {
Some(SocketAddr::V4(address)) if address.ip().is_loopback() => Ok(Self(uri)),
_ => Err(Error::Upstream),
}
}
}
#[derive(Deserialize)]
struct Environment {
litellm_forward_upstream: String,
litellm_forward_host: String,
litellm_forward_port: u16,
litellm_forward_upstream_keepalive: f64,
litellm_forward_shutdown_drain: f64,
}
/// Everything forward mode runs with. The Python supervisor resolves each value from the settings
/// Python actually runs with, so none of them is guessed or defaulted here.
/// Setting `LITELLM_FORWARD_UPSTREAM` requires every other `LITELLM_FORWARD_*` setting below.
#[derive(Debug, Clone, PartialEq)]
pub struct Settings {
pub upstream: Upstream,
pub host: String,
pub port: u16,
pub upstream_keepalive: Duration,
pub shutdown_drain: Duration,
}
impl Settings {
/// `None` when `LITELLM_FORWARD_UPSTREAM` is unset, which means the gateway is not forwarding.
pub fn from_env() -> Result<Option<Self>, Error> {
Self::from_vars(std::env::vars_os())
}
pub fn from_vars(
variables: impl IntoIterator<Item = (OsString, OsString)>,
) -> Result<Option<Self>, Error> {
let variables: Vec<(String, String)> = variables
.into_iter()
.filter_map(|(key, value)| Some((key.into_string().ok()?, value.into_string().ok()?)))
.collect();
if !variables.iter().any(|(key, _)| key == UPSTREAM) {
return Ok(None);
}
let environment: Environment = envy::from_iter(variables).map_err(|error| match error {
envy::Error::MissingValue(field) => Error::MissingSetting(field.to_uppercase()),
envy::Error::Custom(message) => Error::InvalidSetting(message),
})?;
Ok(Some(Self {
upstream: environment.litellm_forward_upstream.parse()?,
host: environment.litellm_forward_host,
port: environment.litellm_forward_port,
upstream_keepalive: positive_seconds(
"LITELLM_FORWARD_UPSTREAM_KEEPALIVE",
environment.litellm_forward_upstream_keepalive,
)?,
shutdown_drain: seconds(
"LITELLM_FORWARD_SHUTDOWN_DRAIN",
environment.litellm_forward_shutdown_drain,
)?,
}))
}
}
fn seconds(variable: &str, value: f64) -> Result<Duration, Error> {
Duration::try_from_secs_f64(value).map_err(|_| {
Error::InvalidSetting(format!(
"{variable} must be a non-negative number of seconds, got {value}"
))
})
}
/// The keep-alive also bounds the wait for a client's first request, so zero would close every
/// client connection before it could send one.
fn positive_seconds(variable: &str, value: f64) -> Result<Duration, Error> {
match seconds(variable, value)? {
Duration::ZERO => Err(Error::InvalidSetting(format!(
"{variable} must be a positive number of seconds, got {value}"
))),
duration => Ok(duration),
}
}

86
litellm-rust/crates/gateway-forward/src/websocket.rs Normal file
View file

@ -0,0 +1,86 @@
use axum::{
body::Body,
extract::Request,
http::{Method, Version},
response::Response,
};
use hyper::{
body::Incoming,
upgrade::{OnUpgrade, Upgraded},
};
use hyper_util::rt::TokioIo;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use crate::{Rejection, connections::Connections, rewrite};
pub(crate) fn take_handshake(
connections: &Connections,
request: &mut Request,
) -> Result<Option<OnUpgrade>, Rejection> {
if request.method() != Method::GET
|| request.version() != Version::HTTP_11
|| !rewrite::is_websocket_handshake(request.headers())
{
return Ok(None);
}
if connections.is_draining() {
return Err(Rejection::ShuttingDown);
}
Ok(Some(hyper::upgrade::on(request)))
}
pub(crate) async fn switch(
connections: &Connections,
downstream: OnUpgrade,
mut response: Response<Incoming>,
) -> Result<Response, Rejection> {
if !rewrite::is_websocket_handshake(response.headers()) {
return Err(Rejection::UnrequestedUpgrade);
}
let upstream = connections
.until_closed(hyper::upgrade::on(&mut response))
.await
.ok_or(Rejection::ShuttingDown)?
.map_err(Rejection::UpgradeFailed)?;
connections.spawn(async move {
match downstream.await {
Ok(downstream) => relay(downstream, upstream).await,
Err(_) => tracing::debug!("websocket upgrade did not complete"),
}
});
Ok(rewrite::switched_response(response).map(|_| Body::empty()))
}
// uvicorn never half-closes a WebSocket, so Python's EOF means the socket is gone and the relay
// ends. Only a client that is gone ends it from the other side: the client's EOF just half-closes
// Python's side, and a write Python refuses means it closed with its last frames still in flight
// to the client.
async fn relay(downstream: Upgraded, upstream: Upgraded) {
let (mut client_read, mut client_write) = tokio::io::split(TokioIo::new(downstream));
let (mut python_read, mut python_write) = tokio::io::split(TokioIo::new(upstream));
let to_python = async {
let mut buffer = vec![0; 8192];
loop {
let read = match client_read.read(&mut buffer).await {
Ok(0) => break,
Ok(read) => read,
Err(_) => return,
};
if python_write.write_all(&buffer[..read]).await.is_err() {
return std::future::pending().await;
}
}
python_write.shutdown().await.ok();
std::future::pending().await
};
let to_client = async {
tokio::io::copy(&mut python_read, &mut client_write)
.await
.ok();
client_write.shutdown().await.ok();
};
tokio::select! {
() = to_python => tracing::debug!("websocket client disconnected"),
() = to_client => {}
}
}

223
litellm-rust/crates/gateway-forward/tests/client_identity.rs Normal file
View file

@ -0,0 +1,223 @@
mod support;
use std::sync::{
Arc,
atomic::{AtomicUsize, Ordering},
};
use axum::{body::Body, http::StatusCode, response::Response};
use litellm_gateway_forward::{Forward, PEER_HEADER};
use rstest::rstest;
use support::{DRAIN, Gateway, Raw, UVICORN_KEEPALIVE, raw_upstream, recording, within};
use tokio::net::{TcpListener, TcpStream};
fn ok() -> Response {
Response::new(Body::from("ok"))
}
#[rstest]
#[case::ordinary("")]
#[case::connection_nominated(
"connection: X-Forwarded-For, X-Forwarded-Proto, Forwarded, X-Litellm-Rust-Peer\r\n"
)]
#[tokio::test]
async fn python_sees_forwarding_headers_untouched_and_the_tcp_peer_rust_saw(
#[case] connection: &str,
) {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(format!(
"GET /sso/callback HTTP/1.1\r\n\
host: public.example\r\n\
x-forwarded-for: 192.0.2.1\r\n\
x-forwarded-for: 192.0.2.2\r\n\
x-forwarded-proto: https\r\n\
forwarded: for=192.0.2.1;proto=https\r\n\
x-litellm-rust-peer: 10.1.2.3:1\r\n{connection}\r\n",
))
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
let request = received.recv().await.unwrap();
let forwarded_for: Vec<_> = request.headers.get_all("x-forwarded-for").iter().collect();
assert_eq!(forwarded_for, ["192.0.2.1", "192.0.2.2"]);
assert_eq!(request.headers["x-forwarded-proto"], "https");
assert_eq!(request.headers["forwarded"], "for=192.0.2.1;proto=https");
let peers: Vec<_> = request.headers.get_all(PEER_HEADER).iter().collect();
assert_eq!(peers, [raw.local_address().to_string().as_str()]);
}
/// Python builds redirects and OAuth metadata from Host, and without one from the address it
/// listens on, which behind Rust is its private loopback port. So a request without a usable Host
/// carries the public address the client connected to, except HTTP/1.1, which h11 rejects.
#[rstest]
#[case::http11_without_host("HTTP/1.1", "", false)]
#[case::http10_without_host("HTTP/1.0", "", true)]
#[case::http11_with_an_empty_host("HTTP/1.1", "host:\r\n", true)]
#[tokio::test]
async fn a_request_without_a_usable_host_never_names_the_private_upstream(
#[case] version: &str,
#[case] host: &str,
#[case] forwarded: bool,
) {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(format!("GET /sso/key/generate {version}\r\n{host}\r\n"))
.await;
let status = within(raw.response_head()).await.unwrap().status;
let hosts = received.try_recv().ok().map(|request| {
request
.headers
.get_all("host")
.iter()
.map(|value| value.to_str().unwrap().to_owned())
.collect::<Vec<_>>()
});
let expected = if forwarded {
(StatusCode::OK, Some(vec![gateway.address.to_string()]))
} else {
(StatusCode::BAD_REQUEST, None)
};
assert_eq!((status, hosts), expected);
}
#[rstest]
#[case::uppercase_with_port("PUBLIC.Example:4000")]
#[case::ipv6_literal("[::1]:4000")]
#[case::ipv4_with_port("192.0.2.10:8443")]
#[tokio::test]
async fn python_sees_the_host_header_byte_for_byte(#[case] host: &str) {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(format!("GET /ui HTTP/1.1\r\nHost: {host}\r\n\r\n"))
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
let request = received.recv().await.unwrap();
let hosts: Vec<_> = request.headers.get_all("host").iter().collect();
assert_eq!(hosts, [host]);
}
/// Python's IP allowlists and rate limits read the peer exactly as uvicorn would have seen it on
/// the public socket, so a dual-stack listener reports IPv4 clients in their mapped form.
#[rstest]
#[case::ipv6("[::1]:0", "[::1]", "[::1]")]
#[case::ipv4_on_a_dual_stack_listener("[::]:0", "127.0.0.1", "[::ffff:127.0.0.1]")]
#[tokio::test]
async fn ipv6_listeners_report_the_peer_bracketed_with_its_port(
#[case] bind: &str,
#[case] connect: &str,
#[case] peer_ip: &str,
) {
let (upstream, mut received) = recording(ok).await;
let forward = Forward::new(
&format!("http://{upstream}").parse().unwrap(),
UVICORN_KEEPALIVE,
);
let public = TcpListener::bind(bind).await.unwrap();
let port = public.local_addr().unwrap().port();
tokio::spawn(forward.serve(public, std::future::pending(), DRAIN));
let mut raw = Raw::new(
TcpStream::connect(format!("{connect}:{port}"))
.await
.unwrap(),
);
raw.send("GET / HTTP/1.1\r\nhost: public.example\r\n\r\n")
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
assert_eq!(
received.recv().await.unwrap().headers[PEER_HEADER],
format!("{peer_ip}:{}", raw.local_address().port()).as_str()
);
}
/// Without the peer header Python would take Rust's loopback hop as the client, which allowlists
/// often trust, so no spelling, repetition or Connection token may remove or replace it.
#[rstest]
#[case::named_as_hop_by_hop(
"connection: x-litellm-rust-peer\r\nx-litellm-rust-peer: 127.0.0.1:1\r\n"
)]
#[case::mixed_case_copies(
"X-LiteLLM-Rust-Peer: 127.0.0.1:1\r\nx-litellm-rust-peer: 192.0.2.7:2\r\n"
)]
#[case::comma_joined("x-litellm-rust-peer: 127.0.0.1:1, 192.0.2.7:2\r\n")]
#[tokio::test]
async fn a_client_can_neither_strip_nor_spoof_the_peer(#[case] spoof: &str) {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(format!(
"GET /key/info HTTP/1.1\r\nhost: public.example\r\n{spoof}\r\n"
))
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
let request = received.recv().await.unwrap();
let peers: Vec<_> = request.headers.get_all(PEER_HEADER).iter().collect();
assert_eq!(peers, [raw.local_address().to_string().as_str()]);
}
/// Pooled upstream connections carry requests from many clients in turn, so the peer belongs to
/// each request and never to the connection it travels on.
#[rstest]
#[tokio::test]
async fn each_request_carries_its_own_peer_over_a_shared_upstream_connection() {
let connections = Arc::new(AtomicUsize::new(0));
let accepted = connections.clone();
let upstream = raw_upstream(move |mut python| {
accepted.fetch_add(1, Ordering::SeqCst);
async move {
while let Some(head) = python.request_head().await {
let peer = head.header(PEER_HEADER).unwrap().to_vec();
python
.send(
[
format!("HTTP/1.1 200 OK\r\ncontent-length: {}\r\n\r\n", peer.len())
.as_bytes(),
&peer,
]
.concat(),
)
.await;
}
}
})
.await;
let gateway = Gateway::start(upstream).await;
let mut clients = [
gateway.connect().await,
gateway.connect().await,
gateway.connect().await,
];
for index in [0, 1, 2, 2, 1, 0] {
let raw = &mut clients[index];
raw.send("GET /v1/models HTTP/1.1\r\nhost: public.example\r\n\r\n")
.await;
let head = within(raw.response_head()).await.unwrap();
let peer = raw.body(head.content_length().unwrap()).await;
assert_eq!(peer, raw.local_address().to_string().into_bytes());
}
assert_eq!(connections.load(Ordering::SeqCst), 1);
}

192
litellm-rust/crates/gateway-forward/tests/download.rs Normal file
View file

@ -0,0 +1,192 @@
mod support;
use std::{
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
time::Duration,
};
use axum::{body::Body, extract::Request, http::StatusCode, response::Response};
use http_body_util::BodyExt;
use rstest::rstest;
use support::{Gateway, client, raw_upstream, recording, within};
use tokio::sync::{Notify, mpsc};
#[rstest]
#[tokio::test]
async fn streams_a_body_without_a_content_length_as_python_produces_it() {
let first: &[u8] = &[0xff, 0xfb, 0x90, 0x64, 0x00, 0x0d, 0x0a];
let rest = vec![0xa5; 64 * 1024];
let finish = Arc::new(Notify::new());
let finished = finish.clone();
let tail = rest.clone();
let upstream = raw_upstream(move |mut python| {
let finished = finished.clone();
let tail = tail.clone();
async move {
python.request_head().await.unwrap();
python
.send([&b"HTTP/1.1 200 OK\r\ncontent-type: audio/mpeg\r\ntransfer-encoding: chunked\r\n\r\n7\r\n"[..], first, b"\r\n"].concat())
.await;
finished.notified().await;
python
.send([format!("{:x}\r\n", tail.len()).as_bytes(), &tail, b"\r\n0\r\n\r\n"].concat())
.await;
}
})
.await;
let gateway = Gateway::start(upstream).await;
let mut response = within(
client().request(
Request::post(gateway.url("/v1/audio/speech"))
.body(Body::empty())
.unwrap(),
),
)
.await
.unwrap();
assert!(!response.headers().contains_key("content-length"));
let frame = within(response.body_mut().frame()).await.unwrap().unwrap();
assert_eq!(frame.into_data().unwrap(), first);
finish.notify_one();
let body = within(response.into_body().collect())
.await
.unwrap()
.to_bytes();
assert_eq!(body, rest);
}
#[rstest]
#[tokio::test]
async fn range_requests_and_partial_content_pass_unchanged() {
let (upstream, mut received) = recording(|| -> Response {
Response::builder()
.status(StatusCode::PARTIAL_CONTENT)
.header("content-range", "bytes 2-5/10")
.header("accept-ranges", "bytes")
.body(Body::from("2345"))
.unwrap()
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("GET /v1/videos/video_1/content HTTP/1.1\r\nhost: x\r\nrange: bytes=2-5\r\nif-range: \"v1\"\r\n\r\n")
.await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.status, StatusCode::PARTIAL_CONTENT);
assert_eq!(head.header("content-range"), Some(&b"bytes 2-5/10"[..]));
assert_eq!(raw.body(head.content_length().unwrap()).await, b"2345");
let request = received.recv().await.unwrap();
assert_eq!(request.headers["range"], "bytes=2-5");
assert_eq!(request.headers["if-range"], "\"v1\"");
}
#[rstest]
#[tokio::test]
async fn a_large_download_is_paced_by_the_client_and_keeps_its_content_length() {
const LENGTH: usize = 128 * 1024 * 1024;
let written = Arc::new(AtomicUsize::new(0));
let writing = written.clone();
let upstream = raw_upstream(move |mut python| {
let writing = writing.clone();
async move {
python.request_head().await.unwrap();
python
.send(format!("HTTP/1.1 200 OK\r\ncontent-type: application/octet-stream\r\ncontent-length: {LENGTH}\r\n\r\n"))
.await;
let chunk = [0xa5; 64 * 1024];
while writing.load(Ordering::SeqCst) < LENGTH {
python.send(chunk).await;
writing.fetch_add(chunk.len(), Ordering::SeqCst);
}
}
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("GET /v1/files/file-1/content HTTP/1.1\r\nhost: x\r\n\r\n")
.await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.content_length(), Some(LENGTH));
assert_eq!(head.header("transfer-encoding"), None);
tokio::time::sleep(Duration::from_secs(1)).await;
assert!(
written.load(Ordering::SeqCst) < LENGTH / 4,
"the forwarder buffered a download the client was not reading"
);
let mut remaining = LENGTH;
while remaining > 0 {
remaining -= within(raw.body(remaining.min(1 << 20))).await.len();
}
}
/// The UI revalidates its static assets, and a 304 must carry exactly the validators Python sent,
/// with no length invented for a body it never has.
#[rstest]
#[tokio::test]
async fn a_conditional_request_gets_python_304_with_its_validators_only() {
let (event, mut events) = mpsc::unbounded_channel();
let upstream = raw_upstream(move |mut python| {
let event = event.clone();
async move {
while let Some(head) = python.request_head().await {
event
.send(head.header("if-none-match").map(<[u8]>::to_vec))
.unwrap();
python
.send("HTTP/1.1 304 Not Modified\r\netag: \"5f1-abc\"\r\ncache-control: no-cache\r\nvary: Accept-Encoding\r\ndate: Tue, 03 Nov 2026 10:00:00 GMT\r\n\r\n")
.await;
}
}
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
for _ in 0..2 {
raw.send("GET /ui/_next/static/chunks/main.js HTTP/1.1\r\nhost: x\r\nif-none-match: \"5f1-abc\"\r\n\r\n")
.await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.status, StatusCode::NOT_MODIFIED);
let mut names: Vec<_> = head.headers.iter().map(|(name, _)| name.as_str()).collect();
names.sort_unstable();
assert_eq!(names, ["cache-control", "date", "etag", "vary"]);
assert_eq!(head.header("etag"), Some(&b"\"5f1-abc\""[..]));
assert_eq!(
head.header("date"),
Some(&b"Tue, 03 Nov 2026 10:00:00 GMT"[..])
);
assert_eq!(
within(events.recv()).await.unwrap().as_deref(),
Some(&b"\"5f1-abc\""[..])
);
}
}
/// A client knows a download was cut short only because fewer bytes than the Content-Length
/// arrived before the connection closed, so the length must survive and the connection must end.
#[rstest]
#[tokio::test]
async fn a_download_python_cuts_short_closes_the_client_connection_before_its_length() {
let upstream = raw_upstream(|mut python| async move {
python.request_head().await.unwrap();
python
.send([&b"HTTP/1.1 200 OK\r\ncontent-type: application/octet-stream\r\ncontent-length: 1000\r\n\r\n"[..], &[0xa5; 500]].concat())
.await;
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("GET /v1/files/file-1/content HTTP/1.1\r\nhost: x\r\n\r\n")
.await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.status, StatusCode::OK);
assert_eq!(head.content_length(), Some(1000));
assert_eq!(within(raw.until_closed()).await, [0xa5; 500]);
}

855
litellm-rust/crates/gateway-forward/tests/exchange.rs Normal file
View file

@ -0,0 +1,855 @@
mod support;
use std::{
convert::Infallible,
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
};
use axum::{
Router,
body::{Body, Bytes},
extract::Request,
http::{HeaderValue, StatusCode},
response::Response,
};
use litellm_gateway_forward::PEER_HEADER;
use rstest::rstest;
use support::{Gateway, Head, client, raw_upstream, recording, serve, within};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
sync::mpsc,
};
fn ok() -> Response {
Response::new(Body::from("ok"))
}
/// An upstream that reports each bodiless request head exactly as it arrived on the wire.
async fn heads() -> (std::net::SocketAddr, mpsc::UnboundedReceiver<Head>) {
let (sender, receiver) = mpsc::unbounded_channel();
let address = raw_upstream(move |mut python| {
let sender = sender.clone();
async move {
while let Some(head) = python.request_head().await {
sender.send(head).unwrap();
python
.send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok")
.await;
}
}
})
.await;
(address, receiver)
}
#[rstest]
#[tokio::test]
async fn forwards_the_request_as_the_client_sent_it_minus_hop_headers() {
let (upstream, mut received) = recording(|| {
Response::builder()
.status(StatusCode::UNAUTHORIZED)
.header("connection", "x-upstream-hop")
.header("x-upstream-hop", "remove")
.header("x-response", "keep")
.body(Body::from("python error"))
.unwrap()
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(
"PATCH /v1/messages?x=a%2Fb&x=c HTTP/1.1\r\n\
host: public.example\r\n\
authorization: Bearer test\r\n\
connection: x-remove, keep-alive\r\n\
x-remove: remove\r\n\
proxy-authorization: remove\r\n\
te: trailers\r\n\
content-length: 7\r\n\r\npayload",
)
.await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.status, StatusCode::UNAUTHORIZED);
assert_eq!(head.header("x-response"), Some(&b"keep"[..]));
assert_eq!(head.header("x-upstream-hop"), None);
let body = raw.body(head.content_length().unwrap()).await;
assert_eq!(body, b"python error");
let request = received.recv().await.unwrap();
assert_eq!(request.method, "PATCH");
assert_eq!(request.uri, "/v1/messages?x=a%2Fb&x=c");
assert_eq!(request.headers["host"], "public.example");
assert_eq!(request.headers["authorization"], "Bearer test");
assert!(!request.headers.contains_key("x-remove"));
assert!(!request.headers.contains_key("proxy-authorization"));
assert!(!request.headers.contains_key("te"));
assert!(!request.headers.contains_key("connection"));
assert_eq!(request.body, "payload");
}
#[rstest]
#[case::host("HoSt", "public.example")]
#[case::length("Content-Length", "7")]
#[case::authorization("AUTHORIZATION", "Bearer test")]
#[case::cookie("Cookie", "session=test")]
#[case::api_key("X-Api-Key", "test")]
#[case::azure_key("Api-Key", "test")]
#[case::litellm_key("Litellm-Api-Key", "test")]
#[case::forwarded("Forwarded", "for=192.0.2.1;proto=https")]
#[case::forwarded_for("X-Forwarded-For", "192.0.2.1")]
#[case::forwarded_proto("X-Forwarded-Proto", "https")]
#[case::forwarded_host("X-Forwarded-Host", "public.example")]
#[case::content_type("Content-Type", "application/octet-stream")]
#[case::content_encoding("Content-Encoding", "identity")]
#[tokio::test]
async fn connection_cannot_remove_end_to_end_headers(#[case] name: &str, #[case] value: &str) {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
let host = if name.eq_ignore_ascii_case("host") {
""
} else {
"host: public.example\r\n"
};
let length = if name.eq_ignore_ascii_case("content-length") {
""
} else {
"content-length: 7\r\n"
};
raw.send(format!(
"POST / HTTP/1.1\r\n{host}{length}{name}: {value}\r\nconnection: x-remove\r\nconnection: {name}, keep-alive\r\nx-remove: remove\r\n\r\npayload"
)).await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
let request = within(received.recv()).await.unwrap();
assert_eq!(request.headers[name], value);
assert_eq!(request.body, "payload");
assert!(!request.headers.contains_key("connection"));
assert!(!request.headers.contains_key("x-remove"));
}
#[rstest]
#[tokio::test]
async fn connection_cannot_remove_response_credentials_or_representation_metadata() {
let (upstream, _received) = recording(|| {
Response::builder()
.status(StatusCode::UNAUTHORIZED)
.header("connection", "Content-Length, Set-Cookie, WWW-Authenticate, Authentication-Info, Content-Type, Content-Encoding, x-remove")
.header("content-length", "2")
.header("set-cookie", "session=; Max-Age=0")
.header("www-authenticate", "Bearer")
.header("authentication-info", "nextnonce=test")
.header("content-type", "text/plain")
.header("content-encoding", "identity")
.header("x-remove", "remove")
.body(Body::from("ok"))
.unwrap()
}).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("GET / HTTP/1.1\r\nhost: x\r\n\r\n").await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.status, StatusCode::UNAUTHORIZED);
assert_eq!(head.content_length(), Some(2));
assert_eq!(head.header("set-cookie"), Some(&b"session=; Max-Age=0"[..]));
assert_eq!(head.header("www-authenticate"), Some(&b"Bearer"[..]));
assert_eq!(
head.header("authentication-info"),
Some(&b"nextnonce=test"[..])
);
assert_eq!(head.header("content-type"), Some(&b"text/plain"[..]));
assert_eq!(head.header("content-encoding"), Some(&b"identity"[..]));
assert_eq!(head.header("x-remove"), None);
assert_eq!(raw.body(2).await, b"ok");
}
#[rstest]
#[case::get("GET")]
#[case::post("POST")]
#[case::put("PUT")]
#[case::delete("DELETE")]
#[case::options("OPTIONS")]
#[case::extension("PURGE")]
#[case::webdav("PROPFIND")]
#[case::lowercase("get")]
#[case::punctuated("M-SEARCH")]
#[tokio::test]
async fn forwards_every_method(#[case] method: &str) {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let response = client()
.request(
Request::builder()
.method(method)
.uri(gateway.url("/mcp"))
.body(Body::from("body"))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let request = received.recv().await.unwrap();
assert_eq!(request.method, method);
assert_eq!(request.body, "body");
}
#[rstest]
#[case::double_slash("//v1//models")]
#[case::encoded_slash("/v1/files/a%2Fb/content")]
#[case::encoded_hash_in_query("/v1/files?purpose=batch%23retry&limit=1")]
#[case::encoded_hash_in_path("/v1/files/report%231.pdf/content")]
#[case::encoded_hash_before_query("/gemini/v1beta/files/a%23b?key=k")]
#[case::encoded_query("/v1/models?q=%E2%9C%93&empty=&flag")]
#[case::utf8_bytes_escaped("/mcp/%F0%9F%A6%80")]
#[case::lowercase_encoded_slash("/v1/files/a%2fb/content")]
#[case::dot_segments("/openai/./v1/../v1/models")]
#[case::encoded_dot_segments("/openai/%2e%2e/v1/%2E/models")]
#[case::matrix_params("/langfuse/api;v=1/public;x=y/traces")]
#[case::plus_and_encoded_space("/gemini/v1beta/models?q=a+b&r=a%20b&s=%2B")]
#[case::empty_query("/v1/models?")]
#[case::question_marks_in_query("/v1/models??a=b?c")]
#[case::colons_in_segments(
"/v1/projects/p/locations/l/publishers/google/models/m:streamGenerateContent?alt=sse"
)]
#[case::sub_delims("/~u/@:!$&'()*+,=/x")]
#[case::raw_bytes_browsers_leave_unescaped("/a{b}|c^d\\e\"f?g={}|^`")]
#[case::invalid_percent_escapes("/a%zz%00b?%")]
#[tokio::test]
async fn forwards_the_request_target_byte_for_byte(#[case] target: &str) {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(format!("GET {target} HTTP/1.1\r\nhost: x\r\n\r\n"))
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
assert_eq!(received.recv().await.unwrap().uri, target);
}
#[rstest]
#[tokio::test]
async fn forwards_a_query_string_as_long_as_uvicorn_accepts() {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
let target = format!("/vertex_ai/v1/models?key=%2F{}&end", "a%20b".repeat(3000));
raw.send(format!("GET {target} HTTP/1.1\r\nhost: x\r\n\r\n"))
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
assert_eq!(received.recv().await.unwrap().uri, target.as_str());
}
#[rstest]
#[case::in_query("/v1/files?purpose=batch#retry&limit=1", "/v1/files?purpose=batch")]
#[case::in_path("/v1/files/report#1.pdf/content", "/v1/files/report")]
#[case::before_query("/gemini/v1beta/files/a#b?key=k", "/gemini/v1beta/files/a")]
#[tokio::test]
async fn forwards_only_the_path_and_query_before_a_fragment(
#[case] target: &str,
#[case] expected: &str,
) {
let (upstream, mut received) = heads().await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(format!("GET {target} HTTP/1.1\r\nhost: x\r\n\r\n"))
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
assert_eq!(within(received.recv()).await.unwrap().path, expected);
}
#[rstest]
#[tokio::test]
async fn absolute_form_targets_reach_python_as_origin_form() {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(
"GET http://elsewhere.example/v1/models?x=1 HTTP/1.1\r\nhost: elsewhere.example\r\n\r\n",
)
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
let request = received.recv().await.unwrap();
assert_eq!(request.uri, "/v1/models?x=1");
}
#[rstest]
#[tokio::test]
async fn relays_a_response_with_more_headers_than_hyper_allows_by_default() {
let upstream = serve(axum::Router::new().fallback(|| async {
let headers = (0..150).map(|index| (format!("llm_provider-x-{index}"), "v"));
let mut response = ok();
for (name, value) in headers {
response.headers_mut().insert(
axum::http::HeaderName::from_bytes(name.as_bytes()).unwrap(),
HeaderValue::from_static(value),
);
}
response
}))
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("GET /v1/chat/completions HTTP/1.1\r\nhost: public.example\r\n\r\n")
.await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.status, StatusCode::OK);
assert_eq!(head.header("llm_provider-x-149"), Some(&b"v"[..]));
}
/// uvicorn's h11 parser has no header count limit, only a 16 KiB cap on the whole head.
#[rstest]
#[tokio::test]
async fn forwards_a_request_with_as_many_headers_as_uvicorn_accepts() {
let (sender, mut received) = mpsc::unbounded_channel();
let upstream = raw_upstream(move |python| {
let sender = sender.clone();
async move {
let mut stream = python.into_stream();
let mut head = Vec::new();
while !head.ends_with(b"\r\n\r\n") {
head.push(stream.read_u8().await.unwrap());
}
sender.send(head).unwrap();
stream
.write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 0\r\n\r\n")
.await
.unwrap();
}
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
let headers: String = (0..1000).map(|index| format!("x-{index}: v\r\n")).collect();
raw.send(format!(
"GET /v1/chat/completions HTTP/1.1\r\nhost: x\r\n{headers}\r\n"
))
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
let head = String::from_utf8(within(received.recv()).await.unwrap()).unwrap();
assert!(head.contains(&headers));
}
async fn header_fields_python_receives(fields: &str) -> Vec<(String, Vec<u8>)> {
let (upstream, mut received) = heads().await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(format!(
"GET /anthropic/v1/messages HTTP/1.1\r\nhost: x\r\n{fields}\r\n"
))
.await;
within(raw.response_head()).await.unwrap();
within(received.recv())
.await
.unwrap()
.headers
.into_iter()
.filter(|(name, _)| name.starts_with("x-") && name != PEER_HEADER)
.collect()
}
fn fields(expected: &[(&str, &str)]) -> Vec<(String, Vec<u8>)> {
expected
.iter()
.map(|(name, value)| ((*name).to_owned(), value.as_bytes().to_vec()))
.collect()
}
/// Python reads `scope["headers"]` as the list the client sent, in that order.
#[rstest]
#[ignore = "divergence: HeaderMap::remove can move the headers after a dropped hop-by-hop one"]
#[tokio::test]
async fn dropping_a_hop_by_hop_header_keeps_the_others_in_order() {
let received = header_fields_python_receives(
"x-a: 1\r\nconnection: keep-alive\r\nx-b: 2\r\nx-c: 3\r\nx-d: 4\r\n",
)
.await;
assert_eq!(
received,
fields(&[("x-a", "1"), ("x-b", "2"), ("x-c", "3"), ("x-d", "4")])
);
}
/// uvicorn speaks only HTTP/1.1, so an HTTP/2 prior-knowledge client never reaches Python through
/// Rust either. hyper closes the connection where h11 answers 400, which no HTTP/2 client reads.
#[rstest]
#[tokio::test]
async fn an_http2_client_never_reaches_python() {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n\x00\x00\x00\x04\x00\x00\x00\x00\x00")
.await;
assert!(within(raw.until_closed()).await.is_empty());
assert!(received.try_recv().is_err());
}
#[rstest]
#[ignore = "divergence: HeaderMap groups every value of a name at its first position"]
#[tokio::test]
async fn interleaved_repeated_headers_reach_python_interleaved() {
let received = header_fields_python_receives("x-a: 1\r\nx-b: 2\r\nx-a: 3\r\n").await;
assert_eq!(
received,
fields(&[("x-a", "1"), ("x-b", "2"), ("x-a", "3")])
);
}
/// h11 trims optional whitespace around a value and keeps everything inside it, including empty
/// values and names using every token character.
#[rstest]
#[tokio::test]
async fn python_sees_header_values_as_h11_would_parse_them() {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("GET / HTTP/1.1\r\nhost: x\r\nx-empty:\r\nx-blank: \r\nx-inner: a \t b \t\r\nX-Token!#$%&'*+.^_`|~: v\r\nx-empty: second\r\n\r\n")
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
let request = received.recv().await.unwrap();
let values = |name: &str| -> Vec<&[u8]> {
request
.headers
.get_all(name)
.iter()
.map(HeaderValue::as_bytes)
.collect()
};
assert_eq!(values("x-empty"), [&b""[..], b"second"]);
assert_eq!(values("x-blank"), [&b""[..]]);
assert_eq!(values("x-inner"), [&b"a \t b"[..]]);
assert_eq!(values("x-token!#$%&'*+.^_`|~"), [&b"v"[..]]);
}
/// RFC 9112 section 5.2 lets a server either reject obs-fold or unfold it, and h11 unfolds.
#[rstest]
#[ignore = "divergence: hyper rejects obs-fold with a 400 where uvicorn unfolds it to one space"]
#[tokio::test]
async fn a_folded_header_reaches_python_unfolded() {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("GET / HTTP/1.1\r\nhost: x\r\nx-a: b\r\n c\r\n\r\n")
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
assert_eq!(received.recv().await.unwrap().headers["x-a"], "b c");
}
#[rstest]
#[tokio::test]
async fn keeps_repeated_response_headers_and_opaque_header_bytes() {
let (upstream, mut received) = recording(|| {
Response::builder()
.header("set-cookie", "a=1; Path=/")
.header("set-cookie", "b=2; HttpOnly")
.header("x-opaque", HeaderValue::from_bytes(b"caf\xe9").unwrap())
.body(Body::empty())
.unwrap()
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(b"GET / HTTP/1.1\r\nhost: x\r\nx-opaque: caf\xe9\r\n\r\n")
.await;
let head = within(raw.response_head()).await.unwrap();
let cookies: Vec<_> = head
.headers
.iter()
.filter(|(name, _)| name == "set-cookie")
.map(|(_, value)| value.as_slice())
.collect();
assert_eq!(cookies, [&b"a=1; Path=/"[..], b"b=2; HttpOnly"]);
assert_eq!(head.header("x-opaque"), Some(&b"caf\xe9"[..]));
assert_eq!(
received.recv().await.unwrap().headers["x-opaque"].as_bytes(),
b"caf\xe9"
);
}
/// An HTTP/1.0 client cannot decode chunked framing, so a streamed answer reaches it delimited by
/// the connection closing, as uvicorn sends it.
#[rstest]
#[tokio::test]
async fn an_http10_client_gets_a_streamed_answer_unframed_and_python_gets_its_body() {
let (upstream, mut received) = recording(|| {
Response::new(Body::from_stream(futures_util::stream::iter(
["str", "eamed"].map(|piece| Ok::<_, Infallible>(Bytes::from_static(piece.as_bytes()))),
)))
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("POST /v1/completions HTTP/1.0\r\nhost: x\r\ncontent-length: 7\r\n\r\npayload")
.await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.status, StatusCode::OK);
assert_eq!(head.header("transfer-encoding"), None);
assert_eq!(within(raw.until_closed()).await, b"streamed");
let request = received.recv().await.unwrap();
assert_eq!(request.method, "POST");
assert_eq!(request.body, "payload");
}
#[rstest]
#[tokio::test]
async fn answers_pipelined_requests_in_order() {
let upstream = serve(
Router::new().fallback(|request: Request| async move { request.uri().path().to_owned() }),
)
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("GET /first HTTP/1.1\r\nhost: x\r\n\r\nHEAD /second HTTP/1.1\r\nhost: x\r\n\r\nGET /third HTTP/1.1\r\nhost: x\r\n\r\n")
.await;
let first = within(raw.response_head()).await.unwrap();
assert_eq!(raw.body(first.content_length().unwrap()).await, b"/first");
let second = within(raw.response_head()).await.unwrap();
assert_eq!(second.content_length(), Some(7));
let third = within(raw.response_head()).await.unwrap();
assert_eq!(raw.body(third.content_length().unwrap()).await, b"/third");
}
#[rstest]
#[case::no_content(StatusCode::NO_CONTENT)]
#[case::not_modified(StatusCode::NOT_MODIFIED)]
#[tokio::test]
async fn bodiless_statuses_keep_the_connection_reusable(#[case] status: StatusCode) {
let upstream = serve(Router::new().fallback(move |request: Request| async move {
match request.uri().path() {
"/next" => (StatusCode::OK, "next"),
_ => (status, ""),
}
}))
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("GET /empty HTTP/1.1\r\nhost: x\r\n\r\nGET /next HTTP/1.1\r\nhost: x\r\n\r\n")
.await;
assert_eq!(within(raw.response_head()).await.unwrap().status, status);
let next = within(raw.response_head()).await.unwrap();
assert_eq!(raw.body(next.content_length().unwrap()).await, b"next");
}
#[rstest]
#[tokio::test]
async fn informational_responses_keep_the_connection_reusable() {
let upstream = raw_upstream(|mut python| async move {
while let Some(head) = python.request_head().await {
match head.path.as_str() {
"/hints" => {
python
.send("HTTP/1.1 103 Early Hints\r\nlink: </ui/app.css>; rel=preload\r\n\r\nHTTP/1.1 200 OK\r\ncontent-length: 5\r\n\r\nfinal")
.await
}
_ => {
python
.send("HTTP/1.1 200 OK\r\ncontent-length: 4\r\n\r\nnext")
.await
}
}
}
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("GET /hints HTTP/1.1\r\nhost: x\r\n\r\nGET /next HTTP/1.1\r\nhost: x\r\n\r\n")
.await;
let mut hinted = within(raw.response_head()).await.unwrap();
while hinted.status.is_informational() {
hinted = within(raw.response_head()).await.unwrap();
}
assert_eq!(hinted.status, StatusCode::OK);
assert_eq!(raw.body(hinted.content_length().unwrap()).await, b"final");
let next = within(raw.response_head()).await.unwrap();
assert_eq!(raw.body(next.content_length().unwrap()).await, b"next");
}
#[rstest]
#[case::found(
StatusCode::FOUND,
"https://idp.example/authorize?state=a%2Fb&redirect_uri=https%3A%2F%2Fpublic.example%2Fsso%2Fcallback"
)]
#[case::see_other(StatusCode::SEE_OTHER, "/ui/?login=success")]
#[case::temporary(StatusCode::TEMPORARY_REDIRECT, "./mcp/oauth/authorize?client_id=x")]
#[tokio::test]
async fn redirects_reach_the_client_with_location_and_cookies_unchanged(
#[case] status: StatusCode,
#[case] location: &'static str,
) {
let (upstream, _received) = recording(move || {
Response::builder()
.status(status)
.header("location", location)
.header(
"set-cookie",
"token=a%2Fb; Path=/; HttpOnly; Secure; SameSite=Lax",
)
.header("set-cookie", "state=; Max-Age=0")
.header(
"set-cookie",
"session=x; Expires=Wed, 04 Nov 2026 10:00:00 GMT; Path=/",
)
.body(Body::empty())
.unwrap()
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("GET /sso/key/generate HTTP/1.1\r\nhost: public.example\r\n\r\n")
.await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.status, status);
assert_eq!(head.header("location"), Some(location.as_bytes()));
let cookies: Vec<_> = head
.headers
.iter()
.filter(|(name, _)| name == "set-cookie")
.map(|(_, value)| value.as_slice())
.collect();
assert_eq!(
cookies,
[
&b"token=a%2Fb; Path=/; HttpOnly; Secure; SameSite=Lax"[..],
b"state=; Max-Age=0",
b"session=x; Expires=Wed, 04 Nov 2026 10:00:00 GMT; Path=/",
]
);
}
#[rstest]
#[tokio::test]
async fn non_websocket_upgrades_reach_python_as_plain_requests() {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("GET / HTTP/1.1\r\nhost: x\r\nconnection: upgrade, http2-settings\r\nupgrade: h2c\r\nhttp2-settings: AAMAAABkAAQAoAAAAAIAAAAA\r\n\r\n")
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
let request = received.recv().await.unwrap();
assert!(!request.headers.contains_key("upgrade"));
assert!(!request.headers.contains_key("connection"));
assert!(!request.headers.contains_key("http2-settings"));
}
#[rstest]
#[tokio::test]
async fn forwards_an_asterisk_form_options_request() {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("OPTIONS * HTTP/1.1\r\nhost: x\r\n\r\n").await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
let request = received.recv().await.unwrap();
assert_eq!(request.method, "OPTIONS");
assert_eq!(request.uri, "*");
}
#[rstest]
#[tokio::test]
async fn an_authority_form_target_is_a_400_python_never_sees() {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("CONNECT api.example:443 HTTP/1.1\r\nhost: api.example:443\r\n\r\n")
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::BAD_REQUEST
);
assert!(received.try_recv().is_err());
}
#[rstest]
#[tokio::test]
async fn a_client_connection_close_ends_only_the_client_connection() {
let connections = Arc::new(AtomicUsize::new(0));
let accepted = connections.clone();
let upstream = raw_upstream(move |mut python| {
accepted.fetch_add(1, Ordering::SeqCst);
async move {
while python.request_head().await.is_some() {
python
.send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok")
.await;
}
}
})
.await;
let gateway = Gateway::start(upstream).await;
for _ in 0..2 {
let mut raw = gateway.connect().await;
raw.send("GET / HTTP/1.1\r\nhost: x\r\nconnection: close\r\n\r\n")
.await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.header("connection"), Some(&b"close"[..]));
assert_eq!(raw.body(2).await, b"ok");
within(raw.closed()).await;
}
assert_eq!(connections.load(Ordering::SeqCst), 1);
}
#[rstest]
#[tokio::test]
async fn a_python_connection_close_ends_only_the_upstream_connection() {
let connections = Arc::new(AtomicUsize::new(0));
let accepted = connections.clone();
let upstream = raw_upstream(move |mut python| {
accepted.fetch_add(1, Ordering::SeqCst);
async move {
python.request_head().await.unwrap();
python
.send("HTTP/1.1 500 Internal Server Error\r\ncontent-length: 5\r\nconnection: close\r\n\r\nerror")
.await;
}
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
for _ in 0..2 {
raw.send("GET / HTTP/1.1\r\nhost: x\r\n\r\n").await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.status, StatusCode::INTERNAL_SERVER_ERROR);
assert_eq!(head.header("connection"), None);
assert_eq!(raw.body(5).await, b"error");
}
assert_eq!(connections.load(Ordering::SeqCst), 2);
}
/// uvicorn answers HEAD on a streaming route with chunked framing and no body, so the client
/// must not be told the body is empty, and must not wait for one before its next request.
#[rstest]
#[tokio::test]
async fn head_answered_with_chunked_framing_claims_no_length_and_keeps_the_connection() {
let upstream = raw_upstream(|mut python| async move {
while let Some(head) = python.request_head().await {
match head.path.as_str() {
"/next" => {
python
.send("HTTP/1.1 200 OK\r\ncontent-length: 4\r\n\r\nnext")
.await
}
_ => {
python
.send("HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n")
.await
}
}
}
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("HEAD /mcp HTTP/1.1\r\nhost: x\r\n\r\nGET /next HTTP/1.1\r\nhost: x\r\n\r\n")
.await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.status, StatusCode::OK);
assert_eq!(head.content_length(), None);
let next = within(raw.response_head()).await.unwrap();
assert_eq!(raw.body(next.content_length().unwrap()).await, b"next");
}
const GZIP: &[u8] = &[
0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0x4b, 0x49, 0x2c, 0x49, 0xb4, 0x52,
0x30, 0xe4, 0xe2, 0x02, 0x00, 0xce, 0xcd, 0x78, 0x54, 0x09, 0x00, 0x00, 0x00,
];
#[rstest]
#[case::fixed_length(false)]
#[case::streamed(true)]
#[tokio::test]
async fn content_codings_cross_unchanged_both_ways(#[case] streamed: bool) {
let (upstream, mut received) = recording(move || {
let body = match streamed {
true => Body::from_stream(futures_util::stream::iter(
GZIP.chunks(5)
.map(|piece| Ok::<_, Infallible>(Bytes::from_static(piece))),
)),
false => Body::from(GZIP),
};
Response::builder()
.header("content-encoding", "gzip")
.header("vary", "Accept-Encoding")
.body(body)
.unwrap()
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(
[
format!(
"POST /v1/chat/completions HTTP/1.1\r\nhost: x\r\ncontent-encoding: gzip\r\naccept-encoding: gzip, br;q=0.5\r\ncontent-length: {}\r\n\r\n",
GZIP.len()
)
.as_bytes(),
GZIP,
]
.concat(),
)
.await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.header("content-encoding"), Some(&b"gzip"[..]));
assert_eq!(head.header("vary"), Some(&b"Accept-Encoding"[..]));
let body = match head.content_length() {
Some(length) => raw.body(length).await,
None => within(raw.chunked_body()).await,
};
assert_eq!(body, GZIP);
assert_eq!(head.content_length().is_none(), streamed);
let request = received.recv().await.unwrap();
assert_eq!(request.headers["content-encoding"], "gzip");
assert_eq!(request.headers["accept-encoding"], "gzip, br;q=0.5");
assert_eq!(request.body, GZIP);
}
#[rstest]
#[tokio::test]
async fn python_sees_no_accept_encoding_the_client_did_not_send() {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("GET /metrics/ HTTP/1.1\r\nhost: x\r\n\r\n").await;
within(raw.response_head()).await.unwrap();
assert!(
!received
.recv()
.await
.unwrap()
.headers
.contains_key("accept-encoding")
);
}

259
litellm-rust/crates/gateway-forward/tests/framing.rs Normal file
View file

@ -0,0 +1,259 @@
mod support;
use std::{collections::BTreeSet, net::SocketAddr};
use axum::http::StatusCode;
use rstest::rstest;
use support::{Gateway, Head, Raw, raw_upstream, within};
use tokio::sync::mpsc;
struct Captured {
head: Head,
body: Vec<u8>,
}
/// A keep-alive upstream that decodes framing like h11 and reports every request it reads.
async fn capturing() -> (SocketAddr, mpsc::UnboundedReceiver<Captured>) {
let (sender, receiver) = mpsc::unbounded_channel();
let address = raw_upstream(move |mut python: Raw| {
let sender = sender.clone();
async move {
while let Some(head) = python.request_head().await {
let body = match head.header("transfer-encoding") {
Some(_) => python.chunked_body().await,
None => python.body(content_length(&head)).await,
};
sender.send(Captured { head, body }).unwrap();
python
.send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok")
.await;
}
}
})
.await;
(address, receiver)
}
/// The length h11 reads from every Content-Length field, which may repeat one value.
fn content_length(head: &Head) -> usize {
let lengths: BTreeSet<usize> = head
.headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("content-length"))
.flat_map(|(_, value)| value.split(|byte| *byte == b','))
.map(|length| {
std::str::from_utf8(length.trim_ascii())
.unwrap()
.parse()
.unwrap()
})
.collect();
assert!(
lengths.len() < 2,
"python got conflicting lengths {lengths:?}"
);
lengths.first().copied().unwrap_or(0)
}
/// RFC 9112 section 6.3: Transfer-Encoding overrides Content-Length. Forwarding both would let
/// Python split the stream at a different place than Rust did. Section 6.1 then closes the
/// connection after the response, so a request smuggled behind it never runs.
#[rstest]
#[tokio::test]
async fn a_request_framed_two_ways_reaches_python_framed_by_transfer_encoding_and_closes() {
let (upstream, mut received) = capturing().await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: 3\r\ntransfer-encoding: chunked\r\n\r\n5\r\nhello\r\n0\r\n\r\nGET /next HTTP/1.1\r\nhost: x\r\n\r\n")
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
raw.body(2).await;
within(raw.until_closed()).await;
let upload = received.recv().await.unwrap();
assert_eq!(upload.head.header("content-length"), None);
assert_eq!(
upload.head.header("transfer-encoding"),
Some(&b"chunked"[..])
);
assert_eq!(upload.body, b"hello");
assert!(
received.try_recv().is_err(),
"python ran the smuggled request"
);
}
#[rstest]
#[case::conflicting_content_lengths("content-length: 4\r\ncontent-length: 40\r\n")]
#[case::content_length_list("content-length: 4, 40\r\n")]
#[case::signed_content_length("content-length: +4\r\n")]
#[case::chunked_not_last("transfer-encoding: chunked, identity\r\ncontent-length: 4\r\n")]
#[case::unknown_coding("transfer-encoding: xchunked\r\ncontent-length: 4\r\n")]
#[case::space_before_colon("transfer-encoding : chunked\r\ncontent-length: 4\r\n")]
#[case::folded_transfer_encoding("transfer-encoding:\r\n chunked\r\ncontent-length: 4\r\n")]
#[case::bare_cr("x-a: b\rx-b: c\r\ncontent-length: 4\r\n")]
#[case::nul_byte("x-a: b\x00c\r\ncontent-length: 4\r\n")]
#[tokio::test]
async fn ambiguous_framing_never_reaches_python(#[case] headers: &str) {
let (upstream, mut received) = capturing().await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(format!(
"POST /v1/files HTTP/1.1\r\nhost: x\r\n{headers}\r\nbodyGET /smuggled HTTP/1.1\r\nhost: x\r\n\r\n"
))
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::BAD_REQUEST
);
within(raw.closed()).await;
assert!(received.try_recv().is_err(), "python received a request");
}
#[rstest]
#[case::bare_lf("4\nbody\n0\n\n")]
#[case::invalid_size("zz\r\nbody\r\n0\r\n\r\n")]
#[case::size_overflow("10000000000000004\r\nbody\r\n0\r\n\r\n")]
#[tokio::test]
async fn a_malformed_request_body_is_a_400_not_a_502(#[case] chunks: &str) {
let upstream = raw_upstream(|mut python| async move { python.closed().await }).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(format!(
"POST /v1/files HTTP/1.1\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n{chunks}"
))
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::BAD_REQUEST
);
}
/// hyper decodes only the chunked layer, so any other coding would reach Python still coded but
/// framed as a plain body. uvicorn turns every h11 parse error into a 400, ignoring the 501 h11
/// suggests for codings.
#[rstest]
#[case::listed_before_chunked("transfer-encoding: gzip, chunked\r\n")]
#[case::split_across_fields("transfer-encoding: gzip\r\ntransfer-encoding: chunked\r\n")]
#[tokio::test]
async fn a_transfer_coding_other_than_chunked_is_the_400_uvicorn_sends_python_never_sees(
#[case] headers: &str,
) {
let (upstream, mut received) = capturing().await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(format!(
"POST /v1/files HTTP/1.1\r\nhost: x\r\n{headers}\r\n4\r\nbody\r\n0\r\n\r\n"
))
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::BAD_REQUEST
);
assert!(
received.try_recv().is_err(),
"python received a still-coded body"
);
}
async fn assert_python_reads_the_body_and_the_next_request(headers: &str) {
let (upstream, mut received) = capturing().await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(format!(
"POST /v1/files HTTP/1.1\r\nhost: x\r\n{headers}\r\npayloadGET /next HTTP/1.1\r\nhost: x\r\n\r\n"
))
.await;
for _ in 0..2 {
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
raw.body(2).await;
}
assert_eq!(received.recv().await.unwrap().body, b"payload");
assert_eq!(received.recv().await.unwrap().head.path, "/next");
}
#[rstest]
#[case::repeated_fields("content-length: 7\r\ncontent-length: 7\r\n")]
#[case::leading_zeros("content-length: 007\r\n")]
#[tokio::test]
async fn a_content_length_h11_reads_one_way_reaches_python_framed_the_same(#[case] headers: &str) {
assert_python_reads_the_body_and_the_next_request(headers).await;
}
/// h11 accepts a Content-Length list that repeats one value, as RFC 9110 section 8.6 allows.
#[rstest]
#[ignore = "divergence: hyper's server parses each Content-Length field as one number and 400s a list"]
#[tokio::test]
async fn a_content_length_list_of_one_value_reaches_python() {
assert_python_reads_the_body_and_the_next_request("content-length: 7, 7\r\n").await;
}
#[rstest]
#[case::leading_zeros("0004\r\nbody\r\n0\r\n\r\n", "body")]
#[case::uppercase_hex("A\r\nbody\r\nbody\r\n0\r\n\r\n", "body\r\nbody")]
#[case::whitespace_after_size("4 \r\nbody\r\n0\r\n\r\n", "body")]
#[case::extension_and_trailer("4;a=\"b;c\"\r\nbody\r\n0\r\nx-sum: 1\r\n\r\n", "body")]
#[tokio::test]
async fn chunk_framing_h11_accepts_reaches_python_decoded(
#[case] chunks: &str,
#[case] body: &str,
) {
let (upstream, mut received) = capturing().await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(format!(
"POST /v1/files HTTP/1.1\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n{chunks}GET /next HTTP/1.1\r\nhost: x\r\n\r\n"
))
.await;
for _ in 0..2 {
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
raw.body(2).await;
}
assert_eq!(received.recv().await.unwrap().body, body.as_bytes());
assert_eq!(received.recv().await.unwrap().head.path, "/next");
}
/// h11 reads a head whose lines end in a bare LF, so Python gets the same request through Rust.
#[rstest]
#[tokio::test]
async fn a_head_with_bare_lf_line_endings_reaches_python() {
let (upstream, mut received) = capturing().await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("POST /v1/files HTTP/1.1\nhost: x\nx-a: b\ncontent-length: 4\n\nbody")
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
let upload = received.recv().await.unwrap();
assert_eq!(upload.head.header("x-a"), Some(&b"b"[..]));
assert_eq!(upload.body, b"body");
}
/// h11 decodes a chunked body from an HTTP/1.0 client, and RFC 9112 section 6.1 only asks the
/// server to close the connection afterwards.
#[rstest]
#[ignore = "divergence: hyper rejects Transfer-Encoding on HTTP/1.0 with a 400 that uvicorn never sends"]
#[tokio::test]
async fn an_http10_chunked_upload_reaches_python() {
let (upstream, mut received) = capturing().await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("POST /v1/files HTTP/1.0\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n4\r\nbody\r\n0\r\n\r\n")
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
assert_eq!(received.recv().await.unwrap().body, b"body");
}

102
litellm-rust/crates/gateway-forward/tests/idle_stream.rs Normal file
View file

@ -0,0 +1,102 @@
mod support;
use std::time::Duration;
use axum::{
body::Body,
extract::Request,
http::StatusCode,
response::{IntoResponse, Response},
};
use http_body_util::BodyExt;
use rstest::rstest;
use support::{Gateway, client, raw_upstream, recording, within};
use tokio::{sync::mpsc, time::Instant};
const SSE_HEAD: &str = "HTTP/1.1 200 OK\r\n\
date: Tue, 03 Nov 2026 00:00:00 GMT\r\n\
server: uvicorn\r\n\
cache-control: no-cache, no-transform\r\n\
connection: keep-alive\r\n\
content-type: text/event-stream\r\n\
mcp-session-id: session-1\r\n\
x-accel-buffering: no\r\n\
transfer-encoding: chunked\r\n\r\n";
const PING: &str = ": ping - 2026-11-03 00:20:00.000000+00:00\r\n\r\n";
const MESSAGE: &str = "event: message\r\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/tools/list_changed\"}\r\n\r\n";
fn chunk(data: &str) -> String {
format!("{:x}\r\n{data}\r\n", data.len())
}
/// An MCP client holds a GET stream open for server-initiated messages, which may not come for
/// many minutes, so neither hop may close a stream for being idle.
#[rstest]
#[tokio::test(start_paused = true)]
async fn an_idle_mcp_stream_stays_open_in_both_directions() {
let idle = Duration::from_secs(20 * 60);
let (session, mut sessions) = mpsc::unbounded_channel();
let upstream = raw_upstream(move |mut python| {
let session = session.clone();
async move {
let head = python.request_head().await.unwrap();
session
.send(head.header("mcp-session-id").map(<[u8]>::to_vec))
.unwrap();
python.send(SSE_HEAD).await;
tokio::time::sleep(idle).await;
python.send(chunk(PING)).await;
python.send(chunk(MESSAGE)).await;
python.send("0\r\n\r\n").await;
}
})
.await;
let gateway = Gateway::start(upstream).await;
let started = Instant::now();
let response = client()
.request(
Request::get(gateway.url("/github/mcp"))
.header("accept", "text/event-stream")
.header("mcp-session-id", "session-1")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.headers()["mcp-session-id"], "session-1");
assert_eq!(response.headers()["x-accel-buffering"], "no");
assert!(!response.headers().contains_key("connection"));
assert_eq!(
sessions.recv().await.unwrap().as_deref(),
Some(&b"session-1"[..])
);
let body = response.into_body().collect().await.unwrap().to_bytes();
assert_eq!(body, [PING, MESSAGE].concat());
assert!(started.elapsed() >= idle);
}
#[rstest]
#[tokio::test]
async fn ends_a_session_with_a_bodiless_delete() {
let (upstream, mut received) = recording(|| -> Response {
([("mcp-session-id", "session-1")], StatusCode::NO_CONTENT).into_response()
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("DELETE /mcp HTTP/1.1\r\nhost: x\r\nmcp-session-id: session-1\r\n\r\n")
.await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.status, StatusCode::NO_CONTENT);
assert_eq!(head.header("mcp-session-id"), Some(&b"session-1"[..]));
let request = received.recv().await.unwrap();
assert_eq!(request.method, "DELETE");
assert_eq!(request.headers["mcp-session-id"], "session-1");
assert!(!request.headers.contains_key("transfer-encoding"));
assert!(request.body.is_empty());
}

473
litellm-rust/crates/gateway-forward/tests/lifecycle.rs Normal file
View file

@ -0,0 +1,473 @@
mod support;
use std::{
sync::Arc,
time::{Duration, Instant},
};
use axum::{Router, body::Body, extract::Request, http::StatusCode, response::Response};
use http_body_util::BodyExt;
use rstest::rstest;
use support::{
DRAIN, Gateway, Raw, client, raw_upstream, recording, serve, unreachable_upstream, within,
};
use tokio::{
net::{TcpListener, TcpStream},
sync::{Notify, mpsc},
};
fn ok() -> Response {
Response::new(Body::from("ok"))
}
/// Python stops a provider stream when it sees the client go away, so a client disconnect must
/// close Rust's connection to Python at every stage of the exchange.
#[rstest]
#[case::waiting_for_headers(
"POST /v1/chat/completions HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n",
false
)]
#[case::mid_stream(
"POST /v1/chat/completions HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n",
true
)]
#[case::mid_upload(
"POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: 1000000\r\n\r\npartial",
false
)]
#[case::waiting_for_websocket_handshake(
"GET /v1/realtime HTTP/1.1\r\nhost: x\r\nconnection: Upgrade\r\nupgrade: websocket\r\nsec-websocket-version: 13\r\nsec-websocket-key: dGhlIHNhbXBsZSBub25jZQ==\r\n\r\n",
false
)]
#[tokio::test]
async fn client_disconnect_closes_the_upstream_connection(
#[case] request: &'static str,
#[case] headers_sent: bool,
) {
let (event, mut events) = mpsc::unbounded_channel();
let upstream = raw_upstream(move |mut python: Raw| {
let event = event.clone();
async move {
python.request_head().await.unwrap();
if headers_sent {
python
.send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n6\r\ndata: \r\n")
.await;
}
event.send("received").unwrap();
python.closed().await;
event.send("closed").unwrap();
}
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(request).await;
assert_eq!(within(events.recv()).await, Some("received"));
if headers_sent {
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
}
drop(raw);
assert_eq!(within(events.recv()).await, Some("closed"));
}
/// Python treats EOF on its socket as the client leaving, so Rust must not half-close the
/// upstream connection once it has sent the request body.
#[rstest]
#[tokio::test]
async fn python_sees_no_eof_while_its_response_is_pending() {
let (verdict, mut verdicts) = mpsc::unbounded_channel();
let upstream = raw_upstream(move |mut python: Raw| {
let verdict = verdict.clone();
async move {
let head = python.request_head().await.unwrap();
python.body(head.content_length().unwrap()).await;
let eof = tokio::time::timeout(Duration::from_millis(200), python.closed())
.await
.is_ok();
verdict.send(eof).unwrap();
python
.send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok")
.await;
}
})
.await;
let gateway = Gateway::start(upstream).await;
let response = within(
client().request(
Request::post(gateway.url("/mcp"))
.body(Body::from("{}"))
.unwrap(),
),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(verdicts.recv().await, Some(false));
}
/// uvicorn closes a connection that sat idle past its keep-alive timeout. A request Rust sends on
/// that connection as uvicorn closes it is lost, so Rust must retire idle connections first.
#[rstest]
#[case::reused_while_fresh(Duration::from_millis(0))]
#[case::past_the_pool_timeout(Duration::from_millis(250))]
#[case::past_python_keepalive(Duration::from_millis(450))]
#[tokio::test]
async fn never_reuses_a_connection_python_is_closing_for_idleness(#[case] idle: Duration) {
let keepalive = Duration::from_millis(400);
let upstream = raw_upstream(move |mut python: Raw| async move {
let mut idle_since = None::<Instant>;
while python.request_head().await.is_some() {
if idle_since.is_some_and(|since| since.elapsed() >= keepalive) {
return;
}
python
.send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok")
.await;
idle_since = Some(Instant::now());
}
})
.await;
let gateway = Gateway::with_keepalive(upstream, keepalive).await;
let client = client();
let first = within(client.request(Request::get(gateway.url("/")).body(Body::empty()).unwrap()))
.await
.unwrap();
first.into_body().collect().await.unwrap();
tokio::time::sleep(idle).await;
let second =
within(client.request(Request::get(gateway.url("/")).body(Body::empty()).unwrap()))
.await
.unwrap();
assert_eq!(second.status(), StatusCode::OK);
}
#[rstest]
#[case::health("/health/liveliness")]
#[case::messages("/v1/messages")]
#[tokio::test]
async fn unreachable_python_is_a_502(#[case] path: &str) {
let gateway = Gateway::start(unreachable_upstream()).await;
let response =
within(client().request(Request::get(gateway.url(path)).body(Body::empty()).unwrap()))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
}
#[rstest]
#[case::bad_status_line(b"HTTP/1.1 OK\r\ncontent-length: 0\r\n\r\n".to_vec())]
#[case::header_without_colon(b"HTTP/1.1 200 OK\r\nbroken header\r\ncontent-length: 0\r\n\r\n".to_vec())]
#[case::not_http(b"SSH-2.0-OpenSSH_9.6\r\n".to_vec())]
#[case::conflicting_lengths(
b"HTTP/1.1 200 OK\r\ncontent-length: 2\r\ncontent-length: 3\r\n\r\nok".to_vec()
)]
#[case::oversized_head(
[&b"HTTP/1.1 200 OK\r\nx-big: "[..], &vec![b'a'; 1024 * 1024], b"\r\n\r\n"].concat()
)]
#[case::closed_without_reply(Vec::new())]
#[tokio::test]
async fn an_upstream_that_breaks_http_is_a_502_and_the_client_connection_survives(
#[case] reply: Vec<u8>,
) {
let upstream = raw_upstream(move |mut python: Raw| {
let reply = reply.clone();
async move {
python.request_head().await.unwrap();
if reply.is_empty() {
return;
}
python.send(reply).await;
std::future::pending::<()>().await;
}
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
for _ in 0..2 {
raw.send("GET / HTTP/1.1\r\nhost: x\r\n\r\n").await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.status, StatusCode::BAD_GATEWAY);
assert_eq!(head.content_length(), Some(0));
}
}
#[rstest]
#[tokio::test]
async fn python_restarting_is_a_502_until_it_listens_again() {
let address = unreachable_upstream();
let gateway = Gateway::start(address).await;
let get = || client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap());
assert_eq!(
within(get()).await.unwrap().status(),
StatusCode::BAD_GATEWAY
);
let listener = TcpListener::bind(address).await.unwrap();
let router = Router::new().fallback(|| async { "up" });
tokio::spawn(async move { axum::serve(listener, router).await });
assert_eq!(within(get()).await.unwrap().status(), StatusCode::OK);
}
/// A request Python already read may have run, so it is answered with a 502 and never replayed.
#[rstest]
#[tokio::test]
async fn a_request_python_read_before_closing_is_a_502_and_never_replayed() {
let (event, mut seen) = mpsc::unbounded_channel();
let upstream = raw_upstream(move |mut python: Raw| {
let event = event.clone();
async move {
while let Some(head) = python.request_head().await {
python.body(head.content_length().unwrap()).await;
event.send(head.path.clone()).unwrap();
if head.path == "/v1/chat/completions" {
return;
}
python
.send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok")
.await;
}
}
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("POST /warm HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n")
.await;
let warm = within(raw.response_head()).await.unwrap();
raw.body(warm.content_length().unwrap()).await;
raw.send("POST /v1/chat/completions HTTP/1.1\r\nhost: x\r\ncontent-length: 2\r\n\r\n{}")
.await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.status, StatusCode::BAD_GATEWAY);
let paths: Vec<_> = std::iter::from_fn(|| seen.try_recv().ok()).collect();
assert_eq!(paths, ["/warm", "/v1/chat/completions"]);
}
#[rstest]
#[tokio::test]
async fn shutdown_lets_an_in_flight_stream_finish() {
let release = Arc::new(Notify::new());
let released = release.clone();
let upstream = raw_upstream(move |mut python: Raw| {
let released = released.clone();
async move {
python.request_head().await.unwrap();
python
.send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n5\r\nfirst\r\n")
.await;
released.notified().await;
python.send("4\r\nlast\r\n0\r\n\r\n").await;
}
})
.await;
let gateway = Gateway::start(upstream).await;
let response =
within(client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap()))
.await
.unwrap();
let server = gateway.shutdown();
tokio::time::sleep(Duration::from_millis(100)).await;
assert!(!server.is_finished());
release.notify_one();
let body = within(response.into_body().collect())
.await
.unwrap()
.to_bytes();
assert_eq!(body, "firstlast");
within(server).await.unwrap();
}
#[rstest]
#[tokio::test]
async fn shutdown_stops_waiting_for_a_stream_still_open_after_the_drain() {
let upstream = raw_upstream(|mut python: Raw| async move {
python.request_head().await.unwrap();
python
.send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n5\r\nfirst\r\n")
.await;
std::future::pending::<()>().await;
})
.await;
let gateway = Gateway::start(upstream).await;
let response =
within(client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap()))
.await
.unwrap();
let started = std::time::Instant::now();
within(gateway.shutdown()).await.unwrap();
assert!(started.elapsed() >= DRAIN);
drop(response);
}
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_keep_alive_connections_each_get_their_own_responses() {
let upstream = serve(Router::new().fallback(|request: Request| async move {
let index: usize = request.uri().query().unwrap().parse().unwrap();
Response::builder()
.header("etag", format!("\"{index}\""))
.body(Body::from(index.to_string().repeat(1 + index * 1000)))
.unwrap()
}))
.await;
let gateway = Gateway::start(upstream).await;
let connections = 6;
let tasks: Vec<_> = (0..connections)
.map(|first| {
let address = gateway.address;
tokio::spawn(async move {
let mut raw = Raw::new(TcpStream::connect(address).await.unwrap());
for index in (first..50).step_by(connections) {
raw.send(format!(
"GET /ui/chunk.js?{index} HTTP/1.1\r\nhost: x\r\n\r\n"
))
.await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.header("etag"), Some(format!("\"{index}\"").as_bytes()));
let body = raw.body(head.content_length().unwrap()).await;
assert_eq!(
body,
index.to_string().repeat(1 + index * 1000).into_bytes()
);
}
})
})
.collect();
for task in tasks {
within(task).await.unwrap();
}
}
// Python's `--keepalive_timeout` governed idle client connections before Rust took the public
// port, and operators tune it against their load balancer's idle timeout, so it still does.
// Without it, a client that vanishes without closing its connection holds a socket forever and
// the gateway eventually runs out of file descriptors.
const IDLE: Duration = Duration::from_millis(400);
#[rstest]
#[tokio::test]
async fn closes_a_client_connection_idle_past_python_keepalive() {
let (upstream, _received) = recording(ok).await;
let gateway = Gateway::with_keepalive(upstream, IDLE).await;
let mut raw = gateway.connect().await;
for _ in 0..2 {
tokio::time::sleep(IDLE / 2).await;
raw.send("GET /health/liveliness HTTP/1.1\r\nhost: x\r\n\r\n")
.await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.status, StatusCode::OK);
raw.body(head.content_length().unwrap()).await;
}
let idle_since = Instant::now();
within(raw.closed()).await;
assert!(idle_since.elapsed() >= IDLE * 3 / 4);
}
/// uvicorn would hold a connection that never sends a request, which lets a client park sockets
/// for free, so the same limit applies before the first request too.
#[rstest]
#[tokio::test]
async fn closes_a_client_connection_that_never_sends_a_request() {
let (upstream, _received) = recording(ok).await;
let gateway = Gateway::with_keepalive(upstream, IDLE).await;
let mut raw = gateway.connect().await;
within(raw.closed()).await;
}
/// The limit only covers the wait for the next request. A response that goes quiet for longer
/// belongs to Python, which decides how long it may take.
#[rstest]
#[tokio::test]
async fn keeps_a_connection_open_while_a_response_goes_quiet() {
let upstream = raw_upstream(|mut python: Raw| async move {
python.request_head().await.unwrap();
python
.send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n")
.await;
tokio::time::sleep(IDLE * 3).await;
python.send("4\r\ndone\r\n0\r\n\r\n").await;
})
.await;
let gateway = Gateway::with_keepalive(upstream, IDLE).await;
let mut raw = gateway.connect().await;
raw.send("GET /v1/chat/completions HTTP/1.1\r\nhost: x\r\n\r\n")
.await;
within(raw.response_head()).await.unwrap();
assert_eq!(within(raw.chunked_body()).await, b"done");
}
/// Likewise for an upload that pauses longer than the limit before it finishes.
#[rstest]
#[tokio::test]
async fn keeps_a_connection_open_while_an_upload_goes_quiet() {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::with_keepalive(upstream, IDLE).await;
let mut raw = gateway.connect().await;
raw.send("POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: 8\r\n\r\nhalf")
.await;
tokio::time::sleep(IDLE * 3).await;
raw.send("done").await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
assert_eq!(received.recv().await.unwrap().body, "halfdone");
}
/// uvicorn answers a request it rejects (413, 401) as soon as it reads the head and then closes,
/// leaving the upload unread. The client must get that answer, not a 502 or a reset.
#[rstest]
#[case::sized("content-length: 16777216\r\n", vec![b'x'; 256 * 1024])]
#[case::chunked(
"transfer-encoding: chunked\r\n",
[&b"40000\r\n"[..], &[b'x'; 256 * 1024], b"\r\n"].concat()
)]
#[tokio::test]
async fn python_rejecting_an_upload_unread_and_closing_still_answers_the_client(
#[case] framing: &'static str,
#[case] partial_upload: Vec<u8>,
) {
let upstream = raw_upstream(|mut python: Raw| async move {
python.request_head().await.unwrap();
python
.send("HTTP/1.1 413 Payload Too Large\r\ncontent-type: application/json\r\ncontent-length: 26\r\nconnection: close\r\n\r\n{\"error\":\"file too large\"}")
.await;
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(format!(
"POST /v1/files HTTP/1.1\r\nhost: x\r\n{framing}\r\n"
))
.await;
raw.send(partial_upload).await;
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.status, StatusCode::PAYLOAD_TOO_LARGE);
assert_eq!(
raw.body(head.content_length().unwrap()).await,
b"{\"error\":\"file too large\"}"
);
}

360
litellm-rust/crates/gateway-forward/tests/roundtrip.rs Normal file
View file

@ -0,0 +1,360 @@
mod support;
use std::{
convert::Infallible,
sync::{Arc, Mutex},
};
use axum::{
body::{Body, Bytes},
http::{HeaderMap, HeaderName, HeaderValue, Method, Request, StatusCode},
response::Response,
};
use http_body_util::BodyExt;
use litellm_gateway_forward::PEER_HEADER;
use proptest::{prelude::*, test_runner::TestRunner};
use support::{Gateway, client, recording, within};
#[derive(Clone, Debug)]
struct Message {
headers: Vec<(String, Vec<Vec<u8>>)>,
body: Vec<u8>,
chunked: bool,
}
#[derive(Clone, Debug)]
struct Exchange {
method: Method,
target: String,
request: Message,
status: StatusCode,
response: Message,
}
/// Any value h11 reads back unchanged: visible ASCII, obs-text and interior whitespace.
fn header_value() -> impl Strategy<Value = Vec<u8>> {
prop::collection::vec(
prop_oneof![8 => 0x21u8..=0x7e, 2 => 0x80u8..=0xff, 1 => Just(b' '), 1 => Just(b'\t')],
0..24,
)
.prop_map(|value| value.trim_ascii().to_vec())
}
fn method() -> impl Strategy<Value = Method> {
prop_oneof![
Just(Method::GET),
Just(Method::POST),
Just(Method::PUT),
Just(Method::PATCH),
Just(Method::DELETE),
Just(Method::OPTIONS),
"[A-Za-z0-9!#$%&'*+.^_`|~-]{1,16}"
.prop_filter("HEAD and CONNECT frame differently", |method| {
!["HEAD", "CONNECT"].contains(&method.as_str())
})
.prop_map(|method| Method::from_bytes(method.as_bytes()).unwrap()),
]
}
fn message() -> impl Strategy<Value = Message> {
(
prop::collection::btree_map(
"x-[a-z0-9!#$%&'*+.^_`|~-]{1,12}",
prop::collection::vec(header_value(), 1..4),
0..12,
),
prop::collection::vec(any::<u8>(), 0..32 * 1024),
any::<bool>(),
)
.prop_map(|(headers, body, chunked)| Message {
headers: headers.into_iter().collect(),
body,
chunked,
})
}
fn target() -> impl Strategy<Value = String> {
let segment = prop_oneof![
"[A-Za-z0-9._~!$&'()*+,;=:@-]{0,8}",
"%[0-9A-Fa-f]{2}",
Just(".".to_owned()),
Just("..".to_owned()),
Just(String::new()),
];
(
prop::collection::vec(segment, 0..6),
prop::option::of("[A-Za-z0-9=&._~%/?:@!$'()*+,;-]{0,40}"),
)
.prop_map(|(segments, query)| {
let path = format!("/{}", segments.join("/"));
match query {
Some(query) => format!("{path}?{query}"),
None => path,
}
})
}
fn exchange() -> impl Strategy<Value = Exchange> {
let status = (200u16..600)
.prop_filter("bodiless status", |code| ![204, 205, 304].contains(code))
.prop_map(|code| StatusCode::from_u16(code).unwrap());
(method(), target(), message(), status, message()).prop_map(
|(method, target, request, status, response)| Exchange {
method,
target,
request,
status,
response,
},
)
}
fn body(message: &Message) -> Body {
if !message.chunked {
return Body::from(message.body.clone());
}
let chunks: Vec<Result<Bytes, Infallible>> = message
.body
.chunks(1000)
.map(|chunk| Ok(Bytes::copy_from_slice(chunk)))
.collect();
Body::from_stream(futures_util::stream::iter(chunks))
}
fn apply(headers: &mut HeaderMap, message: &Message) {
if message.chunked {
headers.insert("transfer-encoding", HeaderValue::from_static("chunked"));
}
for (name, values) in &message.headers {
for value in values {
headers.append(
HeaderName::from_bytes(name.as_bytes()).unwrap(),
HeaderValue::from_bytes(value).unwrap(),
);
}
}
}
fn assert_carried(headers: &HeaderMap, message: &Message, framing_and_hop: &[&str]) {
for (name, values) in &message.headers {
let carried: Vec<&[u8]> = headers
.get_all(name)
.iter()
.map(HeaderValue::as_bytes)
.collect();
let sent: Vec<&[u8]> = values.iter().map(Vec::as_slice).collect();
assert_eq!(carried, sent, "header {name}");
}
for name in headers.keys() {
assert!(
framing_and_hop.contains(&name.as_str())
|| message
.headers
.iter()
.any(|(sent, _)| sent == name.as_str()),
"unexpected header {name}"
);
}
}
/// Any request a client can send and any response Python can give cross the forwarder unchanged,
/// apart from the framing and peer headers the forwarder owns.
#[test]
fn every_exchange_crosses_the_forwarder_unchanged() {
let runtime = tokio::runtime::Runtime::new().unwrap();
let current = Arc::new(Mutex::new(None::<Exchange>));
let responding = current.clone();
let (gateway, received) = runtime.block_on(async {
let (upstream, received) = recording(move || {
let exchange = responding.lock().unwrap().clone().unwrap();
let mut response = Response::new(body(&exchange.response));
*response.status_mut() = exchange.status;
apply(response.headers_mut(), &exchange.response);
response
})
.await;
(
Gateway::start(upstream).await,
tokio::sync::Mutex::new(received),
)
});
let client = client();
TestRunner::new(ProptestConfig::with_cases(256))
.run(&exchange(), |exchange| {
*current.lock().unwrap() = Some(exchange.clone());
runtime.block_on(async {
while received.lock().await.try_recv().is_ok() {}
let mut request = Request::builder()
.method(exchange.method.clone())
.uri(gateway.url(&exchange.target))
.body(body(&exchange.request))
.unwrap();
apply(request.headers_mut(), &exchange.request);
let response = support::within(client.request(request)).await.unwrap();
let python = received.lock().await.recv().await.unwrap();
assert_eq!(python.method, exchange.method);
assert_eq!(python.uri, exchange.target.as_str());
assert_eq!(python.body, exchange.request.body);
assert_carried(
&python.headers,
&exchange.request,
&["host", "content-length", "transfer-encoding", PEER_HEADER],
);
assert_eq!(python.headers.get_all(PEER_HEADER).iter().count(), 1);
assert_eq!(response.status(), exchange.status);
let headers = response.headers().clone();
assert_carried(
&headers,
&exchange.response,
&["date", "content-length", "transfer-encoding"],
);
let body = support::within(response.into_body().collect())
.await
.unwrap();
assert_eq!(body.to_bytes(), exchange.response.body);
});
Ok(())
})
.unwrap();
}
#[derive(Clone, Debug)]
enum Framing {
Unframed,
Length,
Chunked(usize),
}
#[derive(Clone, Debug)]
struct WireRequest {
method: Method,
target: String,
headers: Vec<(String, Vec<u8>)>,
body: Vec<u8>,
framing: Framing,
}
fn wire_request() -> impl Strategy<Value = WireRequest> {
let framing = prop_oneof![
Just(Framing::Unframed),
Just(Framing::Length),
(1usize..2048).prop_map(Framing::Chunked),
];
(
method(),
target(),
prop::collection::vec(
("[xX]-[A-Za-z0-9!#$%&'*+.^_`|~-]{1,12}", header_value()),
0..24,
),
prop::collection::vec(any::<u8>(), 0..8 * 1024),
framing,
)
.prop_map(|(method, target, headers, body, framing)| WireRequest {
method,
target,
headers,
body: match framing {
Framing::Unframed => Vec::new(),
_ => body,
},
framing,
})
}
fn encode(request: &WireRequest) -> Vec<u8> {
let framing = match request.framing {
Framing::Unframed => String::new(),
Framing::Length => format!("Content-Length: {}\r\n", request.body.len()),
Framing::Chunked(_) => "Transfer-Encoding: chunked\r\n".to_owned(),
};
let fields = request
.headers
.iter()
.flat_map(|(name, value)| [name.as_bytes(), b": ", value, b"\r\n"].concat());
let body = match request.framing {
Framing::Chunked(size) => request
.body
.chunks(size)
.flat_map(|chunk| {
[format!("{:x}\r\n", chunk.len()).as_bytes(), chunk, b"\r\n"].concat()
})
.chain(*b"0\r\n\r\n")
.collect(),
_ => request.body.clone(),
};
[
format!(
"{} {} HTTP/1.1\r\nHost: x\r\n",
request.method, request.target
)
.into_bytes(),
fields.collect(),
framing.into_bytes(),
b"\r\n".to_vec(),
body,
]
.concat()
}
fn by_name<'a>(fields: impl Iterator<Item = (&'a str, &'a [u8])>) -> Vec<(String, &'a [u8])> {
let mut fields: Vec<_> = fields
.map(|(name, value)| (name.to_ascii_lowercase(), value))
.collect();
fields.sort_by(|(left, _), (right, _)| left.cmp(right));
fields
}
/// Pipelined requests in any framing, written in arbitrary pieces, reach Python one by one with
/// the bytes h11 would have read from the client directly.
#[test]
fn pipelined_wire_requests_reach_python_as_sent() {
let runtime = tokio::runtime::Runtime::new().unwrap();
let (gateway, received) = runtime.block_on(async {
let (upstream, received) = recording(|| Response::new(Body::from("ok"))).await;
(
Gateway::start(upstream).await,
tokio::sync::Mutex::new(received),
)
});
let requests = (prop::collection::vec(wire_request(), 1..5), 1usize..4096);
TestRunner::new(ProptestConfig::with_cases(256))
.run(&requests, |(requests, piece)| {
runtime.block_on(async {
let mut received = received.lock().await;
while received.try_recv().is_ok() {}
let mut raw = gateway.connect().await;
let wire: Vec<u8> = requests.iter().flat_map(encode).collect();
for bytes in wire.chunks(piece) {
raw.send(bytes).await;
}
for request in &requests {
let head = within(raw.response_head()).await.unwrap();
assert_eq!(head.status, StatusCode::OK, "{request:?}");
assert_eq!(raw.body(2).await, b"ok");
let python = within(received.recv()).await.unwrap();
assert_eq!(python.method, request.method);
assert_eq!(python.uri, request.target.as_str());
assert_eq!(python.body, request.body);
assert_eq!(python.headers["host"], "x");
let forwarded = python.headers.iter().filter(|(name, _)| {
!["host", "content-length", "transfer-encoding", PEER_HEADER]
.contains(&name.as_str())
});
assert_eq!(
by_name(forwarded.map(|(name, value)| (name.as_str(), value.as_bytes()))),
by_name(
request
.headers
.iter()
.map(|(name, value)| (name.as_str(), value.as_slice()))
)
);
}
});
Ok(())
})
.unwrap();
}

110
litellm-rust/crates/gateway-forward/tests/settings.rs Normal file
View file

@ -0,0 +1,110 @@
use std::{ffi::OsString, time::Duration};
use litellm_gateway_forward::{Error, Settings};
use rstest::rstest;
const COMPLETE: [(&str, &str); 5] = [
("LITELLM_FORWARD_UPSTREAM", "http://127.0.0.1:53211"),
("LITELLM_FORWARD_HOST", "0.0.0.0"),
("LITELLM_FORWARD_PORT", "4000"),
("LITELLM_FORWARD_UPSTREAM_KEEPALIVE", "2"),
("LITELLM_FORWARD_SHUTDOWN_DRAIN", "40"),
];
fn parse(variables: &[(&str, &str)]) -> Result<Option<Settings>, Error> {
Settings::from_vars(
variables
.iter()
.map(|(key, value)| (OsString::from(key), OsString::from(value))),
)
}
fn without(variable: &str) -> Vec<(&'static str, &'static str)> {
COMPLETE
.into_iter()
.filter(|(key, _)| *key != variable)
.collect()
}
fn with<'a>(variable: &'a str, value: &'a str) -> Vec<(&'a str, &'a str)> {
without(variable)
.into_iter()
.chain([(variable, value)])
.collect()
}
#[rstest]
fn reads_every_variable() {
assert_eq!(
parse(&COMPLETE).unwrap(),
Some(Settings {
upstream: "http://127.0.0.1:53211".parse().unwrap(),
host: "0.0.0.0".into(),
port: 4000,
upstream_keepalive: Duration::from_secs(2),
shutdown_drain: Duration::from_secs(40),
})
);
}
#[rstest]
fn without_an_upstream_the_gateway_is_not_forwarding() {
assert_eq!(parse(&without("LITELLM_FORWARD_UPSTREAM")).unwrap(), None);
}
#[rstest]
#[case::host("LITELLM_FORWARD_HOST")]
#[case::port("LITELLM_FORWARD_PORT")]
#[case::upstream_keepalive("LITELLM_FORWARD_UPSTREAM_KEEPALIVE")]
#[case::shutdown_drain("LITELLM_FORWARD_SHUTDOWN_DRAIN")]
fn a_missing_variable_fails_naming_it(#[case] variable: &str) {
let error = parse(&without(variable)).unwrap_err();
assert!(error.to_string().contains(variable), "{error}");
}
#[rstest]
#[case::port_overflow("LITELLM_FORWARD_PORT", "65536")]
#[case::port_text("LITELLM_FORWARD_PORT", "http")]
#[case::keepalive_text("LITELLM_FORWARD_UPSTREAM_KEEPALIVE", "soon")]
#[case::keepalive_negative("LITELLM_FORWARD_UPSTREAM_KEEPALIVE", "-1")]
#[case::keepalive_zero("LITELLM_FORWARD_UPSTREAM_KEEPALIVE", "0")]
#[case::drain_text("LITELLM_FORWARD_SHUTDOWN_DRAIN", "soon")]
#[case::drain_negative("LITELLM_FORWARD_SHUTDOWN_DRAIN", "-1")]
#[case::drain_infinite("LITELLM_FORWARD_SHUTDOWN_DRAIN", "inf")]
fn an_invalid_variable_fails_naming_it(#[case] variable: &str, #[case] value: &str) {
let error = parse(&with(variable, value)).unwrap_err();
assert!(error.to_string().contains(variable), "{error}");
}
#[rstest]
fn fractional_drain_seconds_are_used_as_given() {
let settings = parse(&with("LITELLM_FORWARD_SHUTDOWN_DRAIN", "12.5"))
.unwrap()
.unwrap();
assert_eq!(settings.shutdown_drain, Duration::from_millis(12_500));
}
#[rstest]
#[case::keepalive("KEEPALIVE_TIMEOUT")]
#[case::graceful_shutdown("GRACEFUL_SHUTDOWN_TIMEOUT")]
#[case::host("HOST")]
#[case::port("PORT")]
fn python_settings_are_never_read(#[case] variable: &str) {
let mut variables = COMPLETE.to_vec();
variables.push((variable, "9"));
assert_eq!(parse(&variables).unwrap(), parse(&COMPLETE).unwrap());
}
#[rstest]
fn a_non_loopback_upstream_fails_startup() {
let error = parse(&with("LITELLM_FORWARD_UPSTREAM", "http://192.0.2.1:4000")).unwrap_err();
assert!(
error.to_string().contains("LITELLM_FORWARD_UPSTREAM"),
"{error}"
);
}

41
litellm-rust/crates/gateway-forward/tests/slow_head.rs Normal file
View file

@ -0,0 +1,41 @@
mod support;
use std::time::Duration;
use axum::{body::Body, extract::Request, http::StatusCode};
use http_body_util::BodyExt;
use rstest::rstest;
use support::{Gateway, client, raw_upstream};
use tokio::time::Instant;
/// Reasoning models with `stream=false`, image and video generation, and model health checks send
/// no response headers until the provider finishes, which can take many minutes.
#[rstest]
#[tokio::test(start_paused = true)]
async fn waits_for_response_headers_as_long_as_python_takes() {
let thinking = Duration::from_secs(30 * 60);
let upstream = raw_upstream(move |mut python| async move {
python.request_head().await.unwrap();
tokio::time::sleep(thinking).await;
python
.send("HTTP/1.1 200 OK\r\ncontent-length: 5\r\n\r\nimage")
.await;
})
.await;
let gateway = Gateway::start(upstream).await;
let started = Instant::now();
let response = client()
.request(
Request::post(gateway.url("/v1/images/generations"))
.body(Body::from("{}"))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
assert!(started.elapsed() >= thinking);
let body = response.into_body().collect().await.unwrap().to_bytes();
assert_eq!(body, "image");
}

313
litellm-rust/crates/gateway-forward/tests/stream.rs Normal file
View file

@ -0,0 +1,313 @@
mod support;
use std::{
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
time::Duration,
};
use axum::{body::Body, extract::Request, http::StatusCode};
use http_body_util::BodyExt;
use hyper::body::Incoming;
use rstest::rstest;
use support::{Gateway, client, raw_upstream, within};
use tokio::{
io::AsyncWriteExt,
sync::{Notify, mpsc},
};
fn chunk(bytes: &[u8]) -> Vec<u8> {
[format!("{:x}\r\n", bytes.len()).as_bytes(), bytes, b"\r\n"].concat()
}
const EVENTS: [&str; 3] = ["data: 1\n\n", "data: 2\n\n", "data: [DONE]\n\n"];
#[rstest]
#[tokio::test]
async fn each_event_reaches_the_client_before_python_produces_the_next(
#[values(true, false)] chunked: bool,
) {
let produce = Arc::new(Notify::new());
let next = produce.clone();
let upstream = raw_upstream(move |mut python| {
let next = next.clone();
async move {
python.request_head().await.unwrap();
let framing = match chunked {
true => "transfer-encoding: chunked".to_owned(),
false => format!("content-length: {}", EVENTS.concat().len()),
};
python
.send(format!(
"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\n{framing}\r\n\r\n"
))
.await;
for event in EVENTS {
match chunked {
true => python.send(chunk(event.as_bytes())).await,
false => python.send(event).await,
}
next.notified().await;
}
if chunked {
python.send("0\r\n\r\n").await;
}
}
})
.await;
let gateway = Gateway::start(upstream).await;
let mut response = within(
client().request(
Request::post(gateway.url("/v1/chat/completions"))
.body(Body::from(r#"{"stream":true}"#))
.unwrap(),
),
)
.await
.unwrap();
for event in EVENTS {
let frame = within(response.body_mut().frame()).await.unwrap().unwrap();
assert_eq!(frame.into_data().unwrap(), event);
produce.notify_one();
}
let rest = within(response.into_body().collect()).await.unwrap();
assert!(rest.to_bytes().is_empty());
}
#[rstest]
#[tokio::test]
async fn relays_a_binary_framed_stream_unchanged() {
let message: Vec<u8> = [
&[
0x00, 0x00, 0x00, 0x3d, 0x00, 0x00, 0x00, 0x20, 0x8a, 0x0d, 0x0a, 0xff,
][..],
b"\x0b:event-type\x07\x00\x05chunk\x0d:content-type\x07\x00\x10application/json",
b"0\r\n\r\n{\"bytes\":\"eyJ0eXBlIjoiY29udGVudCJ9\"}",
&[0xc3, 0x28, 0x00, 0x0d],
]
.concat();
let sent = message.clone();
let upstream = raw_upstream(move |mut python| {
let sent = sent.clone();
async move {
python.request_head().await.unwrap();
python
.send("HTTP/1.1 200 OK\r\ncontent-type: application/vnd.amazon.eventstream\r\ntransfer-encoding: chunked\r\n\r\n")
.await;
for piece in sent.chunks(7) {
python.send(chunk(piece)).await;
}
python.send("0\r\n\r\n").await;
}
})
.await;
let gateway = Gateway::start(upstream).await;
let response = within(
client().request(
Request::post(gateway.url("/bedrock/model/m/converse-stream"))
.body(Body::empty())
.unwrap(),
),
)
.await
.unwrap();
assert_eq!(
response.headers()["content-type"],
"application/vnd.amazon.eventstream"
);
let body = within(response.into_body().collect())
.await
.unwrap()
.to_bytes();
assert_eq!(body, message);
}
#[rstest]
#[tokio::test]
async fn relays_a_body_delimited_by_the_upstream_closing() {
let upstream = raw_upstream(|mut python| async move {
python.request_head().await.unwrap();
python
.send("HTTP/1.1 200 OK\r\nconnection: close\r\n\r\nuntil eof")
.await;
})
.await;
let gateway = Gateway::start(upstream).await;
let response = client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap());
let body = within(response).await.unwrap().into_body().collect();
assert_eq!(within(body).await.unwrap().to_bytes(), "until eof");
}
async fn drain(mut body: Incoming) -> (Vec<u8>, Result<(), hyper::Error>) {
let mut delivered = Vec::new();
loop {
match within(body.frame()).await {
None => return (delivered, Ok(())),
Some(Err(error)) => return (delivered, Err(error)),
Some(Ok(frame)) => delivered.extend_from_slice(&frame.into_data().unwrap()),
}
}
}
#[rstest]
#[case::after_the_head("transfer-encoding: chunked", "", "")]
#[case::between_events("transfer-encoding: chunked", "6\r\ndata: \r\n", "data: ")]
#[case::mid_event("transfer-encoding: chunked", "a\r\ndata:", "data:")]
#[case::short_of_its_content_length("content-length: 100", "data: 1\n\n", "data: 1\n\n")]
#[tokio::test]
async fn an_upstream_dying_mid_stream_fails_the_client_body_instead_of_ending_it(
#[case] framing: &'static str,
#[case] sent: &'static str,
#[case] delivered: &'static str,
) {
let upstream = raw_upstream(move |mut python| async move {
python.request_head().await.unwrap();
python
.send(format!(
"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\n{framing}\r\n\r\n{sent}"
))
.await;
})
.await;
let gateway = Gateway::start(upstream).await;
let response =
within(client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap()))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let (received, ended) = drain(response.into_body()).await;
assert_eq!(received, delivered.as_bytes());
assert!(ended.is_err());
}
#[rstest]
#[tokio::test]
async fn response_headers_reach_the_client_before_python_sends_any_body() {
let release = Arc::new(Notify::new());
let released = release.clone();
let upstream = raw_upstream(move |mut python| {
let released = released.clone();
async move {
python.request_head().await.unwrap();
python
.send("HTTP/1.1 200 OK\r\ncontent-type: text/event-stream; charset=utf-8\r\nx-accel-buffering: no\r\nx-litellm-call-id: call-1\r\ntransfer-encoding: chunked\r\n\r\n")
.await;
released.notified().await;
python.send(chunk(b"data: 1\n\n")).await;
python.send("0\r\n\r\n").await;
}
})
.await;
let gateway = Gateway::start(upstream).await;
let response =
within(client().request(Request::post(gateway.url("/")).body(Body::empty()).unwrap()))
.await
.unwrap();
assert_eq!(response.headers()["x-litellm-call-id"], "call-1");
assert_eq!(response.headers()["x-accel-buffering"], "no");
release.notify_one();
let body = within(response.into_body().collect())
.await
.unwrap()
.to_bytes();
assert_eq!(body, "data: 1\n\n");
}
/// A connection closed without any response looks retryable to client SDKs, which would re-run a
/// request Python already ran. Python dying right after its head races the relay, hence the loop.
#[rstest]
#[tokio::test]
async fn a_head_python_sent_right_before_dying_still_reaches_the_client() {
let upstream = raw_upstream(|mut python| async move {
python.request_head().await.unwrap();
python
.send("HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n")
.await;
})
.await;
let gateway = Gateway::start(upstream).await;
for _ in 0..20 {
let mut raw = gateway.connect().await;
raw.send("POST / HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n")
.await;
let head = within(raw.response_head()).await;
assert_eq!(head.map(|head| head.status), Some(StatusCode::OK));
}
}
#[rstest]
#[tokio::test]
async fn a_client_that_stops_reading_stalls_python_instead_of_being_buffered_for() {
let written = Arc::new(AtomicUsize::new(0));
let counted = written.clone();
let (failed, mut write_failed) = mpsc::unbounded_channel();
let upstream = raw_upstream(move |mut python| {
let counted = counted.clone();
let failed = failed.clone();
async move {
python.request_head().await.unwrap();
let mut python = python.into_stream();
python
.write_all(b"HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n")
.await
.unwrap();
let piece = chunk(&[0x5a; 64 * 1024]);
while python.write_all(&piece).await.is_ok() {
counted.fetch_add(piece.len(), Ordering::SeqCst);
}
failed.send(()).unwrap();
}
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("POST / HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n")
.await;
within(raw.response_head()).await.unwrap();
tokio::time::sleep(Duration::from_secs(1)).await;
let stalled = written.load(Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(250)).await;
assert_eq!(written.load(Ordering::SeqCst), stalled);
assert!(stalled < 64 * 1024 * 1024, "{stalled} bytes in flight");
drop(raw);
within(write_failed.recv()).await.unwrap();
}
#[rstest]
#[tokio::test]
async fn an_http10_client_receives_a_streamed_body_delimited_by_close() {
let upstream = raw_upstream(|mut python| async move {
python.request_head().await.unwrap();
python
.send([
&b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n"[..],
&chunk(b"data: 1\n"),
&chunk(b"data: 2\n"),
b"0\r\n\r\n",
]
.concat())
.await;
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("GET /v1/stream HTTP/1.0\r\nhost: x\r\n\r\n").await;
let response = String::from_utf8(within(raw.until_closed()).await).unwrap();
let (head, body) = response.split_once("\r\n\r\n").unwrap();
assert!(head.starts_with("HTTP/1.0 200 OK"), "{head}");
assert!(
!head.to_ascii_lowercase().contains("transfer-encoding"),
"{head}"
);
assert_eq!(body, "data: 1\ndata: 2\n");
}

311
litellm-rust/crates/gateway-forward/tests/support/mod.rs Normal file
View file

@ -0,0 +1,311 @@
#![allow(
dead_code,
reason = "each test binary uses a different subset of the harness"
)]
use std::{future::Future, net::SocketAddr, time::Duration};
use axum::{
Router,
body::{Body, Bytes},
extract::Request,
http::{HeaderMap, Method, StatusCode, Uri},
};
use http_body_util::BodyExt;
use hyper_util::{
client::legacy::{Client, connect::HttpConnector},
rt::TokioExecutor,
};
use litellm_gateway_forward::Forward;
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{TcpListener, TcpStream},
sync::{mpsc, oneshot},
task::JoinHandle,
};
pub const WAIT: Duration = Duration::from_secs(5);
pub const UVICORN_KEEPALIVE: Duration = Duration::from_secs(5);
pub const DRAIN: Duration = Duration::from_secs(1);
pub async fn within<T>(future: impl Future<Output = T>) -> T {
tokio::time::timeout(WAIT, future).await.expect("timed out")
}
/// A running forwarder in front of `upstream`, listening on a real socket like production.
pub struct Gateway {
pub address: SocketAddr,
stop: oneshot::Sender<()>,
server: JoinHandle<()>,
}
impl Gateway {
pub async fn start(upstream: SocketAddr) -> Self {
Self::with_keepalive(upstream, UVICORN_KEEPALIVE).await
}
pub async fn with_keepalive(upstream: SocketAddr, keepalive: Duration) -> Self {
let forward = Forward::new(&format!("http://{upstream}").parse().unwrap(), keepalive);
let public = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = public.local_addr().unwrap();
let (stop, stopped) = oneshot::channel::<()>();
let server = tokio::spawn(forward.serve(
public,
async {
stopped.await.ok();
},
DRAIN,
));
Self {
address,
stop,
server,
}
}
pub fn url(&self, path_and_query: &str) -> Uri {
format!("http://{}{path_and_query}", self.address)
.parse()
.unwrap()
}
pub async fn connect(&self) -> Raw {
Raw::new(TcpStream::connect(self.address).await.unwrap())
}
/// Starts graceful shutdown and returns the server task, which ends once in-flight work drains
/// or `DRAIN` passes.
pub fn shutdown(self) -> JoinHandle<()> {
self.stop.send(()).unwrap();
self.server
}
}
pub fn client() -> Client<HttpConnector, Body> {
Client::builder(TokioExecutor::new()).build_http()
}
/// What the upstream received, after the forwarder rewrote it.
pub struct Received {
pub method: Method,
pub uri: Uri,
pub headers: HeaderMap,
pub body: Bytes,
}
/// An axum upstream that records every request it fully reads and answers with `respond`.
pub async fn recording(
respond: impl Fn() -> axum::response::Response + Clone + Send + Sync + 'static,
) -> (SocketAddr, mpsc::UnboundedReceiver<Received>) {
let (sender, receiver) = mpsc::unbounded_channel();
let address = serve(Router::new().fallback(move |request: Request| {
let sender = sender.clone();
let respond = respond.clone();
async move {
let (parts, body) = request.into_parts();
let body = body.collect().await.unwrap().to_bytes();
sender
.send(Received {
method: parts.method,
uri: parts.uri,
headers: parts.headers,
body,
})
.unwrap();
respond()
}
}))
.await;
(address, receiver)
}
pub async fn serve(router: Router) -> SocketAddr {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, router).await.unwrap() });
address
}
/// An upstream that hands each accepted connection to `handle` as raw bytes, for behavior a
/// well-behaved HTTP server would never produce.
pub async fn raw_upstream<F, Fut>(handle: F) -> SocketAddr
where
F: Fn(Raw) -> Fut + Send + Sync + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
tokio::spawn(async move {
loop {
let (stream, _) = listener.accept().await.unwrap();
tokio::spawn(handle(Raw::new(stream)));
}
});
address
}
pub fn unreachable_upstream() -> SocketAddr {
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
listener.local_addr().unwrap()
}
pub struct Head {
pub status: StatusCode,
pub headers: Vec<(String, Vec<u8>)>,
pub path: String,
}
impl Head {
pub fn header(&self, name: &str) -> Option<&[u8]> {
self.headers
.iter()
.find(|(key, _)| key.eq_ignore_ascii_case(name))
.map(|(_, value)| value.as_slice())
}
pub fn content_length(&self) -> Option<usize> {
self.header("content-length")
.map(|value| std::str::from_utf8(value).unwrap().parse().unwrap())
}
}
/// One TCP connection driven byte by byte, so tests control framing, timing and disconnects.
/// Bytes read past a message head stay buffered for the next read.
pub struct Raw {
stream: TcpStream,
buffer: Vec<u8>,
}
impl Raw {
pub fn new(stream: TcpStream) -> Self {
Self {
stream,
buffer: Vec::new(),
}
}
pub fn into_stream(self) -> TcpStream {
assert!(self.buffer.is_empty(), "buffered bytes would be lost");
self.stream
}
pub fn local_address(&self) -> SocketAddr {
self.stream.local_addr().unwrap()
}
pub async fn send(&mut self, bytes: impl AsRef<[u8]>) {
self.stream.write_all(bytes.as_ref()).await.unwrap();
}
/// Reads one request head, or `None` when the peer closed the connection first.
pub async fn request_head(&mut self) -> Option<Head> {
self.head(|bytes| {
let mut headers = [httparse::EMPTY_HEADER; 256];
let mut request = httparse::Request::new(&mut headers);
let httparse::Status::Complete(length) = request.parse(bytes).unwrap() else {
return None;
};
Some((
length,
Head {
status: StatusCode::OK,
headers: owned(request.headers),
path: request.path.unwrap().to_owned(),
},
))
})
.await
}
pub async fn response_head(&mut self) -> Option<Head> {
self.head(|bytes| {
let mut headers = [httparse::EMPTY_HEADER; 256];
let mut response = httparse::Response::new(&mut headers);
let httparse::Status::Complete(length) = response.parse(bytes).unwrap() else {
return None;
};
Some((
length,
Head {
status: StatusCode::from_u16(response.code.unwrap()).unwrap(),
headers: owned(response.headers),
path: String::new(),
},
))
})
.await
}
async fn head(&mut self, parse: impl Fn(&[u8]) -> Option<(usize, Head)>) -> Option<Head> {
loop {
if let Some((length, head)) = parse(&self.buffer) {
self.buffer.drain(..length);
return Some(head);
}
if self.fill().await == 0 {
return None;
}
}
}
/// Reads exactly `length` body bytes, using what is already buffered first.
pub async fn body(&mut self, length: usize) -> Vec<u8> {
while self.buffer.len() < length {
assert_ne!(self.fill().await, 0, "closed mid-body");
}
self.buffer.drain(..length).collect()
}
/// Reads a whole chunked body and returns its decoded bytes.
pub async fn chunked_body(&mut self) -> Vec<u8> {
let mut decoded = Vec::new();
loop {
let size_line = self.line().await;
let size =
usize::from_str_radix(size_line.split(';').next().unwrap().trim(), 16).unwrap();
if size == 0 {
while !self.line().await.is_empty() {}
return decoded;
}
decoded.extend(self.body(size).await);
assert_eq!(self.body(2).await, b"\r\n");
}
}
async fn line(&mut self) -> String {
loop {
if let Some(end) = self.buffer.windows(2).position(|window| window == b"\r\n") {
let line = String::from_utf8(self.buffer[..end].to_vec()).unwrap();
self.buffer.drain(..end + 2);
return line;
}
assert_ne!(self.fill().await, 0, "closed mid-line");
}
}
async fn fill(&mut self) -> usize {
let mut chunk = [0; 8192];
let read = self.stream.read(&mut chunk).await.unwrap_or(0);
self.buffer.extend_from_slice(&chunk[..read]);
read
}
/// Everything the peer sends until it closes, including bytes already buffered.
pub async fn until_closed(&mut self) -> Vec<u8> {
while self.fill().await != 0 {}
std::mem::take(&mut self.buffer)
}
/// Resolves once the peer closes the connection, discarding anything it still sends.
pub async fn closed(&mut self) {
let mut chunk = [0; 8192];
while self.stream.read(&mut chunk).await.unwrap_or(0) != 0 {}
}
}
fn owned(headers: &[httparse::Header<'_>]) -> Vec<(String, Vec<u8>)> {
headers
.iter()
.map(|header| (header.name.to_owned(), header.value.to_vec()))
.collect()
}

339
litellm-rust/crates/gateway-forward/tests/upload.rs Normal file
View file

@ -0,0 +1,339 @@
mod support;
use std::{
convert::Infallible,
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
time::Duration,
};
use axum::{
Router,
body::{Body, Bytes},
extract::Request,
http::StatusCode,
response::Response,
};
use http_body_util::BodyExt;
use rstest::rstest;
use support::{Gateway, Raw, client, raw_upstream, recording, serve, within};
use tokio::{
io::AsyncWriteExt,
sync::{Notify, mpsc},
};
fn ok() -> Response {
Response::new(Body::from("ok"))
}
#[rstest]
#[tokio::test]
async fn streams_upload_and_response_incrementally_without_a_body_limit() {
let upstream = serve(
Router::new().fallback(|request: Request| async { Response::new(request.into_body()) }),
)
.await;
let gateway = Gateway::start(upstream).await;
let (sender, receiver) = mpsc::channel::<Result<Bytes, Infallible>>(1);
let stream = futures_util::stream::unfold(receiver, |mut receiver| async {
receiver.recv().await.map(|chunk| (chunk, receiver))
});
sender.send(Ok(Bytes::from_static(b"first"))).await.unwrap();
let mut response = within(
client().request(
Request::post(gateway.url("/v1/files"))
.body(Body::from_stream(stream))
.unwrap(),
),
)
.await
.unwrap();
let frame = within(response.body_mut().frame()).await.unwrap().unwrap();
assert_eq!(frame.into_data().unwrap(), "first");
let large = Bytes::from(vec![b'x'; 8 * 1024 * 1024]);
sender.send(Ok(large.clone())).await.unwrap();
drop(sender);
let rest = within(response.into_body().collect())
.await
.unwrap()
.to_bytes();
assert_eq!(rest, large);
}
/// hyper sends the 100 as soon as the request is forwarded, where uvicorn waits for the app to
/// read the body, so an upload Python rejects on headers alone is still sent in full.
#[rstest]
#[tokio::test]
async fn answers_expect_continue_so_large_uploads_start() {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send("POST /v1/audio/transcriptions HTTP/1.1\r\nhost: x\r\nexpect: 100-continue\r\ncontent-length: 4\r\n\r\n")
.await;
let interim = within(raw.response_head()).await.unwrap();
assert_eq!(interim.status, StatusCode::CONTINUE);
raw.send("data").await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
assert_eq!(received.recv().await.unwrap().body, "data");
}
/// curl sends `expect: 100-continue` for any body over 1 MiB and waits for the 100 before sending
/// the multipart body, which must reach Python intact even where it looks like framing.
#[rstest]
#[tokio::test]
async fn expect_continue_delivers_a_large_multipart_upload_byte_for_byte() {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let boundary = "------------------------d74496d66958873e";
let audio: Vec<u8> = (0..3 * 1024 * 1024)
.map(|index: u32| (index.wrapping_mul(2_654_435_761) >> 24) as u8)
.chain(*b"\r\n0\r\n\r\nPOST /smuggled HTTP/1.1\r\n\r\n--")
.collect();
let body = [
format!("--{boundary}\r\ncontent-disposition: form-data; name=\"model\"\r\n\r\nwhisper-1\r\n--{boundary}\r\ncontent-disposition: form-data; name=\"file\"; filename=\"a b.wav\"\r\ncontent-type: audio/wav\r\n\r\n").into_bytes(),
audio,
format!("\r\n--{boundary}--\r\n").into_bytes(),
]
.concat();
let content_type = format!("multipart/form-data; boundary={boundary}");
let mut raw = gateway.connect().await;
raw.send(format!(
"POST /v1/audio/transcriptions HTTP/1.1\r\nhost: x\r\ncontent-type: {content_type}\r\nexpect: 100-continue\r\ncontent-length: {}\r\n\r\n",
body.len()
))
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::CONTINUE
);
raw.send(&body).await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
let request = received.recv().await.unwrap();
assert_eq!(request.headers["content-type"], content_type.as_str());
assert_eq!(request.body, body);
}
#[rstest]
#[case::get("GET")]
#[case::delete("DELETE")]
#[case::options("OPTIONS")]
#[case::post("POST")]
#[tokio::test]
async fn forwards_a_chunked_body_on_any_method(#[case] method: &str) {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(format!(
"{method} /v1/search HTTP/1.1\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n4\r\nbody\r\n0\r\n\r\n"
))
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
assert_eq!(received.recv().await.unwrap().body, "body");
}
#[rstest]
#[case::no_framing("", "", &[], "")]
#[case::empty_content_length("content-length: 0\r\n", "", &[("content-length", "0")], "")]
#[case::empty_chunked(
"transfer-encoding: chunked\r\n",
"0\r\n\r\n",
&[("transfer-encoding", "chunked")],
""
)]
#[case::sized(
"content-length: 7\r\n",
"payload",
&[("content-length", "7")],
"payload"
)]
#[case::chunk_extensions_and_trailers(
"transfer-encoding: chunked\r\n",
"4;name=\"v\"\r\nbody\r\n0\r\nx-checksum: 1\r\n\r\n",
&[("transfer-encoding", "chunked")],
"body"
)]
#[tokio::test]
async fn a_body_reaches_python_with_the_framing_the_client_chose(
#[values("POST", "GET")] method: &str,
#[case] framing: &str,
#[case] wire_body: &str,
#[case] python_framing: &[(&str, &str)],
#[case] python_body: &str,
) {
let (upstream, mut received) = recording(ok).await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(format!(
"{method} /v1/files HTTP/1.1\r\nhost: x\r\n{framing}\r\n{wire_body}"
))
.await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
let request = received.recv().await.unwrap();
let forwarded: Vec<(&str, &str)> = ["content-length", "transfer-encoding"]
.into_iter()
.flat_map(|name| {
request
.headers
.get_all(name)
.iter()
.map(move |value| (name, value.to_str().unwrap()))
})
.collect();
assert_eq!(request.method, method);
assert_eq!(forwarded, python_framing);
assert_eq!(request.body, python_body);
}
#[rstest]
#[tokio::test]
async fn a_large_upload_is_paced_by_python_reading_it() {
const LENGTH: usize = 128 * 1024 * 1024;
let start_reading = Arc::new(Notify::new());
let reading = start_reading.clone();
let upstream = raw_upstream(move |mut python| {
let reading = reading.clone();
async move {
let head = python.request_head().await.unwrap();
reading.notified().await;
let mut remaining = head.content_length().unwrap();
while remaining > 0 {
remaining -= python.body(remaining.min(1 << 20)).await.len();
}
python
.send("HTTP/1.1 200 OK\r\ncontent-length: 0\r\n\r\n")
.await;
}
})
.await;
let gateway = Gateway::start(upstream).await;
let (read, mut write) = gateway.connect().await.into_stream().into_split();
let sent = Arc::new(AtomicUsize::new(0));
let sending = sent.clone();
let upload = tokio::spawn(async move {
write
.write_all(
format!("POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: {LENGTH}\r\n\r\n")
.as_bytes(),
)
.await
.unwrap();
let chunk = [0x5a; 64 * 1024];
while sending.load(Ordering::SeqCst) < LENGTH {
write.write_all(&chunk).await.unwrap();
sending.fetch_add(chunk.len(), Ordering::SeqCst);
}
write
});
tokio::time::sleep(Duration::from_secs(1)).await;
assert!(
sent.load(Ordering::SeqCst) < LENGTH / 4,
"the forwarder buffered an upload Python was not reading"
);
start_reading.notify_one();
let write = tokio::time::timeout(Duration::from_secs(60), upload)
.await
.unwrap()
.unwrap();
let mut raw = Raw::new(read.reunite(write).unwrap());
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::OK
);
}
#[rstest]
#[case::between_chunks("5\r\nhello\r\n")]
#[case::mid_chunk("100\r\nhello")]
#[tokio::test]
async fn a_client_abandoning_a_chunked_upload_never_ends_python_body_cleanly(
#[case] sent: &'static str,
) {
let (forwarded, mut received) = mpsc::unbounded_channel();
let upstream = raw_upstream(move |mut python| {
let forwarded = forwarded.clone();
async move {
python.request_head().await.unwrap();
forwarded.send(python.until_closed().await).unwrap();
}
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(format!(
"POST /v1/files HTTP/1.1\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n{sent}"
))
.await;
tokio::time::sleep(Duration::from_millis(100)).await;
drop(raw);
let body = within(received.recv()).await.unwrap();
assert!(body.starts_with(b"5\r\nhello\r\n"));
assert!(
!body.ends_with(b"0\r\n\r\n"),
"python saw a truncated upload end cleanly"
);
}
#[rstest]
#[tokio::test]
async fn an_early_answer_during_an_upload_keeps_both_connections_in_sync() {
let (event, mut events) = mpsc::unbounded_channel();
let upstream = raw_upstream(move |mut python| {
let event = event.clone();
async move {
while let Some(head) = python.request_head().await {
let length = head.content_length();
event.send(head.path.into_bytes()).unwrap();
let Some(length) = length else {
python
.send("HTTP/1.1 200 OK\r\ncontent-length: 4\r\n\r\nnext")
.await;
continue;
};
python
.send("HTTP/1.1 413 Payload Too Large\r\ncontent-length: 2\r\n\r\n{}")
.await;
event.send(python.body(length).await).unwrap();
}
}
})
.await;
let gateway = Gateway::start(upstream).await;
let upload = b"GET /smuggled HTTP/1.1\r\nhost: x\r\n\r\n".repeat(1024);
let (first_half, second_half) = upload.split_at(upload.len() / 2);
let mut raw = gateway.connect().await;
raw.send(format!(
"POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: {}\r\n\r\n",
upload.len()
))
.await;
raw.send(first_half).await;
let rejected = within(raw.response_head()).await.unwrap();
assert_eq!(rejected.status, StatusCode::PAYLOAD_TOO_LARGE);
assert_eq!(raw.body(2).await, b"{}");
raw.send(second_half).await;
raw.send("GET /next HTTP/1.1\r\nhost: x\r\n\r\n").await;
let next = within(raw.response_head()).await.unwrap();
assert_eq!(raw.body(next.content_length().unwrap()).await, b"next");
assert_eq!(within(events.recv()).await.unwrap(), b"/v1/files");
assert_eq!(within(events.recv()).await.unwrap(), upload);
assert_eq!(within(events.recv()).await.unwrap(), b"/next");
}

21
litellm-rust/crates/gateway-forward/tests/upstream.rs Normal file
View file

@ -0,0 +1,21 @@
use litellm_gateway_forward::Upstream;
use rstest::rstest;
#[rstest]
#[case::https("https://127.0.0.1:4000")]
#[case::public("http://192.0.2.1:4000")]
#[case::hostname("http://localhost:4000")]
#[case::ipv6("http://[::1]:4000")]
#[case::path("http://127.0.0.1:4000/api")]
#[case::query("http://127.0.0.1:4000/?x=y")]
#[case::credentials("http://user@127.0.0.1:4000")]
fn rejects_anything_but_a_loopback_http_origin(#[case] upstream: &str) {
assert!(upstream.parse::<Upstream>().is_err());
}
#[rstest]
#[case::standard("http://127.0.0.1:4000")]
#[case::loopback_range("http://127.0.0.2:4000/")]
fn accepts_a_loopback_http_origin(#[case] upstream: &str) {
assert!(upstream.parse::<Upstream>().is_ok());
}

719
litellm-rust/crates/gateway-forward/tests/websocket.rs Normal file
View file

@ -0,0 +1,719 @@
mod support;
use std::{
net::SocketAddr,
sync::Arc,
time::{Duration, Instant},
};
use axum::{
Router,
http::{HeaderValue, StatusCode},
response::IntoResponse,
};
use futures_util::{SinkExt, StreamExt, future::join_all};
use litellm_gateway_forward::PEER_HEADER;
use rstest::rstest;
use support::{DRAIN, Gateway, Head, Raw, raw_upstream, serve, within};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpStream,
sync::{Barrier, mpsc},
};
use tokio_tungstenite::{
WebSocketStream, accept_hdr_async, client_async,
tungstenite::{
Error, Message,
client::IntoClientRequest,
handshake::{
derive_accept_key,
server::{ErrorResponse, Request as Handshake, Response as Accepted},
},
protocol::{CloseFrame, frame::coding::CloseCode},
},
};
const KEY: &str = "dGhlIHNhbXBsZSBub25jZQ==";
async fn connect(
address: SocketAddr,
connection: &'static str,
) -> Result<WebSocketStream<TcpStream>, Box<Error>> {
let mut request = format!("ws://{address}/v1/realtime?model=m")
.into_client_request()
.unwrap();
let headers = request.headers_mut();
headers.insert("authorization", HeaderValue::from_static("Bearer test"));
headers.insert(
"sec-websocket-protocol",
HeaderValue::from_static("realtime"),
);
headers.insert("connection", HeaderValue::from_static(connection));
let stream = TcpStream::connect(address).await.unwrap();
let (socket, response) = client_async(request, stream).await.map_err(Box::new)?;
assert_eq!(response.headers()["sec-websocket-protocol"], "realtime");
Ok(socket)
}
/// Greets, echoes every message, and closes with 4000 when told to "bye".
#[expect(
clippy::result_large_err,
reason = "tungstenite's handshake callback returns ErrorResponse"
)]
async fn echo() -> SocketAddr {
raw_upstream(|python| async move {
let callback = |request: &Handshake, mut response: Accepted| {
assert_eq!(request.uri(), "/v1/realtime?model=m");
assert_eq!(request.headers()["authorization"], "Bearer test");
assert!(request.headers().contains_key(PEER_HEADER));
response.headers_mut().insert(
"sec-websocket-protocol",
HeaderValue::from_static("realtime"),
);
Ok::<_, ErrorResponse>(response)
};
let mut socket = accept_hdr_async(python.into_stream(), callback)
.await
.unwrap();
socket.send(Message::text("greeting")).await.unwrap();
while let Some(Ok(message)) = socket.next().await {
if message == Message::text("bye") {
socket
.close(Some(CloseFrame {
code: CloseCode::Library(4000),
reason: "done".into(),
}))
.await
.unwrap();
continue;
}
if message.is_close() {
break;
}
socket.send(message).await.unwrap();
}
})
.await
}
fn handshake(target: &str, headers: &str) -> String {
format!(
"GET {target} HTTP/1.1\r\nhost: public.example\r\nconnection: Upgrade\r\nupgrade: websocket\r\nsec-websocket-version: 13\r\nsec-websocket-key: {KEY}\r\n{headers}\r\n"
)
}
/// The 101 uvicorn writes for `handshake`, with `headers` added.
fn switching(handshake: &Head, headers: &str) -> String {
let accept = derive_accept_key(handshake.header("sec-websocket-key").unwrap());
format!(
"HTTP/1.1 101 Switching Protocols\r\nupgrade: websocket\r\nconnection: Upgrade\r\nsec-websocket-accept: {accept}\r\n{headers}\r\n"
)
}
fn values(head: &Head, name: &str) -> Vec<String> {
head.headers
.iter()
.filter(|(key, _)| key.eq_ignore_ascii_case(name))
.map(|(_, value)| String::from_utf8(value.clone()).unwrap())
.collect()
}
/// Accepts with `headers` in the 101, hands the handshake head to the test, then echoes bytes.
async fn byte_echo(headers: &'static str) -> (SocketAddr, mpsc::UnboundedReceiver<Head>) {
let (sender, handshakes) = mpsc::unbounded_channel();
let upstream = raw_upstream(move |mut python: Raw| {
let sender = sender.clone();
async move {
let handshake = python.request_head().await.unwrap();
python.send(switching(&handshake, headers)).await;
sender.send(handshake).unwrap();
echo_bytes(python).await;
}
})
.await;
(upstream, handshakes)
}
async fn echo_bytes(python: Raw) {
let mut stream = python.into_stream();
let (mut read, mut write) = stream.split();
tokio::io::copy(&mut read, &mut write).await.ok();
}
async fn assert_echoes(socket: &mut TcpStream, message: &[u8]) {
socket.write_all(message).await.unwrap();
let mut echoed = vec![0; message.len()];
within(socket.read_exact(&mut echoed)).await.unwrap();
assert_eq!(echoed, message);
}
async fn open(gateway: &Gateway, handshake: String) -> (TcpStream, Head) {
let mut raw = gateway.connect().await;
raw.send(handshake).await;
let switched = within(raw.response_head()).await.unwrap();
assert_eq!(switched.status, StatusCode::SWITCHING_PROTOCOLS);
(raw.into_stream(), switched)
}
#[rstest]
#[tokio::test]
async fn relays_messages_both_ways(
#[values(
"Upgrade",
"keep-alive, Upgrade",
"upgrade",
"Upgrade, Sec-WebSocket-Key, Sec-WebSocket-Version"
)]
connection: &'static str,
#[values(
Message::text("hello"),
Message::binary(vec![0, 1, 2]),
Message::binary(vec![7; 4 * 1024 * 1024])
)]
message: Message,
) {
let gateway = Gateway::start(echo().await).await;
let mut socket = within(connect(gateway.address, connection)).await.unwrap();
assert_eq!(
socket.next().await.unwrap().unwrap(),
Message::text("greeting")
);
socket.send(message.clone()).await.unwrap();
assert_eq!(within(socket.next()).await.unwrap().unwrap(), message);
socket.close(None).await.unwrap();
while socket.next().await.is_some() {}
drop(socket);
within(gateway.shutdown()).await.unwrap();
}
#[rstest]
#[tokio::test]
async fn relays_a_close_python_initiates() {
let gateway = Gateway::start(echo().await).await;
let mut socket = within(connect(gateway.address, "Upgrade")).await.unwrap();
socket.next().await.unwrap().unwrap();
socket.send(Message::text("bye")).await.unwrap();
let Message::Close(Some(frame)) = within(socket.next()).await.unwrap().unwrap() else {
panic!("expected python's close frame");
};
assert_eq!(frame.code, CloseCode::Library(4000));
assert_eq!(frame.reason, "done");
}
#[rstest]
#[tokio::test]
async fn passes_an_upstream_rejection_through() {
let upstream = serve(
Router::new()
.fallback(|| async { (StatusCode::FORBIDDEN, "origin not allowed").into_response() }),
)
.await;
let gateway = Gateway::start(upstream).await;
let error = within(connect(gateway.address, "Upgrade"))
.await
.unwrap_err();
let Error::Http(response) = *error else {
panic!("expected the upstream rejection, got {error}");
};
assert_eq!(response.status(), StatusCode::FORBIDDEN);
assert_eq!(response.body().as_deref(), Some(&b"origin not allowed"[..]));
}
/// The idle timeout only covers a connection waiting for its next HTTP request, and a realtime
/// session may sit silent between turns.
#[rstest]
#[tokio::test]
async fn keeps_a_quiet_socket_open_past_the_keepalive() {
let gateway = Gateway::with_keepalive(echo().await, Duration::from_millis(400)).await;
let mut socket = within(connect(gateway.address, "Upgrade")).await.unwrap();
socket.next().await.unwrap().unwrap();
tokio::time::sleep(Duration::from_millis(1200)).await;
socket.send(Message::text("still here")).await.unwrap();
assert_eq!(
within(socket.next()).await.unwrap().unwrap(),
Message::text("still here")
);
}
/// Browsers carry the API key in a subprotocol, Python accepts with the first one it was offered,
/// and uvicorn negotiates permessage-deflate, so every handshake header must cross unchanged.
#[rstest]
#[tokio::test]
async fn handshake_headers_cross_unchanged_both_ways() {
let (upstream, mut handshakes) = byte_echo(
"sec-websocket-protocol: realtime\r\nsec-websocket-extensions: permessage-deflate; server_max_window_bits=12; client_max_window_bits=12\r\nset-cookie: a=1\r\nset-cookie: b=2\r\nx-litellm-call-id: abc\r\n",
)
.await;
let gateway = Gateway::start(upstream).await;
let target = "/v1/realtime?model=gpt-realtime&api-key=sk%2Fa%2Bb%3D&dup=1&dup=2";
let (_socket, switched) = open(
&gateway,
handshake(
target,
"sec-websocket-protocol: realtime, openai-insecure-api-key.sk-websocket-test, openai-beta.realtime-v1\r\nsec-websocket-protocol: second-line\r\nsec-websocket-extensions: permessage-deflate; client_max_window_bits\r\ncookie: token=abc\r\n",
),
)
.await;
let handshake = within(handshakes.recv()).await.unwrap();
assert_eq!(handshake.path, target);
assert_eq!(
values(&handshake, "sec-websocket-protocol"),
[
"realtime, openai-insecure-api-key.sk-websocket-test, openai-beta.realtime-v1",
"second-line"
]
);
assert_eq!(
values(&handshake, "sec-websocket-extensions"),
["permessage-deflate; client_max_window_bits"]
);
assert_eq!(values(&handshake, "cookie"), ["token=abc"]);
assert_eq!(values(&switched, "sec-websocket-protocol"), ["realtime"]);
assert_eq!(
values(&switched, "sec-websocket-extensions"),
["permessage-deflate; server_max_window_bits=12; client_max_window_bits=12"]
);
assert_eq!(values(&switched, "set-cookie"), ["a=1", "b=2"]);
assert_eq!(values(&switched, "x-litellm-call-id"), ["abc"]);
}
#[rstest]
#[tokio::test]
async fn invents_no_subprotocol_or_extension() {
let (upstream, mut handshakes) = byte_echo("").await;
let gateway = Gateway::start(upstream).await;
let (_socket, switched) = open(&gateway, handshake("/v1/realtime", "")).await;
let handshake = within(handshakes.recv()).await.unwrap();
for name in ["sec-websocket-protocol", "sec-websocket-extensions"] {
assert!(handshake.header(name).is_none());
assert!(switched.header(name).is_none());
}
}
#[rstest]
#[case::capitalized_upgrade("Upgrade: WebSocket\r\nConnection: keep-alive, Upgrade\r\n")]
#[case::uppercase("UPGRADE: WEBSOCKET\r\nCONNECTION: UPGRADE\r\n")]
#[case::split_connection("Upgrade: websocket\r\nConnection: keep-alive\r\nConnection: Upgrade\r\n")]
#[tokio::test]
async fn accepts_any_spelling_of_the_upgrade_headers(#[case] upgrade: &str) {
let (upstream, _handshakes) = byte_echo("").await;
let gateway = Gateway::start(upstream).await;
let (mut socket, _) = open(
&gateway,
format!(
"GET /v1/realtime HTTP/1.1\r\nhost: x\r\n{upgrade}sec-websocket-version: 13\r\nsec-websocket-key: {KEY}\r\n\r\n"
),
)
.await;
socket.write_all(b"ping").await.unwrap();
let mut echoed = [0; 4];
within(socket.read_exact(&mut echoed)).await.unwrap();
assert_eq!(&echoed, b"ping");
}
/// Compressed frames set RSV1 and every client frame is masked, so the socket is relayed as bytes,
/// never re-framed.
#[rstest]
#[tokio::test]
async fn relays_socket_bytes_exactly() {
let (upstream, _handshakes) = byte_echo("").await;
let gateway = Gateway::start(upstream).await;
let (socket, _) = open(&gateway, handshake("/v1/realtime", "")).await;
let sent: Vec<u8> = (0..3_000_017u32)
.map(|index| index.wrapping_mul(2_654_435_761).to_be_bytes()[1])
.collect();
let (mut read, mut write) = socket.into_split();
let writing = sent.clone();
let writer = tokio::spawn(async move {
for chunk in writing.chunks(7919) {
write.write_all(chunk).await.unwrap();
}
write
});
let mut echoed = vec![0; sent.len()];
within(read.read_exact(&mut echoed)).await.unwrap();
assert!(echoed == sent);
drop(writer.await.unwrap());
}
/// uvicorn writes the 101, an error event and a close frame back to back and then closes, so they
/// share a segment with the FIN, and a client may send its first frame with the handshake.
#[rstest]
#[tokio::test]
async fn keeps_bytes_that_arrive_with_the_handshake() {
let (early, mut received) = mpsc::unbounded_channel();
let upstream = raw_upstream(move |mut python: Raw| {
let early = early.clone();
async move {
let handshake = python.request_head().await.unwrap();
let mut segment = switching(&handshake, "").into_bytes();
segment.extend(b"\x81\x05error\x88\x08\x03\xf0denied");
python.send(segment).await;
early.send(within(python.body(6)).await).unwrap();
}
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(handshake("/v1/realtime", "") + "early!").await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::SWITCHING_PROTOCOLS
);
assert_eq!(within(received.recv()).await.unwrap(), b"early!");
assert_eq!(
within(raw.until_closed()).await,
b"\x81\x05error\x88\x08\x03\xf0denied"
);
}
/// Python frees a realtime session's budget and parallel slot when its socket closes, so a client
/// that disappears without a close frame must still close Python's socket.
#[rstest]
#[case::fin(false)]
#[case::reset(true)]
#[tokio::test]
async fn a_vanished_client_closes_the_python_socket(#[case] reset: bool) {
let (event, mut events) = mpsc::unbounded_channel();
let upstream = raw_upstream(move |mut python: Raw| {
let event = event.clone();
async move {
let handshake = python.request_head().await.unwrap();
python.send(switching(&handshake, "")).await;
event.send("open").unwrap();
python.closed().await;
event.send("closed").unwrap();
}
})
.await;
let gateway = Gateway::start(upstream).await;
let (socket, _) = open(&gateway, handshake("/v1/realtime", "")).await;
assert_eq!(within(events.recv()).await, Some("open"));
if reset {
#[expect(deprecated, reason = "a zero linger is how a test sends a TCP reset")]
socket.set_linger(Some(Duration::ZERO)).unwrap();
}
drop(socket);
assert_eq!(within(events.recv()).await, Some("closed"));
}
#[rstest]
#[case::fin(false)]
#[case::reset(true)]
#[tokio::test]
async fn a_vanished_python_ends_the_client_socket(#[case] reset: bool) {
let upstream = raw_upstream(move |mut python: Raw| async move {
let handshake = python.request_head().await.unwrap();
python.send(switching(&handshake, "")).await;
let socket = python.into_stream();
if reset {
#[expect(deprecated, reason = "a zero linger is how a test sends a TCP reset")]
socket.set_linger(Some(Duration::ZERO)).unwrap();
}
drop(socket);
})
.await;
let gateway = Gateway::start(upstream).await;
let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await;
let mut rest = Vec::new();
assert!(matches!(
within(socket.read_to_end(&mut rest)).await,
Ok(0) | Err(_)
));
}
#[rstest]
#[tokio::test]
async fn shutdown_closes_open_sockets_after_drain() {
let gateway = Gateway::start(echo().await).await;
let mut socket = within(connect(gateway.address, "Upgrade")).await.unwrap();
socket.next().await.unwrap().unwrap();
let started = Instant::now();
within(gateway.shutdown()).await.unwrap();
assert!(started.elapsed() >= DRAIN);
assert!(!matches!(socket.next().await, Some(Ok(Message::Text(_)))));
}
#[rstest]
#[tokio::test]
async fn shutdown_relays_the_close_python_sends_while_draining() {
let gateway = Gateway::start(echo().await).await;
let mut socket = within(connect(gateway.address, "Upgrade")).await.unwrap();
socket.next().await.unwrap().unwrap();
let server = gateway.shutdown();
tokio::time::sleep(Duration::from_millis(100)).await;
socket.send(Message::text("bye")).await.unwrap();
let Message::Close(Some(frame)) = within(socket.next()).await.unwrap().unwrap() else {
panic!("expected python's close frame");
};
assert_eq!(frame.code, CloseCode::Library(4000));
while socket.next().await.is_some() {}
drop(socket);
let started = Instant::now();
within(server).await.unwrap();
assert!(started.elapsed() < DRAIN / 2);
}
#[rstest]
#[tokio::test]
async fn refuses_new_sockets_while_draining() {
let gateway = Gateway::start(echo().await).await;
let mut open = within(connect(gateway.address, "Upgrade")).await.unwrap();
open.next().await.unwrap().unwrap();
let address = gateway.address;
let server = gateway.shutdown();
tokio::time::sleep(Duration::from_millis(100)).await;
assert!(
TcpStream::connect(address).await.is_err() || connect(address, "Upgrade").await.is_err()
);
open.send(Message::text("still open")).await.unwrap();
assert_eq!(
within(open.next()).await.unwrap().unwrap(),
Message::text("still open")
);
drop(open);
within(server).await.unwrap();
}
/// uvicorn never half-closes: once Python closes the socket nothing more can cross it, so the
/// relay must end instead of waiting for the client to close its side too.
#[rstest]
#[tokio::test]
async fn relay_ends_when_python_closes_while_the_client_holds_its_socket() {
let upstream = raw_upstream(move |mut python: Raw| async move {
let handshake = python.request_head().await.unwrap();
python.send(switching(&handshake, "")).await;
})
.await;
let gateway = Gateway::start(upstream).await;
let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await;
assert_eq!(within(socket.read(&mut [0; 1])).await.unwrap(), 0);
let started = Instant::now();
within(gateway.shutdown()).await.unwrap();
assert!(started.elapsed() < DRAIN / 2);
drop(socket);
}
/// uvicorn answers a handshake Python closes before accepting with a 403, and one Python denies
/// through `websocket.http.response` with any status, headers and body, then drops the socket. The
/// client gets that response whole and keeps a connection that still reaches Python.
#[rstest]
#[case::closed_before_accept("403 Forbidden", None, "")]
#[case::denial_response(
"401 Unauthorized",
Some("Bearer realm=\"litellm\""),
"{\"error\":\"invalid api key\"}"
)]
#[tokio::test]
async fn a_denied_handshake_reaches_the_client_whole(
#[case] status: &'static str,
#[case] authenticate: Option<&'static str>,
#[case] body: &'static str,
) {
let upstream = raw_upstream(move |mut python: Raw| async move {
python.request_head().await.unwrap();
let authenticate = authenticate
.map(|value| format!("www-authenticate: {value}\r\n"))
.unwrap_or_default();
python
.send(format!(
"HTTP/1.1 {status}\r\ndate: Tue, 03 Nov 2026 08:00:00 GMT\r\nconnection: close\r\ncontent-length: {}\r\ncontent-type: text/plain; charset=utf-8\r\n{authenticate}\r\n{body}",
body.len()
))
.await;
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
for _ in 0..2 {
raw.send(handshake("/v1/realtime", "")).await;
let denied = within(raw.response_head()).await.unwrap();
assert_eq!(denied.status.as_str(), &status[..3]);
assert_eq!(
values(&denied, "www-authenticate"),
Vec::from_iter(authenticate.map(String::from))
);
assert!(denied.header("sec-websocket-accept").is_none());
let length = denied.content_length().unwrap();
assert_eq!(within(raw.body(length)).await, body.as_bytes());
}
}
/// Only Python's EOF ends a relay, so a client that half-closes after its close frame still gets
/// everything Python sends before closing its own side.
#[rstest]
#[tokio::test]
async fn a_client_half_close_still_receives_what_python_sends_after_it() {
let upstream = raw_upstream(|mut python: Raw| async move {
let handshake = python.request_head().await.unwrap();
python.send(switching(&handshake, "")).await;
let received = python.until_closed().await;
python.send(received).await;
})
.await;
let gateway = Gateway::start(upstream).await;
let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await;
let close = b"\x88\x82\x01\x02\x03\x04\x02\xea";
socket.write_all(close).await.unwrap();
socket.shutdown().await.unwrap();
let mut replied = Vec::new();
within(socket.read_to_end(&mut replied)).await.unwrap();
assert_eq!(replied, close);
}
/// Python authenticates every handshake before accepting it, so sessions opened together must all
/// reach Python before any of them is accepted, and each must stay on its own socket.
#[rstest]
#[tokio::test]
async fn opens_many_sockets_at_once_without_crossing_them() {
const SOCKETS: usize = 64;
let everyone = Arc::new(Barrier::new(SOCKETS));
let upstream = raw_upstream(move |mut python: Raw| {
let everyone = everyone.clone();
async move {
let handshake = python.request_head().await.unwrap();
everyone.wait().await;
python.send(switching(&handshake, "")).await;
echo_bytes(python).await;
}
})
.await;
let gateway = Gateway::start(upstream).await;
let gateway = &gateway;
within(join_all((0..SOCKETS).map(|index| async move {
let (mut socket, _) = open(gateway, handshake("/v1/realtime", "")).await;
assert_echoes(&mut socket, format!("socket {index}").as_bytes()).await;
})))
.await;
}
/// Python resolves the key, budget and model before it accepts, and nothing bounds how long that
/// takes, so the idle timeout must not cut a handshake waiting on its 101.
#[rstest]
#[tokio::test]
async fn waits_for_a_101_slower_than_the_keepalive() {
let upstream = raw_upstream(|mut python: Raw| async move {
let handshake = python.request_head().await.unwrap();
tokio::time::sleep(Duration::from_millis(1200)).await;
python.send(switching(&handshake, "")).await;
echo_bytes(python).await;
})
.await;
let gateway = Gateway::with_keepalive(upstream, Duration::from_millis(400)).await;
let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await;
assert_echoes(&mut socket, b"late").await;
}
/// A 101 the client did not ask for, or to a protocol other than the one it asked for, would hand
/// the client a socket it cannot speak, so it is a 502 and Python's socket is closed.
#[rstest]
#[case::plain_request("GET /v1/models HTTP/1.1\r\nhost: x\r\n\r\n".to_owned(), "websocket")]
#[case::other_protocol(handshake("/v1/realtime", ""), "h2c")]
#[tokio::test]
async fn an_unrequested_switch_is_a_502_that_closes_python(
#[case] request: String,
#[case] protocol: &'static str,
) {
let (closed, mut closes) = mpsc::unbounded_channel();
let upstream = raw_upstream(move |mut python: Raw| {
let closed = closed.clone();
async move {
python.request_head().await.unwrap();
python
.send(format!(
"HTTP/1.1 101 Switching Protocols\r\nconnection: upgrade\r\nupgrade: {protocol}\r\n\r\n"
))
.await;
python.closed().await;
closed.send(()).unwrap();
}
})
.await;
let gateway = Gateway::start(upstream).await;
let mut raw = gateway.connect().await;
raw.send(request).await;
assert_eq!(
within(raw.response_head()).await.unwrap().status,
StatusCode::BAD_GATEWAY
);
within(closes.recv()).await.unwrap();
}
/// A realtime session streams audio for as long as the client reads it, so a client that stops
/// reading must stall Python's writes instead of piling them up in the gateway.
#[rstest]
#[tokio::test]
async fn a_client_that_stops_reading_stalls_python() {
let (sent, mut totals) = mpsc::unbounded_channel();
let upstream = raw_upstream(move |mut python: Raw| {
let sent = sent.clone();
async move {
let handshake = python.request_head().await.unwrap();
python.send(switching(&handshake, "")).await;
python.body(2).await;
let mut stream = python.into_stream();
let chunk = vec![0x82; 64 * 1024];
let mut total = 0;
while let Ok(Ok(())) =
tokio::time::timeout(Duration::from_millis(300), stream.write_all(&chunk)).await
{
total += chunk.len();
}
sent.send(total).unwrap();
}
})
.await;
let gateway = Gateway::start(upstream).await;
let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await;
socket.write_all(b"go").await.unwrap();
let total = within(totals.recv()).await.unwrap();
assert!(
total < 64 << 20,
"python wrote {total} bytes into a stalled socket"
);
}
/// Python closes a session that is still streaming audio by sending an error event and a close
/// frame and dropping its socket with the client's next frames unread, which resets it. Those last
/// bytes reached the client when it talked to uvicorn directly, so they must reach it through the
/// relay too, even though writing the client's frames to Python now fails.
#[rstest]
#[tokio::test]
async fn python_closing_mid_stream_still_delivers_its_last_frames() {
const LAST: &[u8] = b"\x81\x05error\x88\x02\x03\xf0";
let upstream = raw_upstream(|mut python: Raw| async move {
let handshake = python.request_head().await.unwrap();
python.send(switching(&handshake, "")).await;
python.body(1).await;
tokio::time::sleep(Duration::from_millis(20)).await;
python.send(LAST).await;
})
.await;
let gateway = Gateway::start(upstream).await;
for _ in 0..20 {
let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await;
let (mut read, mut write) = socket.split();
let streaming = async {
while write.write_all(b"\x82\x80\0\0\0\0").await.is_ok() {
tokio::time::sleep(Duration::from_micros(200)).await;
}
std::future::pending().await
};
let mut received = Vec::new();
let receiving = async {
let mut chunk = [0; 4096];
while let Ok(read @ 1..) = read.read(&mut chunk).await {
received.extend_from_slice(&chunk[..read]);
}
};
within(async {
tokio::select! {
() = streaming => {}
() = receiving => {}
}
})
.await;
assert_eq!(received, LAST);
}
}

7
litellm-rust/crates/gateway/AGENTS.md
View file

@ -5,8 +5,11 @@
- Test mounting and server lifecycle through public HTTP behavior. Test endpoint semantics in the owning crate
## Redis
- `build_redis` builds the Redis clients once at boot (`src/redis.rs`). Components get a clone as a field, never a per-request extractor or a pool. redis-rs multiplexes one connection per node, so there is no connection pool and pool settings like `max_connections` are accepted and ignored
- `build_redis` builds the Redis clients once at boot (`src/redis/mod.rs`). Components get a clone as a field, never a per-request extractor or a pool. redis-rs multiplexes one connection per node, so there is no connection pool and pool settings like `max_connections` are accepted and ignored
- Which Redis backs which role mirrors the proxy: coordination is `general_settings.coordination_redis`, else the `cache_params` Redis when `cache: true` and `type: redis`, else `REDIS_*` env. The transaction buffer reuses coordination, else env when `use_redis_transaction_buffer` is on. The response cache gets its own client so large values never queue in front of coordination commands
- `litellm-config` follows `litellm/_redis.py`'s merge of config blocks and `REDIS_*` env. A redis-py setting we can't honor (sentinel, IAM, custom TLS files, ...) fails boot with its name, never silently dropped
- `src/redis/layers.rs` is the Rust copy of `_merge_redis_kwargs` in `litellm/_redis.py`, which only this host needs because the Python host merges in Python. It puts the configured kwargs over their `REDIS_*` values and fills the `get_secret` fallbacks from the `SETTINGS` rows in `redis-resolve`, under the same conditions Python reads them, then hands the result to `litellm_redis_resolve::interpret`. Values come from a `SecretSnapshot` filled once at boot, read like `get_secret` (secret manager, then environment, with `true` and `false` read as bools), and the only straight environment read is Azure's `REDIS_USERNAME`
- A surface is a `ConfiguredRedis` constructor: `coordination` for `general_settings.coordination_redis`, which rejects a malformed `port` like the proxy's `CoordinationRedisParams` pydantic model, `cache_params` for `litellm_settings.cache_params`, which drops the cache-only keys and reads the `redis_startup_nodes` alias, and `redis_cache_kwargs` for a `RedisCache.redis_kwargs` dict in the contract tests. A new surface is a new constructor, never a flag
- Nulls depend on the surface. `coordination_redis` drops every null like `model_dump(exclude_none=True)`, `cache_params` keeps a null as an explicit `None` that hides the `REDIS_*` value unless `Cache` or `RedisCache` names that parameter, and `RedisCache.redis_kwargs` keeps every null. A missing `os.environ/` reference is a null, as `load_config` resolves it
- A setting the Rust client cannot honor warns and is dropped, and only a security check it cannot enforce fails boot, per `redis-resolve/AGENTS.md`
- Boot never waits for a configured Redis (clients connect lazily). Only the env-only coordination fallback is pinged, for 2s, and treated as absent when unreachable, as in Python
- `/health/readiness` never touches Redis, matching the proxy

8
litellm-rust/crates/gateway/Cargo.toml
View file

@ -12,6 +12,7 @@ envy = "0.4.2"
http-body-util = "0.1"
litellm-inference.workspace = true
litellm-gateway-inference.workspace = true
litellm-gateway-forward.workspace = true
litellm-gateway-mcp.workspace = true
tokio-util = "0.7"
thiserror.workspace = true
@ -19,9 +20,14 @@ futures-util.workspace = true
litellm-gateway-auth.workspace = true
litellm-gateway-ui.workspace = true
litellm-config.workspace = true
litellm-core-utils.workspace = true
litellm-error.workspace = true
litellm-circuit-breaker.workspace = true
litellm-redis-client.workspace = true
litellm-redis-connect.workspace = true
litellm-redis-resolve.workspace = true
litellm-redis-types.workspace = true
rustls.workspace = true
litellm-http.workspace = true
litellm-llms.workspace = true
litellm-secrets.workspace = true
@ -31,10 +37,12 @@ serde.workspace = true
tracing.workspace = true
tokio = { workspace = true, features = ["signal"] }
uuid.workspace = true
veil.workspace = true
tower-sessions-moka-store = "0.15.0"
[dev-dependencies]
litellm-redis-testing.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
base64.workspace = true
rstest.workspace = true
tempfile.workspace = true

11
litellm-rust/crates/gateway/src/error.rs
View file

@ -12,13 +12,12 @@ pub enum Error {
Mcp(#[from] litellm_gateway_mcp::ConnectError),
#[error(transparent)]
Config(#[from] litellm_config::Error),
#[error("invalid Redis setting {setting}: {reason}")]
RedisInvalid {
setting: &'static str,
reason: &'static str,
},
#[error(transparent)]
RedisConnect(#[from] litellm_redis_client::ConnectError),
Redis(#[from] litellm_redis_resolve::Error),
#[error("Redis setting {setting} is not plain data")]
RedisNotPlainData { setting: String },
#[error(transparent)]
RedisConnect(#[from] litellm_redis_connect::Error),
#[error("MCP requires working gateway authentication")]
Auth(#[from] litellm_gateway_auth::Error),
}

22
litellm-rust/crates/gateway/src/lib.rs
View file

@ -68,15 +68,29 @@ pub fn build_inference(config: &Config) -> Result<Arc<Gateway>, Error> {
Ok(Arc::new(Gateway::new(resources, http, secrets, router)?))
}
pub async fn build_redis(config: &Config) -> Result<RedisClients, Error> {
let environment = secrets::environment_values(&config.environment_variables)?;
let lookup = |name: &str| {
pub fn install_crypto_provider() {
if rustls::crypto::CryptoProvider::get_default().is_none() {
let _ = rustls::crypto::ring::default_provider().install_default();
}
}
pub async fn build_redis(
config: &Config,
pool: &HttpClientPool,
http: &litellm_http::HttpClientConfig,
) -> Result<RedisClients, Error> {
let environment = Arc::new(secrets::environment_values(&config.environment_variables)?);
let lookup: litellm_redis_connect::SharedEnvironment = Arc::new(move |name: &str| {
environment
.get(name)
.map(|value| value.expose().to_owned())
.or_else(|| std::env::var(name).ok())
});
let resources = redis::Resources {
http: pool.client(http, ClientVariant::Provider)?,
lookup: lookup.clone(),
};
redis::connect(config, &lookup).await
redis::connect(config, &*lookup, &resources).await
}
pub async fn build_mcp(

30
litellm-rust/crates/gateway/src/main.rs
View file

@ -10,7 +10,7 @@ use serde_json::json;
mod settings;
use settings::{Settings, UiSettings};
use settings::{LogSettings, Settings, UiSettings};
struct StderrSink {
level: Level,
@ -41,12 +41,23 @@ impl Sink for StderrSink {
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
litellm_gateway::install_crypto_provider();
if let Some(forward) = litellm_gateway_forward::Settings::from_env()? {
let log: LogSettings = settings::from_iter(std::env::vars_os())?;
Logger::new(StderrSink {
level: log.log_level(),
})
.install_global()?;
litellm_gateway_forward::run(forward, shutdown_signal()).await?;
return Ok(());
}
let settings: Settings = settings::from_iter(std::env::vars_os())?;
let level = settings.log_level();
Logger::new(StderrSink { level }).install_global()?;
let config = Config::load(settings.litellm_config)?;
let inference = litellm_gateway::build_inference(&config)?;
let _redis = litellm_gateway::build_redis(&config).await?;
let _redis =
litellm_gateway::build_redis(&config, &inference.resources.pool, &inference.http).await?;
let ui = match std::env::var_os("LITELLM_UI_PATH") {
Some(directory) => {
let ui_settings: UiSettings = settings::from_iter(std::env::vars_os())?;
@ -126,7 +137,20 @@ async fn shutdown_signal() {
}
}
}
#[cfg(not(unix))]
// Windows support exists for local CLI runs only, the proxy is never deployed on Windows. The supervisor
// stops the gateway with CTRL_BREAK_EVENT, which has to drain the same way SIGTERM does.
#[cfg(windows)]
{
match tokio::signal::windows::ctrl_break() {
Ok(mut ctrl_break) => {
tokio::select! { _ = tokio::signal::ctrl_c() => (), _ = ctrl_break.recv() => () }
}
Err(_) => {
let _ = tokio::signal::ctrl_c().await;
}
}
}
#[cfg(not(any(unix, windows)))]
{
let _ = tokio::signal::ctrl_c().await;
}

144
litellm-rust/crates/gateway/src/redis/configured.rs Normal file
View file

@ -0,0 +1,144 @@
use std::collections::BTreeMap;
use litellm_config::{RedisSettings, Value};
use litellm_core_utils::settings::Lookup;
use litellm_redis_resolve::{Expected, Kwargs};
use litellm_redis_types::{ProviderHandle, Resolution};
use veil::Redact;
use super::{layers, secrets::SecretSnapshot};
use crate::Error;
const CACHE_ONLY: &[&str] = &[
"type",
"mode",
"ttl",
"default_in_memory_ttl",
"default_in_redis_ttl",
"supported_call_types",
"similarity_threshold",
"max_messages",
"_backend",
];
/// The redis-py kwargs one surface hands `_get_redis_client_logic`, before `REDIS_*` and secrets.
#[derive(Clone, Redact)]
pub(crate) struct ConfiguredRedis {
#[redact(with = "[REDACTED]")]
pub(super) fields: BTreeMap<String, Value>,
pub(super) nulls: Nulls,
pub(super) credential_provider: Option<ProviderHandle>,
}
/// Whether a null setting reaches redis-py as an explicit `None`, which outranks its `REDIS_*` value.
#[derive(Clone, Copy, Debug)]
pub(super) enum Nulls {
/// `coordination_redis` reaches redis-py through `model_dump(exclude_none=True)`.
Dropped,
/// `Cache` drops a null only for the parameters it names and passes any other key on.
KeptUnlessNamedByCache,
/// `RedisCache.redis_kwargs` already holds only the nulls redis-py receives.
#[cfg(test)]
Kept,
}
impl ConfiguredRedis {
/// `general_settings.coordination_redis`, which the proxy validates as a pydantic model first.
/// An `os.environ/` port is only known after merging, where it is read like any other port.
pub(crate) fn coordination(
settings: &RedisSettings,
) -> Result<Self, litellm_redis_resolve::Error> {
if let Some(port) = settings.fields().get("port").and_then(text)
&& !port.starts_with("os.environ/")
&& port.parse::<u16>().is_err()
{
return Err(litellm_redis_resolve::Error::Invalid {
setting: "port".to_owned(),
expected: Expected::Integer,
});
}
Ok(Self {
fields: settings.fields().clone(),
nulls: Nulls::Dropped,
credential_provider: None,
})
}
/// `litellm_settings.cache_params` with `type: redis`, as `litellm.Cache` hands it to `RedisCache`.
pub(crate) fn cache_params(settings: &RedisSettings) -> Self {
let fields = settings.fields();
let alias = fields
.get("redis_startup_nodes")
.filter(|_| !fields.contains_key("startup_nodes"))
.map(|nodes| ("startup_nodes".to_owned(), nodes.clone()));
let fields = fields
.iter()
.filter(|(key, _)| {
key.as_str() != "redis_startup_nodes" && !CACHE_ONLY.contains(&key.as_str())
})
.map(|(key, value)| (key.clone(), value.clone()))
.chain(alias)
.collect();
Self {
fields,
nulls: Nulls::KeptUnlessNamedByCache,
credential_provider: None,
}
}
/// `RedisCache.redis_kwargs`, the dict every Python surface ends in.
#[cfg(test)]
pub(crate) fn redis_cache_kwargs(fields: BTreeMap<String, Value>) -> Self {
Self {
fields,
nulls: Nulls::Kept,
credential_provider: None,
}
}
/// The handle the host holds for the `credential_provider` object in these settings.
#[cfg(test)]
pub(crate) fn with_credential_provider(self, handle: ProviderHandle) -> Self {
Self {
credential_provider: Some(handle),
..self
}
}
/// No configured settings, so only `REDIS_*` and the other environment fallbacks apply.
pub(crate) fn environment_only() -> Self {
Self {
fields: BTreeMap::new(),
nulls: Nulls::Dropped,
credential_provider: None,
}
}
/// `_merge_redis_kwargs` in `litellm/_redis.py`: the kwargs over their `REDIS_*` values with every
/// `get_secret` fallback filled.
pub(crate) fn merged(&self, secrets: &SecretSnapshot) -> Result<Kwargs, Error> {
layers::merged(self, secrets)
}
/// The environment is only for what Python reads straight from `os.environ`, `REDIS_USERNAME` for Azure.
pub(crate) fn resolve(
&self,
secrets: &SecretSnapshot,
environment: &dyn Lookup,
) -> Result<Resolution, Error> {
Ok(litellm_redis_resolve::interpret(
&self.merged(secrets)?,
self.credential_provider,
environment.get("REDIS_USERNAME").as_deref(),
)?)
}
}
pub(super) fn text(value: &Value) -> Option<String> {
match value {
Value::String(text) => Some(text.clone()),
Value::Number(number) => Some(number.to_string()),
Value::Bool(flag) => Some(flag.to_string()),
_ => None,
}
}

226
litellm-rust/crates/gateway/src/redis/layers.rs Normal file
View file

@ -0,0 +1,226 @@
use litellm_config::Value;
use litellm_core_utils::settings::Lookup;
use litellm_redis_resolve::{Entry, Kwargs, SETTINGS, is_present, setting};
use super::configured::{ConfiguredRedis, Nulls};
use crate::Error;
const NAMED_BY_CACHE: &[&str] = &[
"host",
"port",
"password",
"socket_timeout",
"startup_nodes",
"namespace",
"gcp_service_account",
"gcp_ssl_ca_certs",
];
/// How `_merge_redis_kwargs` reads a fallback: `get_secret` turns `true` and `false` into a bool, and
/// `get_secret_str` treats that bool as unset.
#[derive(Clone, Copy)]
enum Read {
Secret,
SecretStr,
}
const TARGETS: &[(&str, Read)] = &[
("startup_nodes", Read::Secret),
("sentinel_nodes", Read::Secret),
("sentinel_password", Read::SecretStr),
("service_name", Read::Secret),
];
const IDENTITIES: &[(&str, Read)] = &[
("gcp_service_account", Read::SecretStr),
("gcp_ssl_ca_certs", Read::SecretStr),
("azure_redis_ad_token", Read::Secret),
];
const AZURE: &[(&str, Read)] = &[
("azure_client_id", Read::SecretStr),
("azure_tenant_id", Read::SecretStr),
("azure_client_secret", Read::SecretStr),
];
const AWS: &[(&str, Read)] = &[("aws_iam_region", Read::SecretStr)];
/// The configured settings over their `REDIS_*` values, `{**environment_kwargs, **env_overrides}`
/// in `_merge_redis_kwargs`, with the `get_secret` fallbacks it reads filled in. A null left in the result is
/// an explicit `None` that hid its environment value.
pub(super) fn merged(configured: &ConfiguredRedis, lookup: &dyn Lookup) -> Result<Kwargs, Error> {
let block: Kwargs = configured
.fields
.iter()
.map(|(key, value)| (key, environment_reference(value, lookup)))
.filter(|(key, value)| !value.is_null() || keeps_null(configured.nulls, key))
.map(|(key, value)| {
let value = serde_json::to_value(&value).map_err(|_| Error::RedisNotPlainData {
setting: key.clone(),
})?;
let entry = Entry {
value,
source: key.clone(),
};
Ok((key.clone(), entry))
})
.collect::<Result<_, Error>>()?;
let names_a_target = ["host", "startup_nodes", "sentinel_nodes"]
.iter()
.any(|key| is_set(&block, key));
let drop_environment_url = names_a_target && !is_set(&block, "url");
let layered: Kwargs = environment(lookup)
.filter(|(key, _)| !(key == "url" && drop_environment_url))
.chain(block)
.collect();
let fallbacks = fallbacks(&layered, configured.credential_provider.is_some(), lookup);
Ok(layered.into_iter().chain(fallbacks).collect())
}
// The fallbacks are read under the same conditions as _merge_redis_kwargs, so a secret Python never asks
// for is never filled in
fn fallbacks(layered: &Kwargs, has_provider: bool, lookup: &dyn Lookup) -> Kwargs {
let fill = |kwargs: &[(&str, Read)]| -> Kwargs {
kwargs
.iter()
.filter_map(|(kwarg, read)| fallback(layered, kwarg, *read, lookup))
.collect()
};
let targets = fill(TARGETS);
if has_provider || is_set(layered, "credential_provider") {
return targets;
}
let identities = fill(IDENTITIES);
let value = |kwarg: &str| {
layered
.get(kwarg)
.filter(|entry| is_present(&entry.value))
.or_else(|| identities.get(kwarg))
.map(|entry| &entry.value)
};
let gcp = value("gcp_service_account").is_some();
let azure = value("azure_redis_ad_token")
.is_some_and(|token| text(token).is_some_and(|token| token.eq_ignore_ascii_case("true")));
let aws = layered
.get("aws_iam_auth")
.is_some_and(|entry| coerces_to_true(&entry.value));
let azure_settings = match azure && !gcp {
true => fill(AZURE),
false => Kwargs::new(),
};
let aws_settings = match aws && !gcp && !azure {
true => fill(AWS),
false => Kwargs::new(),
};
targets
.into_iter()
.chain(identities)
.chain(azure_settings)
.chain(aws_settings)
.collect()
}
fn fallback(
layered: &Kwargs,
kwarg: &str,
read: Read,
lookup: &dyn Lookup,
) -> Option<(String, Entry)> {
if layered
.get(kwarg)
.is_some_and(|entry| is_present(&entry.value))
{
return None;
}
let names = setting(kwarg).map_or(&[][..], |setting| setting.fallback);
names.iter().find_map(|name| {
let value = lookup
.get(name)
.filter(|value| !value.is_empty())
.map(secret)?;
match (read, &value) {
(Read::SecretStr, serde_json::Value::Bool(_)) => None,
_ => {
let entry = Entry {
value,
source: (*name).to_owned(),
};
Some((kwarg.to_owned(), entry))
}
}
})
}
// get_secret reads an environment value of true or false, in any case and padding, as a bool
fn secret_flag(value: &str) -> Option<bool> {
match value.trim().to_ascii_lowercase().as_str() {
"true" => Some(true),
"false" => Some(false),
_ => None,
}
}
fn secret(value: String) -> serde_json::Value {
secret_flag(&value).map_or(serde_json::Value::String(value), serde_json::Value::Bool)
}
// _coerces_to_true: text is read like _str_to_bool, anything else by its truthiness
fn coerces_to_true(value: &serde_json::Value) -> bool {
match value {
serde_json::Value::String(text) => {
matches!(text.to_ascii_lowercase().as_str(), "true" | "1" | "yes")
}
serde_json::Value::Null => false,
serde_json::Value::Bool(flag) => *flag,
serde_json::Value::Number(number) => number.as_f64().is_some_and(|number| number != 0.0),
serde_json::Value::Array(items) => !items.is_empty(),
serde_json::Value::Object(entries) => !entries.is_empty(),
}
}
fn text(value: &serde_json::Value) -> Option<String> {
match value {
serde_json::Value::String(text) => Some(text.clone()),
serde_json::Value::Number(number) => Some(number.to_string()),
serde_json::Value::Bool(flag) => Some(flag.to_string()),
_ => None,
}
}
fn keeps_null(nulls: Nulls, key: &str) -> bool {
match nulls {
Nulls::Dropped => false,
Nulls::KeptUnlessNamedByCache => !NAMED_BY_CACHE.contains(&key),
#[cfg(test)]
Nulls::Kept => true,
}
}
fn environment(lookup: &dyn Lookup) -> impl Iterator<Item = (String, Entry)> + '_ {
SETTINGS.iter().filter_map(|setting| {
let source = setting.environment_name()?;
let value = lookup.get(&source)?;
let entry = Entry {
value: secret(value),
source,
};
Some((setting.kwarg.to_owned(), entry))
})
}
// load_config resolves os.environ/ to None when unset, so a missing reference is a null like any other
fn environment_reference(value: &Value, lookup: &dyn Lookup) -> Value {
match value {
Value::String(text) => match text.strip_prefix("os.environ/") {
Some(name) => lookup.get(name).map_or(Value::Null, |value| {
secret_flag(&value).map_or(Value::String(value), Value::Bool)
}),
None => value.clone(),
},
_ => value.clone(),
}
}
fn is_set(kwargs: &Kwargs, key: &str) -> bool {
kwargs.get(key).is_some_and(|entry| !entry.value.is_null())
}

410
litellm-rust/crates/gateway/src/redis.rs → litellm-rust/crates/gateway/src/redis/mod.rs
View file

@ -1,19 +1,31 @@
use std::time::Duration;
use std::{sync::Arc, time::Duration};
use litellm_circuit_breaker::{Breaker, Config as BreakerConfig, State};
use litellm_config::{
CacheSettings, Config, Lookup, RedisEndpoint, RedisSettings, ResolvedRedis, Value,
environment_has_target,
};
use litellm_redis_client::{
ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, Namespace, RedisClient,
RedisConfig,
};
use litellm_circuit_breaker::Config as BreakerConfig;
use litellm_config::{CacheSettings, Config, RedisSettings, Value};
use litellm_core_utils::settings::Lookup;
use litellm_redis_client::{RedisClient, RedisEvents};
use litellm_redis_connect::{HostServices, RedisHosts, SharedEnvironment, TracingEvents};
use litellm_redis_resolve::Expected;
use litellm_redis_types::Resolution;
use crate::Error;
mod configured;
mod layers;
mod secrets;
#[cfg(test)]
mod surface_tests;
use configured::ConfiguredRedis;
use secrets::SecretSnapshot;
pub(crate) struct Resources {
pub(crate) http: litellm_http::Client,
pub(crate) lookup: SharedEnvironment,
}
const ENVIRONMENT_PING_TIMEOUT: Duration = Duration::from_secs(2);
type Params = ResolvedRedis;
type Params = Resolution;
pub struct RedisClients {
pub coordination: Option<RedisClient>,
@ -34,12 +46,15 @@ enum Target {
Environment(Params),
}
pub(crate) fn plan(config: &Config, lookup: Lookup<'_>) -> Result<Plan, Error> {
pub(crate) fn plan(config: &Config, lookup: &dyn Lookup) -> Result<Plan, Error> {
let response_cache = redis_cache_params(config)
.map(|settings| settings.resolve_cache(lookup))
.map(|settings| resolve(&ConfiguredRedis::cache_params(settings), lookup))
.transpose()?;
let coordination = match &config.general_settings.coordination_redis {
Some(settings) => Some(Target::Configured(settings.resolve(lookup)?)),
Some(settings) => Some(Target::Configured(resolve(
&ConfiguredRedis::coordination(settings)?,
lookup,
)?)),
None => match (&response_cache, config.litellm_settings.cache == Some(true)) {
(Some(cache), _) => Some(Target::Configured(cache.clone())),
(None, true) => environment(lookup)?.map(Target::Configured),
@ -57,12 +72,30 @@ pub(crate) fn plan(config: &Config, lookup: Lookup<'_>) -> Result<Plan, Error> {
})
}
pub async fn connect(config: &Config, lookup: Lookup<'_>) -> Result<RedisClients, Error> {
pub(crate) async fn connect(
config: &Config,
lookup: &dyn Lookup,
resources: &Resources,
) -> Result<RedisClients, Error> {
let plan = plan(config, lookup)?;
let breaker = breaker_config(lookup)?;
let events: Arc<dyn RedisEvents> = Arc::new(TracingEvents::new(lookup));
let hosts = RedisHosts::default();
let client = |params: Params| -> Result<RedisClient, Error> {
Ok(hosts.client(
params,
HostServices {
events: Some(events.clone()),
http: resources.http.clone(),
environment: resources.lookup.clone(),
credential_provider: None,
},
breaker,
)?)
};
let coordination = match plan.coordination {
Some(Target::Configured(params)) => Some(client(params, breaker)?),
Some(Target::Environment(params)) => match client(params, breaker) {
Some(Target::Configured(params)) => Some(client(params)?),
Some(Target::Environment(params)) => match client(params) {
Ok(client) => reachable(client).await,
Err(_) => {
tracing::warn!(
@ -75,13 +108,10 @@ pub async fn connect(config: &Config, lookup: Lookup<'_>) -> Result<RedisClients
};
let transaction_buffer = match (&coordination, plan.transaction_buffer_fallback) {
(Some(coordination), _) => Some(coordination.clone()),
(None, Some(params)) => Some(client(params, breaker)?),
(None, Some(params)) => Some(client(params)?),
(None, None) => None,
};
let response_cache = plan
.response_cache
.map(|params| client(params, breaker))
.transpose()?;
let response_cache = plan.response_cache.map(client).transpose()?;
Ok(RedisClients {
coordination,
transaction_buffer,
@ -98,17 +128,38 @@ fn redis_cache_params(config: &Config) -> Option<&RedisSettings> {
}
}
fn environment(lookup: Lookup<'_>) -> Result<Option<Params>, Error> {
fn resolve(configured: &ConfiguredRedis, lookup: &dyn Lookup) -> Result<Params, Error> {
let secrets = SecretSnapshot::read(configured.secret_names(), lookup);
let resolution = configured.resolve(&secrets, lookup)?;
for ignored in &resolution.ignored {
tracing::warn!(
"Redis setting {} is not supported by the Rust client and is ignored: {}",
ignored.setting,
ignored.consequence
);
}
Ok(resolution)
}
fn environment(lookup: &dyn Lookup) -> Result<Option<Params>, Error> {
match environment_has_target(lookup) {
true => RedisSettings::default()
.resolve(lookup)
.map(Some)
.map_err(Into::into),
true => Ok(Some(resolve(&ConfiguredRedis::environment_only(), lookup)?)),
false => Ok(None),
}
}
fn optional_environment(lookup: Lookup<'_>) -> Option<Params> {
fn environment_has_target(lookup: &dyn Lookup) -> bool {
[
"REDIS_HOST",
"REDIS_URL",
"REDIS_CLUSTER_NODES",
"REDIS_SENTINEL_NODES",
]
.iter()
.any(|name| lookup.get(name).is_some())
}
fn optional_environment(lookup: &dyn Lookup) -> Option<Params> {
match environment(lookup) {
Ok(params) => params,
Err(error) => {
@ -149,122 +200,101 @@ async fn reachable(client: RedisClient) -> Option<RedisClient> {
}
}
fn client(params: Params, breaker: Option<BreakerConfig>) -> Result<RedisClient, Error> {
tracing::info!(endpoint = ?params.endpoint, namespace = ?params.namespace, "redis client configured");
let endpoint = match params.endpoint {
RedisEndpoint::Url(url) => Endpoint::Url(url),
RedisEndpoint::Host(host) => Endpoint::Host(HostEndpoint {
host: host.host,
port: host.port,
username: host.username,
password: host.password,
tls: host.tls,
db: host.db,
}),
RedisEndpoint::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint {
nodes: cluster
.nodes
.into_iter()
.map(|node| ClusterNode {
host: node.host,
port: node.port,
})
.collect(),
username: cluster.username,
password: cluster.password,
tls: cluster.tls,
}),
};
Ok(RedisClient::new(RedisConfig {
endpoint,
namespace: params.namespace.map(Namespace::new).unwrap_or_default(),
timeout: params.timeout.unwrap_or(DEFAULT_TIMEOUT),
breaker: breaker.map(|config| Breaker::new(config, Box::new(log_transition))),
})?)
fn parse_seconds(text: &str) -> Option<Duration> {
text.parse::<f64>()
.ok()
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
}
fn log_transition(transition: &litellm_circuit_breaker::Transition<State>) {
match transition.to {
State::Open { .. } => tracing::warn!(
generation = transition.generation,
"redis circuit breaker opened, fast-failing redis calls"
),
State::HalfOpen => tracing::info!(
generation = transition.generation,
"redis circuit breaker half open, probing redis"
),
State::Closed { .. } => tracing::info!(
generation = transition.generation,
"redis circuit breaker closed, redis recovered"
),
}
}
fn breaker_config(lookup: Lookup<'_>) -> Result<Option<BreakerConfig>, Error> {
let enabled = lookup("REDIS_CIRCUIT_BREAKER_ENABLED")
fn breaker_config(lookup: &dyn Lookup) -> Result<Option<BreakerConfig>, Error> {
let enabled = lookup
.get("REDIS_CIRCUIT_BREAKER_ENABLED")
.is_none_or(|value| value.eq_ignore_ascii_case("true"));
if !enabled {
return Ok(None);
}
let defaults = BreakerConfig::default();
let seconds = |name: &'static str, default: Duration| -> Result<Duration, Error> {
lookup(name)
.map(|value| {
value
.parse::<f64>()
.ok()
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
.ok_or(Error::RedisInvalid {
setting: name,
reason: "must be a non-negative number of seconds",
})
})
.unwrap_or(Ok(default))
};
let failure_threshold = lookup("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD")
Ok(Some(BreakerConfig {
failure_threshold: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD",
Expected::Integer,
|text| text.parse().ok(),
)?
.unwrap_or(defaults.failure_threshold),
recovery_timeout: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT",
Expected::Seconds,
parse_seconds,
)?
.unwrap_or(defaults.recovery_timeout),
timeout_min_duration: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION",
Expected::Seconds,
parse_seconds,
)?
.unwrap_or(defaults.timeout_min_duration),
}))
}
fn environment_value<T>(
lookup: &dyn Lookup,
name: &str,
expected: Expected,
parse: impl Fn(&str) -> Option<T>,
) -> Result<Option<T>, litellm_redis_resolve::Error> {
lookup
.get(name)
.map(|value| {
value.parse().map_err(|_| Error::RedisInvalid {
setting: "REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD",
reason: "must be an integer",
parse(&value).ok_or_else(|| litellm_redis_resolve::Error::Invalid {
setting: name.to_owned(),
expected,
})
})
.unwrap_or(Ok(defaults.failure_threshold))?;
Ok(Some(BreakerConfig {
failure_threshold,
recovery_timeout: seconds(
"REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT",
defaults.recovery_timeout,
)?,
timeout_min_duration: seconds(
"REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION",
defaults.timeout_min_duration,
)?,
}))
.transpose()
}
#[cfg(test)]
mod tests {
use litellm_core_utils::settings::Lookup;
use std::{collections::HashMap, time::Duration};
use litellm_config::{Config, RedisEndpoint, RedisHost};
use litellm_config::Config;
use litellm_redis_client::{Standalone, StaticAuth, commands};
use litellm_redis_testing::RedisServer;
use litellm_redis_types::{Auth, HealthCheck};
use rstest::rstest;
use super::{Params, Plan, Target, connect, plan};
use std::sync::Arc;
use super::{Params, Plan, Resources, Target, plan};
use crate::Error;
async fn connect(config: &Config, lookup: &dyn Lookup) -> Result<super::RedisClients, Error> {
let resources = Resources {
http: litellm_http::Client::for_test(),
lookup: Arc::new(|_: &str| None),
};
super::connect(config, lookup, &resources).await
}
fn host(name: &str, port: u16) -> Params {
Params {
endpoint: RedisEndpoint::Host(RedisHost {
target: litellm_redis_client::Target::Standalone(Standalone {
host: name.to_owned(),
port,
username: None,
password: None,
tls: false,
tls: None,
db: 0,
}),
auth: Auth::Static(StaticAuth::default()),
namespace: None,
timeout: None,
connect_timeout: None,
health_check: HealthCheck::Unset,
keepalive: true,
ignored: Vec::new(),
}
}
@ -273,7 +303,7 @@ mod tests {
.iter()
.map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
.collect();
plan(&Config::from_yaml(yaml).unwrap(), &|name| {
plan(&Config::from_yaml(yaml).unwrap(), &|name: &str| {
environment.get(name).cloned()
})
}
@ -350,36 +380,44 @@ mod tests {
port_of(&server)
))
.unwrap();
let clients = connect(&config, &|_| None).await.unwrap();
let clients = connect(&config, &|_: &str| None).await.unwrap();
let coordination = clients.coordination.unwrap();
let buffer = clients.transaction_buffer.unwrap();
let cache = clients.response_cache.unwrap();
let unnamespaced = connect(&coordination_config(&server), &|_| None)
let unnamespaced = connect(&coordination_config(&server), &|_: &str| None)
.await
.unwrap()
.coordination
.unwrap();
coordination
.set("rate-limit", "1", Duration::from_secs(60))
.send(commands::set("rate-limit", "1", Duration::from_secs(60)))
.await
.unwrap();
buffer
.set("spend", "2", Duration::from_secs(60))
.send(commands::set("spend", "2", Duration::from_secs(60)))
.await
.unwrap();
assert_eq!(
cache.get::<String>("rate-limit").await.unwrap().as_deref(),
cache
.send(commands::get::<String>("rate-limit"))
.await
.unwrap()
.as_deref(),
Some("1")
);
assert_eq!(
cache.get::<String>("spend").await.unwrap().as_deref(),
cache
.send(commands::get::<String>("spend"))
.await
.unwrap()
.as_deref(),
Some("2")
);
assert_eq!(
unnamespaced
.get::<String>("responses:rate-limit")
.send(commands::get::<String>("responses:rate-limit"))
.await
.unwrap()
.as_deref(),
@ -387,7 +425,7 @@ mod tests {
);
assert_eq!(
unnamespaced
.get::<String>("responses:spend")
.send(commands::get::<String>("responses:spend"))
.await
.unwrap()
.as_deref(),
@ -395,16 +433,21 @@ mod tests {
);
assert!(
unnamespaced
.get::<String>("rate-limit")
.send(commands::get::<String>("rate-limit"))
.await
.unwrap()
.is_none()
);
assert!(
unnamespaced
.send(commands::get::<String>("spend"))
.await
.unwrap()
.is_none()
);
assert!(unnamespaced.get::<String>("spend").await.unwrap().is_none());
}
#[rstest]
#[case::malformed_port(&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")])]
#[case::malformed_cluster_nodes(&[("REDIS_CLUSTER_NODES", "[{\"host\": ")])]
fn redis_migration_regression_malformed_optional_environment_is_absent(
#[case] environment: &[(&str, &str)],
@ -413,54 +456,45 @@ mod tests {
}
#[rstest]
fn redis_migration_regression_explicit_malformed_setting_still_fails() {
#[case::optional_environment("model_list: []", &[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")])]
#[case::transaction_buffer_environment(
"general_settings:\n use_redis_transaction_buffer: true",
&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")]
)]
#[case::response_cache(
"litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: env\n port: not-a-port",
&[]
)]
fn a_malformed_port_outside_the_coordination_block_falls_back_to_6379_like_redis_py(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
) {
let plan = resolved(yaml, environment).unwrap();
let target = match (plan.coordination, plan.transaction_buffer_fallback) {
(Some(Target::Configured(params) | Target::Environment(params)), _)
| (None, Some(params)) => params,
(None, None) => panic!("expected a Redis target"),
};
assert_eq!(target.target, host("env", 6379).target);
}
#[rstest]
#[case::port("port: not-a-port", "port")]
#[case::ssl("ssl: maybe", "ssl")]
fn a_malformed_coordination_block_fails_like_the_pydantic_model(
#[case] setting: &str,
#[case] expected: &str,
) {
let error = resolved(
"general_settings:\n coordination_redis:\n host: coord\n port: not-a-port",
&format!("general_settings:\n coordination_redis:\n host: coord\n {setting}"),
&[],
)
.unwrap_err();
assert!(matches!(
error,
Error::Config(litellm_config::Error::RedisInvalid {
setting: "port",
..
})
));
}
#[rstest]
fn redis_migration_regression_enabled_transaction_buffer_rejects_invalid_environment() {
let error = resolved(
"general_settings:\n use_redis_transaction_buffer: true",
&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")],
)
.unwrap_err();
assert!(matches!(
error,
Error::Config(litellm_config::Error::RedisInvalid {
setting: "port",
..
})
));
}
#[rstest]
fn redis_migration_regression_invalid_response_cache_still_fails() {
let error = resolved(
"litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: cache\n port: not-a-port",
&[],
)
.unwrap_err();
assert!(matches!(
error,
Error::Config(litellm_config::Error::RedisInvalid {
setting: "port",
..
})
));
assert!(
matches!(error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, .. }) if setting == expected)
);
}
#[rstest]
@ -468,7 +502,7 @@ mod tests {
async fn redis_migration_regression_invalid_optional_environment_url_is_absent() {
let config = Config::from_yaml("model_list: []").unwrap();
let clients = connect(&config, &|name| {
let clients = connect(&config, &|name: &str| {
(name == "REDIS_URL").then(|| "not a redis URL".to_owned())
})
.await
@ -506,7 +540,7 @@ mod tests {
let server = RedisServer::start().await.unwrap();
server.stop().await.unwrap();
let clients = connect(&coordination_config(&server), &|_| None)
let clients = connect(&coordination_config(&server), &|_: &str| None)
.await
.unwrap();
let coordination = clients.coordination.unwrap();
@ -521,12 +555,12 @@ mod tests {
#[tokio::test]
async fn the_coordination_client_survives_a_redis_restart_and_backs_the_transaction_buffer() {
let server = RedisServer::start().await.unwrap();
let clients = connect(&coordination_config(&server), &|_| None)
let clients = connect(&coordination_config(&server), &|_: &str| None)
.await
.unwrap();
let coordination = clients.coordination.unwrap();
coordination
.set("written", "1", Duration::from_secs(60))
.send(commands::set("written", "1", Duration::from_secs(60)))
.await
.unwrap();
@ -537,6 +571,30 @@ mod tests {
assert!(buffer.ping().await.is_ok());
}
#[rstest]
#[tokio::test]
async fn an_environment_only_tls_redis_is_reached_with_its_ca() {
crate::install_crypto_provider();
let server = litellm_redis_testing::RedisTlsServer::start(false)
.await
.unwrap();
let ca = server.files().ca.display().to_string();
let port = server.port().to_string();
let config = Config::from_yaml("model_list: []").unwrap();
let clients = connect(&config, &|name: &str| match name {
"REDIS_HOST" => Some(server.host().to_owned()),
"REDIS_PORT" => Some(port.clone()),
"REDIS_SSL" => Some("true".to_owned()),
"REDIS_SSL_CA_CERTS" => Some(ca.clone()),
_ => None,
})
.await
.unwrap();
assert!(clients.coordination.is_some());
}
#[rstest]
#[tokio::test]
async fn an_unreachable_environment_only_redis_leaves_coordination_in_memory() {
@ -545,7 +603,7 @@ mod tests {
server.stop().await.unwrap();
let config = Config::from_yaml("model_list: []").unwrap();
let clients = connect(&config, &|name| match name {
let clients = connect(&config, &|name: &str| match name {
"REDIS_HOST" => Some("127.0.0.1".to_owned()),
"REDIS_PORT" => Some(port.clone()),
_ => None,
@ -564,7 +622,7 @@ mod tests {
let port = port_of(&server);
let config = Config::from_yaml("model_list: []").unwrap();
let clients = connect(&config, &|name| match name {
let clients = connect(&config, &|name: &str| match name {
"REDIS_HOST" => Some("127.0.0.1".to_owned()),
"REDIS_PORT" => Some(port.clone()),
_ => None,

61
litellm-rust/crates/gateway/src/redis/secrets.rs Normal file
View file

@ -0,0 +1,61 @@
use std::collections::{BTreeMap, BTreeSet};
use litellm_core_utils::settings::Lookup;
use litellm_secret_value::SecretValue;
use litellm_config::Value;
use litellm_redis_resolve::SETTINGS;
use super::configured::ConfiguredRedis;
#[derive(Clone, Debug, Default)]
pub(crate) struct SecretSnapshot(BTreeMap<String, Option<SecretValue>>);
impl SecretSnapshot {
pub(crate) fn read(names: BTreeSet<String>, lookup: &dyn Lookup) -> Self {
Self(
names
.into_iter()
.map(|name| {
let value = lookup.get(&name).map(SecretValue::new);
(name, value)
})
.collect(),
)
}
}
impl Lookup for SecretSnapshot {
fn get(&self, name: &str) -> Option<String> {
match self.0.get(name) {
Some(value) => value.as_ref().map(|value| value.expose().to_owned()),
None => {
debug_assert!(
false,
"{name} is missing from ConfiguredRedis::secret_names"
);
tracing::error!(
name,
"a Redis secret was read without being in the snapshot"
);
None
}
}
}
}
impl ConfiguredRedis {
pub(crate) fn secret_names(&self) -> BTreeSet<String> {
let references = self.fields.values().filter_map(|value| match value {
Value::String(text) => text.strip_prefix("os.environ/").map(str::to_owned),
_ => None,
});
let declared = SETTINGS.iter().flat_map(|setting| {
setting
.environment_name()
.into_iter()
.chain(setting.fallback.iter().map(|name| (*name).to_owned()))
});
declared.chain(references).collect()
}
}

964
litellm-rust/crates/gateway/src/redis/surface_tests.rs Normal file
View file

@ -0,0 +1,964 @@
use std::collections::{BTreeMap, HashMap};
use litellm_config::{CacheSettings, Config, Value};
use litellm_redis_resolve::Expected;
use litellm_redis_types::{
Auth, AwsIamAuth, AzureAdAuth, Node, ProviderHandle, Resolution, StaticAuth, Target, TlsOptions,
};
use litellm_secret_value::SecretValue;
use rstest::rstest;
use serde_json::json;
use super::{configured::ConfiguredRedis, secrets::SecretSnapshot};
use crate::Error;
fn resolution(
configured: &ConfiguredRedis,
environment: &HashMap<String, String>,
) -> Result<Resolution, Error> {
let lookup = |name: &str| environment.get(name).cloned();
configured.resolve(
&SecretSnapshot::read(configured.secret_names(), &lookup),
&lookup,
)
}
const PROVIDER: ProviderHandle = ProviderHandle::new(7);
fn redis_cache_kwargs(kwargs: serde_json::Value) -> ConfiguredRedis {
let has_provider = kwargs
.get("credential_provider")
.is_some_and(|provider| !provider.is_null());
let configured = ConfiguredRedis::redis_cache_kwargs(
serde_json::from_value::<BTreeMap<String, Value>>(kwargs).unwrap(),
);
match has_provider {
true => configured.with_credential_provider(PROVIDER),
false => configured,
}
}
fn coordination_resolution(yaml: &str, environment: &[(&str, &str)]) -> Result<Resolution, Error> {
let environment: HashMap<String, String> = environment
.iter()
.map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
.collect();
let config = Config::from_yaml(yaml).unwrap();
let configured = config
.general_settings
.coordination_redis
.unwrap_or_default();
let configured = ConfiguredRedis::coordination(&configured)?;
resolution(&configured, &environment)
}
fn resolve(yaml: &str, environment: &[(&str, &str)]) -> Result<Resolution, Error> {
coordination_resolution(yaml, environment)
}
fn resolve_cache(yaml: &str) -> Result<Resolution, Error> {
resolve_cache_in(yaml, &[])
}
fn resolve_cache_in(yaml: &str, environment: &[(&str, &str)]) -> Result<Resolution, Error> {
let lookup = |name: &str| {
environment
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| (*value).to_owned())
};
let config = Config::from_yaml(yaml).unwrap();
let configured = match config.litellm_settings.cache_params.unwrap() {
CacheSettings::Redis(settings) => ConfiguredRedis::cache_params(&settings),
CacheSettings::Other(_) => panic!("expected a Redis cache"),
};
configured.resolve(
&SecretSnapshot::read(configured.secret_names(), &lookup),
&lookup,
)
}
fn cluster_ports(resolved: &Resolution) -> Vec<u16> {
match &resolved.target {
Target::Cluster(cluster) => cluster.nodes.iter().map(|node| node.port).collect(),
target => panic!("expected a cluster, got {target:?}"),
}
}
#[rstest]
#[case::host(
"general_settings:\n coordination_redis:\n host: redis\n port: 6380\n db: 2\n ssl: true\n namespace: app\n socket_timeout: 1.5",
&[],
"host:redis:6380:2:true:app:1500"
)]
#[case::environment_references_resolve_the_host_and_port(
"general_settings:\n coordination_redis:\n host: os.environ/COORD_HOST\n port: os.environ/COORD_PORT",
&[("COORD_HOST", "referenced"), ("COORD_PORT", "6400")],
"host:referenced:6400:0:false::none"
)]
#[case::a_null_host_falls_back_to_the_environment(
"general_settings:\n coordination_redis:\n host: null",
&[("REDIS_HOST", "environment")],
"host:environment:6379:0:false::none"
)]
fn resolves_redis_through_the_parsed_config(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: &str,
) {
let resolved = resolve(yaml, environment).unwrap();
let description = match resolved.target {
Target::Unix(socket) => format!("unix:{}", socket.path.display()),
Target::Url(url) => format!("url:{}", url.url.expose()),
Target::Sentinel(sentinel) => format!("sentinel:{}", sentinel.service_name),
Target::Standalone(host) => format!(
"host:{}:{}:{}:{}:{}:{}",
host.host,
host.port,
host.db,
host.tls.is_some(),
resolved.namespace.unwrap_or_default(),
resolved
.timeout
.map(|timeout| timeout.as_millis().to_string())
.unwrap_or_else(|| "none".to_owned())
),
Target::Cluster(cluster) => format!(
"cluster:{}",
cluster
.nodes
.iter()
.map(|node| format!("{}:{}", node.host, node.port))
.collect::<Vec<_>>()
.join(",")
),
};
assert_eq!(description, expected);
}
#[rstest]
#[case::tls_ciphers(
"general_settings:\n coordination_redis:\n host: redis\n ssl_ciphers: HIGH",
&[],
"ssl_ciphers"
)]
#[case::ssl_ciphers_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_SSL_CIPHERS", "HIGH")],
"REDIS_SSL_CIPHERS"
)]
#[case::connection_pool(
"general_settings:\n coordination_redis:\n host: redis\n connection_pool: shared",
&[],
"connection_pool"
)]
fn unsupported_settings_resolve_and_are_named_by_their_source(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: &str,
) {
let resolution = coordination_resolution(yaml, environment).unwrap();
assert_eq!(
resolution.target,
Target::Standalone(litellm_redis_types::Standalone {
host: "redis".into(),
port: 6379,
db: 0,
tls: None,
})
);
let [ignored] = resolution.ignored.as_slice() else {
panic!("expected one ignored setting, got {:?}", resolution.ignored);
};
assert_eq!(ignored.setting, expected);
assert!(!ignored.consequence.is_empty());
}
#[rstest]
#[case::kwarg(
"general_settings:\n coordination_redis:\n host: redis\n ssl: true\n ssl_validate_ocsp: true",
&[],
"ssl_validate_ocsp"
)]
#[case::stapled_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_SSL", "true"), ("REDIS_SSL_VALIDATE_OCSP_STAPLED", "yes")],
"REDIS_SSL_VALIDATE_OCSP_STAPLED"
)]
#[case::rediss_url_query(
"general_settings:\n coordination_redis:\n url: rediss://cache:6380/0?ssl_validate_ocsp=true",
&[],
"url ?ssl_validate_ocsp"
)]
fn an_ocsp_check_fails_instead_of_connecting_without_it(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] setting: &str,
) {
let error = coordination_resolution(yaml, environment).unwrap_err();
assert!(
matches!(&error, Error::Redis(litellm_redis_resolve::Error::Unenforceable { setting: named, .. }) if named == setting),
"{error:?}"
);
}
#[rstest]
#[case::kwarg(" host: redis\n ssl: true\n ssl_validate_ocsp: false")]
#[case::rediss_url_query(" url: rediss://cache:6380/0?ssl_validate_ocsp=false")]
fn an_ocsp_flag_set_to_false_asks_for_nothing(#[case] block: &str) {
let yaml = format!("general_settings:\n coordination_redis:\n{block}");
let resolution = coordination_resolution(&yaml, &[]).unwrap();
assert!(resolution.ignored.is_empty(), "{:?}", resolution.ignored);
}
fn static_auth(username: Option<&str>, password: Option<&str>) -> Auth {
Auth::Static(StaticAuth {
username: username.map(litellm_secret_value::SecretValue::new),
password: password.map(litellm_secret_value::SecretValue::new),
})
}
#[rstest]
fn an_explicit_host_keeps_environment_credentials() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: configured",
&[
("REDIS_PASSWORD", "environment-secret"),
("REDIS_SSL", "True"),
],
)
.unwrap();
assert_eq!(resolved.auth, static_auth(None, Some("environment-secret")));
assert!(matches!(resolved.target, Target::Standalone(host) if host.tls.is_some()));
}
fn azure(username: Option<&str>) -> Auth {
Auth::AzureAd(AzureAdAuth {
client_id: Some("client".to_owned()),
tenant_id: None,
client_secret: None,
username: username.map(str::to_owned),
})
}
fn aws(serverless: bool) -> Auth {
Auth::AwsIam(AwsIamAuth {
user_name: "iam-user".to_owned(),
cache_name: "cache-1".to_owned(),
region: "us-east-1".to_owned(),
serverless,
})
}
#[rstest]
#[case::static_credentials(
"general_settings:\n coordination_redis:\n host: redis\n username: user\n password: pass",
&[],
static_auth(Some("user"), Some("pass"))
)]
#[case::gcp_beats_azure_and_aws(
"general_settings:\n coordination_redis:\n host: redis\n password: pass\n azure_redis_ad_token: \"True\"\n aws_iam_auth: true",
&[("REDIS_GCP_SERVICE_ACCOUNT", "projects/-/serviceAccounts/sa@p.iam.gserviceaccount.com")],
Auth::GcpIam { service_account: "projects/-/serviceAccounts/sa@p.iam.gserviceaccount.com".to_owned() }
)]
#[case::azure_beats_aws_and_takes_the_environment_username(
"general_settings:\n coordination_redis:\n host: redis\n username: configured-user\n azure_redis_ad_token: true\n azure_client_id: client\n aws_iam_auth: true",
&[("REDIS_USERNAME", "environment-user")],
azure(Some("environment-user"))
)]
#[case::azure_without_an_environment_username_sends_only_the_token(
"general_settings:\n coordination_redis:\n host: redis\n username: configured-user\n azure_redis_ad_token: true\n azure_client_id: client",
&[],
azure(None)
)]
#[case::azure_needs_the_literal_true(
"general_settings:\n coordination_redis:\n host: redis\n azure_redis_ad_token: \"1\"",
&[],
static_auth(None, None)
)]
#[case::aws_with_a_region_from_the_environment(
"general_settings:\n coordination_redis:\n host: redis\n ssl: true\n password: dropped\n aws_iam_auth: \"yes\"\n aws_iam_user_name: iam-user\n aws_iam_cache_name: cache-1",
&[("AWS_DEFAULT_REGION", "us-east-1")],
aws(false)
)]
#[case::aws_serverless_over_a_rediss_url(
"general_settings:\n coordination_redis:\n url: rediss://cache-1:6379\n aws_iam_auth: true\n aws_iam_user_name: iam-user\n aws_iam_cache_name: cache-1\n aws_iam_region: us-east-1\n aws_iam_serverless: true",
&[],
aws(true)
)]
fn auth_follows_python_precedence(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: Auth,
) {
assert_eq!(resolve(yaml, environment).unwrap().auth, expected);
}
#[rstest]
#[case::coordination(
"general_settings:\n coordination_redis:\n host: redis\n credential_provider: host-object",
&[]
)]
#[case::without_a_target(
"general_settings:\n coordination_redis:\n credential_provider: host-object",
&[("REDIS_HOST", "redis")]
)]
fn a_credential_provider_without_a_host_handle_fails(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
) {
let error = resolve(yaml, environment).unwrap_err();
assert!(
matches!(
error,
Error::Redis(litellm_redis_resolve::Error::ProviderWithoutHandle)
),
"{error}"
);
}
#[rstest]
#[case::without_tls(
"general_settings:\n coordination_redis:\n host: redis\n aws_iam_auth: true\n aws_iam_user_name: u\n aws_iam_cache_name: c\n aws_iam_region: r",
"requires TLS"
)]
#[case::a_plain_url_ignores_ssl(
"general_settings:\n coordination_redis:\n url: redis://cache:6379\n ssl: true\n aws_iam_auth: true\n aws_iam_user_name: u\n aws_iam_cache_name: c\n aws_iam_region: r",
"requires TLS"
)]
#[case::missing_settings(
"general_settings:\n coordination_redis:\n host: redis\n ssl: true\n aws_iam_auth: true\n aws_iam_cache_name: c",
"requires: aws_iam_user_name, aws_iam_region"
)]
fn incomplete_aws_iam_auth_fails(#[case] yaml: &str, #[case] message: &str) {
let error = resolve(yaml, &[]).unwrap_err();
assert!(error.to_string().contains(message), "{error}");
}
#[rstest]
#[case::defaults_verify_the_certificate_but_not_the_hostname(
" ssl: true",
&[],
Some(TlsOptions { ca_certs: None, ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: None, verify_certificate: true, check_hostname: false })
)]
#[case::every_option(
" ssl: true\n ssl_ca_certs: /ca.pem\n ssl_certfile: /client.pem\n ssl_keyfile: /client.key\n ssl_password: key-secret\n ssl_cert_reqs: none\n ssl_check_hostname: true",
&[],
Some(TlsOptions { ca_certs: Some("/ca.pem".into()), ca_data: None, ca_path: None, certfile: Some("/client.pem".into()), keyfile: Some("/client.key".into()), key_password: Some(SecretValue::new("key-secret")), verify_certificate: false, check_hostname: true })
)]
#[case::options_from_the_environment(
" ssl: true",
&[("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional"), ("REDIS_SSL_PASSWORD", "env-key-secret")],
Some(TlsOptions { ca_certs: Some("/env-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: Some(SecretValue::new("env-key-secret")), verify_certificate: true, check_hostname: false })
)]
#[case::gcp_ca_certs_replace_the_configured_ones(
" ssl: true\n ssl_ca_certs: /ca.pem\n gcp_service_account: account\n gcp_ssl_ca_certs: /gcp-ca.pem",
&[],
Some(TlsOptions { ca_certs: Some("/gcp-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: None, verify_certificate: true, check_hostname: false })
)]
#[case::options_without_ssl_stay_plaintext(" ssl_ca_certs: /ca.pem", &[], None)]
fn tls_options_resolve_like_redis_py(
#[case] block: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: Option<TlsOptions>,
) {
let yaml = format!("general_settings:\n coordination_redis:\n host: redis\n{block}");
let resolved = resolve(&yaml, environment).unwrap();
assert!(matches!(resolved.target, Target::Standalone(host) if host.tls == expected));
}
#[rstest]
fn an_unknown_cert_requirement_fails() {
let error = resolve(
"general_settings:\n coordination_redis:\n host: redis\n ssl: true\n ssl_cert_reqs: strict",
&[],
)
.unwrap_err();
assert!(matches!(
error,
Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::CertRequirement }) if setting == "ssl_cert_reqs"
));
}
#[rstest]
#[case::from_config(
"general_settings:\n coordination_redis:\n sentinel_nodes: [[sentinel-a, 26379], [sentinel-b, \"26380\"]]\n service_name: primary\n sentinel_password: sentinel-secret\n db: 2",
&[]
)]
#[case::from_the_environment(
"model_list: []",
&[
("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379], ["sentinel-b", "26380"]]"#),
("REDIS_SERVICE_NAME", "primary"),
("REDIS_SENTINEL_PASSWORD", "sentinel-secret"),
("REDIS_DB", "2"),
]
)]
fn sentinel_nodes_resolve_to_a_sentinel_endpoint(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
) {
let Target::Sentinel(sentinel) = resolve(yaml, environment).unwrap().target else {
panic!("expected a sentinel endpoint");
};
assert_eq!(
sentinel.monitors,
vec![
Node {
host: "sentinel-a".to_owned(),
port: 26379
},
Node {
host: "sentinel-b".to_owned(),
port: 26380
},
]
);
assert_eq!(sentinel.service_name, "primary");
assert_eq!(
sentinel.monitor_auth,
Some(StaticAuth {
username: None,
password: Some(litellm_secret_value::SecretValue::new("sentinel-secret")),
})
);
assert_eq!(sentinel.db, 2);
}
#[rstest]
#[case::a_url_beats_sentinels(
"general_settings:\n coordination_redis:\n url: redis://single:6379\n sentinel_nodes: [[s, 26379]]\n service_name: primary",
"url:redis://single:6379"
)]
fn sentinel_precedence(#[case] yaml: &str, #[case] expected: &str) {
let resolved = resolve(yaml, &[]).unwrap();
let description = match resolved.target {
Target::Url(url) => format!("url:{}", url.url.expose()),
Target::Sentinel(sentinel) => format!("sentinel:{}", sentinel.service_name),
other => panic!("unexpected {other:?}"),
};
assert_eq!(description, expected);
}
#[rstest]
fn sentinels_without_a_service_name_fail() {
let error = resolve(
"general_settings:\n coordination_redis:\n sentinel_nodes: [[s, 26379]]",
&[],
)
.unwrap_err();
assert!(
matches!(
error,
Error::Redis(litellm_redis_resolve::Error::NoSentinelService)
),
"{error:?}"
);
}
#[rstest]
#[case::configured_host(
"general_settings:\n coordination_redis:\n host: ignored\n sentinel_nodes: [[s, 26379]]\n service_name: primary",
&[],
"host"
)]
#[case::environment_port(
"general_settings:\n coordination_redis:\n sentinel_nodes: [[s, 26379]]\n service_name: primary",
&[("REDIS_PORT", "6380")],
"REDIS_PORT"
)]
fn a_host_or_port_cannot_reach_sentinel(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] source: &str,
) {
let error = resolve(yaml, environment).unwrap_err();
assert!(
matches!(&error, Error::Redis(litellm_redis_resolve::Error::SentinelWithHost { setting }) if setting == source),
"{error:?}"
);
}
#[rstest]
#[case::not_pairs(
"general_settings:\n coordination_redis:\n sentinel_nodes: [{host: s, port: 26379}]\n service_name: primary",
&[],
"sentinel_nodes"
)]
#[case::from_the_environment(
"model_list: []",
&[("REDIS_SENTINEL_NODES", "not-json"), ("REDIS_SERVICE_NAME", "primary")],
"REDIS_SENTINEL_NODES"
)]
fn invalid_sentinel_nodes_name_their_source(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] source: &str,
) {
let error = resolve(yaml, environment).unwrap_err();
assert!(
matches!(&error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::SentinelNodes }) if setting == source),
"{error:?}"
);
}
#[rstest]
#[case::from_config("general_settings:\n coordination_redis:\n startup_nodes: [{host: a}]", &[], "startup_nodes")]
#[case::from_the_environment("model_list: []", &[("REDIS_CLUSTER_NODES", "not-json")], "REDIS_CLUSTER_NODES")]
#[case::out_of_range_port_from_the_environment("model_list: []", &[("REDIS_HOST", "redis"), ("REDIS_PORT", "70000")], "REDIS_PORT")]
fn invalid_values_name_their_source(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] source: &str,
) {
let error = resolve(yaml, environment).unwrap_err();
assert!(
matches!(&error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, .. }) if setting == source),
"{error:?}"
);
}
#[rstest]
#[case::falls_back_to_the_socket_timeout(" socket_timeout: 1.5", Some(1500), None)]
#[case::explicit_connect_timeout(
" socket_timeout: 1.5\n socket_connect_timeout: 0.25",
Some(1500),
Some(250)
)]
fn connect_timeout_is_resolved_separately(
#[case] block: &str,
#[case] timeout_ms: Option<u128>,
#[case] connect_ms: Option<u128>,
) {
let resolved = resolve(
&format!("general_settings:\n coordination_redis:\n host: redis\n{block}"),
&[],
)
.unwrap();
assert_eq!(
resolved.timeout.map(|timeout| timeout.as_millis()),
timeout_ms
);
assert_eq!(
resolved.connect_timeout.map(|timeout| timeout.as_millis()),
connect_ms
);
}
#[rstest]
#[case::only_a_port("general_settings:\n coordination_redis:\n port: 6379", &[])]
#[case::a_structured_host_with_an_environment_target(
"general_settings:\n coordination_redis:\n host: []",
&[("REDIS_HOST", "environment"), ("REDIS_URL", "redis://environment:6379")]
)]
fn a_block_without_a_usable_target_fails(#[case] yaml: &str, #[case] environment: &[(&str, &str)]) {
let error = resolve(yaml, environment).unwrap_err();
assert!(
matches!(error, Error::Redis(litellm_redis_resolve::Error::NoTarget)),
"{error:?}"
);
}
#[rstest]
fn the_cache_alias_names_cluster_nodes() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n redis_startup_nodes:\n - host: node-a\n port: 7000",
)
.unwrap();
assert_eq!(cluster_ports(&resolved), vec![7000]);
}
#[rstest]
fn canonical_startup_nodes_beat_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n startup_nodes:\n - host: canonical\n port: 7000\n redis_startup_nodes:\n - host: alias\n port: 7100\n - host: alias\n port: 7101",
)
.unwrap();
assert_eq!(cluster_ports(&resolved), vec![7000]);
}
#[rstest]
fn a_null_canonical_startup_nodes_setting_shadows_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: null\n redis_startup_nodes:\n - host: alias\n port: 7100",
)
.unwrap();
assert!(matches!(resolved.target, Target::Standalone(_)));
}
#[rstest]
fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: os.environ/MISSING_NODES\n redis_startup_nodes:\n - host: alias\n port: 7100",
)
.unwrap();
assert!(matches!(resolved.target, Target::Standalone(_)));
}
#[rstest]
fn unsupported_environment_settings_never_carry_the_value() {
let resolution = coordination_resolution(
"model_list: []",
&[
("REDIS_HOST", "redis"),
("REDIS_SSL_CIPHERS", "private-option-value"),
],
)
.unwrap();
let ignored = format!("{:?}", resolution.ignored);
assert!(ignored.contains("REDIS_SSL_CIPHERS"), "{ignored}");
assert!(!ignored.contains("private-option-value"), "{ignored}");
}
#[rstest]
#[case::negative("-1")]
#[case::infinite("inf")]
fn rejects_bad_socket_timeout(#[case] timeout: &str) {
let yaml = format!(
"general_settings:\n coordination_redis:\n host: redis\n socket_timeout: {timeout}"
);
let error = resolve(&yaml, &[]).unwrap_err();
assert!(matches!(
error,
Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::Seconds }) if setting == "socket_timeout"
));
}
#[rstest]
fn debug_redacts_redis_secrets_and_url() {
let config = Config::from_yaml(
"general_settings:\n coordination_redis:\n url: redis://user:password@redis:6379\n username: credential-user\n password: credential-password\n sentinel_kwargs: {password: credential-passthrough}",
)
.unwrap();
let settings = config.general_settings.coordination_redis.unwrap();
let configured = ConfiguredRedis::coordination(&settings).unwrap();
let resolved = resolution(&configured, &HashMap::new()).unwrap();
let debug = format!("{settings:?} {configured:?} {resolved:?}");
assert!(!debug.contains("redis://user:password"));
assert!(!debug.contains("credential-user"));
assert!(!debug.contains("credential-password"));
assert!(!debug.contains("credential-passthrough"));
assert!(debug.contains("redacted") || debug.contains("REDACTED"));
}
#[rstest]
fn legacy_unknown_settings_do_not_change_resolution() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: redis\n redis_flush_size: 42\n service_name: primary",
&[],
)
.unwrap();
assert!(matches!(resolved.target, Target::Standalone(_)));
}
#[rstest]
#[case::configured(
"general_settings:\n coordination_redis:\n host: redis\n max_connections: 47",
&[],
"max_connections"
)]
#[case::from_the_environment(
"general_settings:\n coordination_redis:\n host: redis",
&[("REDIS_MAX_CONNECTIONS", "47")],
"REDIS_MAX_CONNECTIONS"
)]
fn ignored_settings_are_named_by_their_source(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] source: &str,
) {
let resolution = coordination_resolution(yaml, environment).unwrap();
let settings: Vec<&str> = resolution
.ignored
.iter()
.map(|ignored| ignored.setting.as_str())
.collect();
assert_eq!(settings, vec![source]);
}
fn snapshot(values: &[(&str, &str)]) -> SecretSnapshot {
let values: HashMap<&str, &str> = values.iter().copied().collect();
SecretSnapshot::read(
ConfiguredRedis::environment_only().secret_names(),
&|name: &str| values.get(name).map(|value| (*value).to_owned()),
)
}
#[rstest]
fn secret_manager_values_reach_the_resolver_through_the_snapshot() {
let settings = redis_cache_kwargs(json!({"host": "os.environ/COORD_HOST"}));
let secrets = SecretSnapshot::read(settings.secret_names(), &|name: &str| match name {
"COORD_HOST" => Some("from-the-manager".to_owned()),
"REDIS_PASSWORD" => Some("manager-secret".to_owned()),
_ => None,
});
let resolved = settings.resolve(&secrets, &|_: &str| None).unwrap();
assert!(
matches!(&resolved.target, Target::Standalone(host) if host.host == "from-the-manager")
);
assert_eq!(resolved.auth, static_auth(None, Some("manager-secret")));
}
#[rstest]
#[case::from_the_environment(&[], &[("REDIS_USERNAME", "environment-user")], Some("environment-user"))]
#[case::never_from_the_secret_manager(&[("REDIS_USERNAME", "manager-user")], &[], None)]
fn the_azure_username_comes_from_the_process_environment_like_python(
#[case] secrets: &[(&str, &str)],
#[case] environment: &[(&str, &str)],
#[case] expected: Option<&str>,
) {
let settings = redis_cache_kwargs(json!({"host": "redis", "azure_redis_ad_token": "true"}));
let environment: HashMap<&str, &str> = environment.iter().copied().collect();
let resolved = settings
.resolve(&snapshot(secrets), &|name: &str| {
environment.get(name).map(|value| (*value).to_owned())
})
.unwrap();
assert!(
matches!(&resolved.auth, Auth::AzureAd(azure) if azure.username.as_deref() == expected),
"{:?}",
resolved.auth
);
}
#[rstest]
#[case::coordination_drops_the_null(
"general_settings:\n coordination_redis:\n host: redis\n ssl: null",
false,
true
)]
#[case::cache_passes_an_unnamed_null_on(
"litellm_settings:\n cache_params:\n type: redis\n host: redis\n ssl: null",
true,
false
)]
#[case::cache_drops_a_null_it_names(
"litellm_settings:\n cache_params:\n type: redis\n host: null\n ssl: true",
true,
true
)]
fn a_null_overrides_the_environment_only_where_python_passes_it_on(
#[case] yaml: &str,
#[case] cache: bool,
#[case] tls: bool,
) {
let environment = [("REDIS_HOST", "environment"), ("REDIS_SSL", "true")];
let resolved = match cache {
true => resolve_cache_in(yaml, &environment),
false => resolve(yaml, &environment),
}
.unwrap();
let Target::Standalone(host) = resolved.target else {
panic!("expected a standalone target, got {:?}", resolved.target);
};
assert_eq!(host.tls.is_some(), tls);
}
#[rstest]
#[case::empty_cluster_nodes_use_the_environment(
json!({"startup_nodes": []}),
&[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": 7000}]"#)],
json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": 7000}]"}),
&[("startup_nodes", "REDIS_CLUSTER_NODES")],
)]
#[case::empty_sentinel_nodes_use_the_environment(
json!({"sentinel_nodes": [], "service_name": "primary"}),
&[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#)],
json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "service_name": "primary"}),
&[("sentinel_nodes", "REDIS_SENTINEL_NODES")],
)]
#[case::empty_service_name_uses_the_environment(
json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": ""}),
&[("REDIS_SERVICE_NAME", "primary")],
json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}),
&[("service_name", "REDIS_SERVICE_NAME")],
)]
#[case::environment_host_port_and_password(
json!({}),
&[("REDIS_HOST", "environment"), ("REDIS_PASSWORD", "environment-secret"), ("REDIS_PORT", "6390"), ("REDIS_SSL", "True")],
json!({"port": "6390", "ssl": true, "host": "environment", "password": "environment-secret"}),
&[("host", "REDIS_HOST"), ("password", "REDIS_PASSWORD"), ("port", "REDIS_PORT"), ("ssl", "REDIS_SSL")],
)]
#[case::padded_environment_numbers_parse_like_python_int(
json!({}),
&[("REDIS_HOST", "redis"), ("REDIS_PORT", "6_380"), ("REDIS_DB", "3\n")],
json!({"port": "6_380", "host": "redis", "db": "3\n"}),
&[("db", "REDIS_DB"), ("host", "REDIS_HOST"), ("port", "REDIS_PORT")],
)]
#[case::socket_timeout_from_the_environment(
json!({"host": "redis"}),
&[("REDIS_SOCKET_TIMEOUT", "0.75")],
json!({"socket_timeout": "0.75", "host": "redis"}),
&[("socket_timeout", "REDIS_SOCKET_TIMEOUT")],
)]
#[case::environment_url(
json!({}),
&[("REDIS_URL", "redis://environment:6379/1")],
json!({"url": "redis://environment:6379/1"}),
&[("url", "REDIS_URL")],
)]
#[case::environment_url_beats_the_environment_host(
json!({}),
&[("REDIS_HOST", "ignored"), ("REDIS_URL", "redis://from-url:6379")],
json!({"host": "ignored", "url": "redis://from-url:6379"}),
&[("host", "REDIS_HOST"), ("url", "REDIS_URL")],
)]
#[case::an_explicit_host_drops_the_environment_url(
json!({"host": "configured"}),
&[("REDIS_URL", "redis://environment:6379")],
json!({"host": "configured"}),
&[],
)]
#[case::environment_cluster_nodes(
json!({}),
&[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": "7000"}, {"host": "node-b", "port": 7001}]"#)],
json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": \"7000\"}, {\"host\": \"node-b\", \"port\": 7001}]"}),
&[("startup_nodes", "REDIS_CLUSTER_NODES")],
)]
#[case::sentinel_from_the_environment(
json!({}),
&[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#), ("REDIS_SENTINEL_PASSWORD", "sentinel-secret"), ("REDIS_SERVICE_NAME", "primary")],
json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "sentinel_password": "sentinel-secret", "service_name": "primary"}),
&[("sentinel_nodes", "REDIS_SENTINEL_NODES"), ("sentinel_password", "REDIS_SENTINEL_PASSWORD"), ("service_name", "REDIS_SERVICE_NAME")],
)]
#[case::an_explicit_credential_provider_beats_cloud_auth(
json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}),
&[("REDIS_AZURE_AD_TOKEN", "true")],
json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}),
&[],
)]
#[case::gcp_beats_azure_and_aws(
json!({"aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}),
&[("REDIS_GCP_SERVICE_ACCOUNT", "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com")],
json!({"gcp_service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com", "aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}),
&[("gcp_service_account", "REDIS_GCP_SERVICE_ACCOUNT")],
)]
#[case::azure_beats_aws_and_takes_the_environment_username(
json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}),
&[("REDIS_USERNAME", "environment-user")],
json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}),
&[],
)]
#[case::an_empty_environment_username_sends_only_the_token(
json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis"}),
&[("REDIS_USERNAME", "")],
json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": ""}),
&[("username", "REDIS_USERNAME")],
)]
#[case::aws_iam_over_tls_with_the_default_region(
json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true}),
&[("AWS_DEFAULT_REGION", "us-east-1")],
json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true, "aws_iam_region": "us-east-1"}),
&[("aws_iam_region", "AWS_DEFAULT_REGION")],
)]
#[case::aws_region_beats_the_default_region(
json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true}),
&[("AWS_DEFAULT_REGION", "us-east-1"), ("AWS_REGION", "eu-central-1")],
json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true, "aws_iam_region": "eu-central-1"}),
&[("aws_iam_region", "AWS_REGION")],
)]
#[case::a_null_cert_requirement_skips_verification(
json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}),
&[("REDIS_SSL_CERT_REQS", "required")],
json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}),
&[],
)]
#[case::a_cert_requirement_spelled_like_python_none_fails(
json!({"host": "redis"}),
&[("REDIS_SSL", "true"), ("REDIS_SSL_CERT_REQS", "None")],
json!({"ssl": true, "ssl_cert_reqs": "None", "host": "redis"}),
&[("ssl", "REDIS_SSL"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")],
)]
#[case::tls_options_from_the_environment(
json!({"host": "redis"}),
&[("REDIS_SSL", "true"), ("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional")],
json!({"ssl_ca_certs": "/env-ca.pem", "ssl": true, "ssl_cert_reqs": "optional", "host": "redis"}),
&[("ssl", "REDIS_SSL"), ("ssl_ca_certs", "REDIS_SSL_CA_CERTS"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")],
)]
#[case::environment_references_resolve(
json!({"host": "os.environ/COORD_HOST", "password": "os.environ/COORD_PASSWORD"}),
&[("COORD_HOST", "referenced"), ("COORD_PASSWORD", "referenced-secret")],
json!({"host": "referenced", "password": "referenced-secret"}),
&[],
)]
#[case::malformed_port_from_the_environment(
json!({}),
&[("REDIS_HOST", "redis"), ("REDIS_PORT", "abc")],
json!({"port": "abc", "host": "redis"}),
&[("host", "REDIS_HOST"), ("port", "REDIS_PORT")],
)]
#[case::retry_on_error_from_the_environment(
json!({"host": "redis"}),
&[("REDIS_RETRY_ON_ERROR", "true")],
json!({"retry_on_error": true, "host": "redis"}),
&[("retry_on_error", "REDIS_RETRY_ON_ERROR")],
)]
#[case::ssl_ca_path_from_the_environment(
json!({"host": "redis"}),
&[("REDIS_SSL", "true"), ("REDIS_SSL_CA_PATH", "/etc/ssl/private-cas")],
json!({"ssl": true, "ssl_ca_path": "/etc/ssl/private-cas", "host": "redis"}),
&[("ssl", "REDIS_SSL"), ("ssl_ca_path", "REDIS_SSL_CA_PATH")],
)]
fn merges_like_python_redis(
#[case] kwargs: serde_json::Value,
#[case] environment: &[(&str, &str)],
#[case] merged: serde_json::Value,
#[case] sources: &[(&str, &str)],
) {
let configured = redis_cache_kwargs(kwargs);
let lookup = |name: &str| {
environment
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| (*value).to_owned())
};
let kwargs = configured
.merged(&SecretSnapshot::read(configured.secret_names(), &lookup))
.unwrap();
let values: serde_json::Map<String, serde_json::Value> = kwargs
.iter()
.map(|(kwarg, entry)| (kwarg.clone(), entry.value.clone()))
.collect();
let renamed: Vec<(&str, &str)> = kwargs
.iter()
.filter(|(kwarg, entry)| *kwarg != &entry.source)
.map(|(kwarg, entry)| (kwarg.as_str(), entry.source.as_str()))
.collect();
assert_eq!(serde_json::Value::Object(values), merged);
assert_eq!(renamed, sources);
}

22
litellm-rust/crates/gateway/src/settings.rs
View file

@ -4,6 +4,23 @@ use litellm_secret_value::SecretValue;
use litellm_tracing::Level;
use serde::{Deserialize, de::DeserializeOwned};
#[derive(Deserialize)]
pub(super) struct LogSettings {
rust_log: Option<String>,
}
impl LogSettings {
pub fn log_level(&self) -> Level {
log_level(self.rust_log.as_deref())
}
}
fn log_level(rust_log: Option<&str>) -> Level {
rust_log
.and_then(|value| value.parse().ok())
.unwrap_or(Level::INFO)
}
#[derive(Deserialize)]
#[serde(default)]
pub(super) struct Settings {
@ -26,10 +43,7 @@ impl Default for Settings {
impl Settings {
pub fn log_level(&self) -> Level {
self.rust_log
.as_deref()
.and_then(|value| value.parse().ok())
.unwrap_or(Level::INFO)
log_level(self.rust_log.as_deref())
}
}

2
litellm-rust/crates/host-python/src/lib.rs
View file

@ -37,7 +37,7 @@ pub use owned::PythonOwned;
pub use runtime::{
ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, enter_native, poll_async_value,
ready_future, reserve_process_for_forking, run_async, run_async_value, run_sync,
run_sync_value, runtime_started,
run_sync_value, runtime_handle, runtime_started,
};
pub use services::PythonHostCalls;

5
litellm-rust/crates/host-python/src/runtime.rs
View file

@ -68,6 +68,11 @@ fn runtime() -> PyResult<&'static Runtime> {
Ok(pyo3_async_runtimes::tokio::get_runtime())
}
/// The native runtime, for work that must outlive the call that schedules it.
pub fn runtime_handle() -> PyResult<Handle> {
Ok(runtime()?.handle().clone())
}
#[expect(clippy::disallowed_methods, reason = "this is the gated door")]
fn future_into_py<F, T>(py: Python<'_>, future: F) -> PyResult<Bound<'_, PyAny>>
where

11
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -30,6 +30,14 @@ bytes.workspace = true
futures-util.workspace = true
litellm-cache.workspace = true
litellm-error.workspace = true
litellm-circuit-breaker.workspace = true
litellm-redis-client.workspace = true
litellm-redis-connect.workspace = true
litellm-redis-store.workspace = true
litellm-redis-resolve.workspace = true
litellm-python-compat.workspace = true
litellm-redis-types.workspace = true
tracing.workspace = true
litellm-cache-response.workspace = true
serde.workspace = true
litellm-auth.workspace = true
@ -52,7 +60,8 @@ litellm-router-core.workspace = true
litellm-host-python.workspace = true
litellm-model-catalog.workspace = true
litellm-token-counter = { path = "../token-counter", default-features = false }
pyo3.workspace = true
indexmap = "2.14.0"
pyo3 = { workspace = true, features = ["indexmap"] }
prost.workspace = true
pyo3-async-runtimes.workspace = true
reqwest.workspace = true

46
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -12,6 +12,7 @@ mod marshal;
mod model_registry;
mod preflight;
mod python_settings;
mod redis;
mod router;
mod routes;
mod secrets;
@ -37,6 +38,29 @@ mod _native {
#[pymodule_export]
use crate::model_registry::bundled_model_cost_map;
#[pymodule_export]
use crate::redis::breaker::{NativeAdmission, NativeRedisBreaker, native_redis_breaker_states};
#[pymodule_export]
use crate::redis::client::NativeRedisClient;
#[pymodule_export]
use crate::redis::store::auth_cache::{NativeAuthCache, PyLocalSizes};
#[pymodule_export]
use crate::redis::store::rate_limits::{
NativeRateLimits, PyAtomicApplied, PyAtomicOverLimit, PyGaugeOverLimit, PyGaugesAcquired,
PyGuardedReply, PyWindowCount,
};
#[pymodule_export]
use crate::redis::store::request::{NativeRequestRedis, NativeSend};
#[pymodule_export]
use crate::redis::store::router_usage::NativeRouterUsage;
#[pymodule_export]
use crate::redis::store::sensitive_route_pins::NativeSensitiveRoutePins;
#[pymodule_export]
use crate::redis::store::spend_counters::{NativeSpendCounters, PySpendRead};
#[pymodule_export]
use crate::redis::store::{
NativePending, NativeRedisStoreError, router_cooldowns::NativeRouterCooldowns,
};
#[pymodule_export]
use crate::router::{RouterCall, RouterCoreHandle};
#[pymodule_export]
use crate::routes::audio_transcription::{atranscription, transcription};
@ -111,6 +135,28 @@ mod tests {
"aresponses",
"ResponsesWebSocketConnection",
"NativeDiagnosticProcessor",
"NativeRedisClient",
"NativeRedisBreaker",
"NativeAdmission",
"NativeRequestRedis",
"NativeSend",
"NativeAuthCache",
"NativeRateLimits",
"WindowCount",
"AtomicApplied",
"AtomicOverLimit",
"GuardedReply",
"GaugesAcquired",
"GaugeOverLimit",
"NativePending",
"NativeRouterCooldowns",
"NativeRouterUsage",
"NativeSensitiveRoutePins",
"NativeRedisStoreError",
"NativeSpendCounters",
"SpendRead",
"LocalSizes",
"native_redis_breaker_states",
"NativeClickHouseSpendConfig",
"NativeClickHouseSpendStorage",
"NativeDatabase",

2
litellm-rust/crates/python-bridge/src/python_settings.rs
View file

@ -13,6 +13,8 @@ pub(crate) enum PythonSettings {
UrlPolicy,
#[strum(serialize = "provider_defaults")]
ProviderDefaults,
#[strum(serialize = "router_tiers")]
RouterTiers,
#[strum(serialize = "secret_manager")]
SecretManager,
#[strum(serialize = "secret_manager_binding")]

134
litellm-rust/crates/python-bridge/src/redis/breaker.rs Normal file
View file

@ -0,0 +1,134 @@
use std::{
sync::{Arc, Mutex, PoisonError},
time::{Duration, Instant},
};
use litellm_circuit_breaker::{Admission, Breaker, Config, Failure};
use litellm_redis_client::RedisEvents;
use litellm_redis_types::Endpoint;
use pyo3::{exceptions::PyValueError, prelude::*};
use crate::redis::{
client::{RedisSpec, resolved_endpoint},
events::{ServiceLoggingEvents, state_label},
hosts,
};
/// `litellm.redis.store_settings.RedisBreakerSettings`, filled from `litellm.constants`.
#[derive(FromPyObject)]
pub(crate) struct BreakerSettings {
failure_threshold: u32,
recovery_timeout: f64,
timeout_min_duration: f64,
}
impl BreakerSettings {
pub(crate) fn config(&self) -> PyResult<Config> {
let seconds = |seconds: f64| {
Duration::try_from_secs_f64(seconds)
.map_err(|error| PyValueError::new_err(error.to_string()))
};
Ok(Config {
failure_threshold: self.failure_threshold,
recovery_timeout: seconds(self.recovery_timeout)?,
timeout_min_duration: seconds(self.timeout_min_duration)?,
})
}
}
/// The one health signal for a Redis, shared by its `RedisCache` and its native client, so a Redis the
/// cache found wedged is fast-failed by every store too.
#[pyclass(frozen, name = "NativeRedisBreaker")]
pub(crate) struct NativeRedisBreaker(pub(crate) Breaker);
impl NativeRedisBreaker {
fn build(
py: Python<'_>,
settings: &BreakerSettings,
endpoint: Option<Endpoint>,
) -> PyResult<Self> {
let config = settings.config()?;
let runtime = litellm_host_python::runtime_handle()?;
let events: Arc<dyn RedisEvents> =
Arc::new(ServiceLoggingEvents::transitions_only(py, runtime)?);
Ok(Self(py.detach(|| {
hosts().breaker(endpoint, config, Some(events))
})))
}
}
#[pymethods]
impl NativeRedisBreaker {
#[new]
fn new(py: Python<'_>, settings: BreakerSettings) -> PyResult<Self> {
Self::build(py, &settings, None)
}
/// The breaker of the Redis server `redis_kwargs` names, which every cache and client on that server shares
/// whatever its db, credentials or namespace. The first caller's settings win.
#[staticmethod]
fn shared(
py: Python<'_>,
redis_kwargs: RedisSpec,
settings: BreakerSettings,
) -> PyResult<Self> {
let endpoint = resolved_endpoint(py, &redis_kwargs)?;
Self::build(py, &settings, endpoint)
}
fn admit(&self) -> Option<NativeAdmission> {
self.0
.admit(Instant::now())
.ok()
.map(|admission| NativeAdmission(Mutex::new(Some(admission))))
}
fn failed(&self, timeout: bool) {
self.0.failed(failure(timeout), Instant::now());
}
#[getter]
fn state(&self) -> &'static str {
state_label(&self.0.state())
}
}
/// One admitted call. Settling it twice records nothing more, and dropping an unsettled probe reopens the breaker.
#[pyclass(frozen, name = "NativeAdmission")]
pub(crate) struct NativeAdmission(Mutex<Option<Admission>>);
impl NativeAdmission {
fn take(&self) -> Option<Admission> {
self.0.lock().unwrap_or_else(PoisonError::into_inner).take()
}
}
#[pymethods]
impl NativeAdmission {
fn succeeded(&self) {
if let Some(admission) = self.take() {
admission.succeeded();
}
}
fn failed(&self, timeout: bool) {
if let Some(admission) = self.take() {
admission.failed(failure(timeout), Instant::now());
}
}
}
fn failure(timeout: bool) -> Failure {
match timeout {
true => Failure::Timeout,
false => Failure::Hard,
}
}
#[pyfunction]
pub(crate) fn native_redis_breaker_states(py: Python<'_>) -> Vec<&'static str> {
py.detach(|| hosts().breakers())
.iter()
.map(|breaker| state_label(&breaker.state()))
.collect()
}

47
litellm-rust/crates/python-bridge/src/redis/client/credentials.rs Normal file
View file

@ -0,0 +1,47 @@
use std::sync::Arc;
use litellm_host_python::{PythonContext, attach_blocking};
use litellm_redis_connect::{CredentialSource, CredentialStage, Error, Expiring, SourceFuture};
use litellm_secret_value::SecretValue;
use pyo3::prelude::*;
const MODULE: &str = "litellm.rust_bridge.host.redis_credentials";
/// A redis-py `CredentialProvider` the SDK caller passed in `cache_kwargs`, asked through the host
/// module the way redis-py's async connection asks it.
pub(crate) struct PythonCredentialSource {
provider: Arc<Py<PyAny>>,
context: PythonContext,
}
impl PythonCredentialSource {
pub(crate) fn new(provider: Py<PyAny>, context: PythonContext) -> Self {
Self {
provider: Arc::new(provider),
context,
}
}
}
impl CredentialSource for PythonCredentialSource {
fn fetch(&self) -> SourceFuture<'_> {
let provider = Arc::clone(&self.provider);
let context = self.context.clone();
Box::pin(async move {
let (username, password) = attach_blocking(context, move |py| {
py.import(MODULE)?
.getattr("credentials")?
.call1((provider.bind(py),))?
.extract::<(Option<String>, String)>()
})
.await
.and_then(|credentials| credentials)
.map_err(Error::credential(CredentialStage::Provider))?;
Ok(Expiring {
username,
password: SecretValue::new(password),
expires_at: None,
})
})
}
}

337
litellm-rust/crates/python-bridge/src/redis/client/mod.rs Normal file
View file

@ -0,0 +1,337 @@
mod credentials;
use std::{collections::BTreeMap, sync::Arc, time::Duration};
use litellm_circuit_breaker::Config as BreakerConfig;
use litellm_core_utils::settings::ProcessEnvironment;
use litellm_host_python::PythonContext;
use litellm_python_compat::number::{self, SecondsError};
use litellm_redis_client::{RedisClient, RedisEvents};
use litellm_redis_connect::{CredentialSource, HostServices, SharedEnvironment};
use litellm_redis_resolve::{Kwargs, PythonException};
use litellm_redis_types::{Endpoint, ProviderHandle, Resolution};
use pyo3::{
exceptions::{PyOverflowError, PyTypeError, PyValueError},
prelude::*,
types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyMapping, PyString, PyTuple},
};
use self::credentials::PythonCredentialSource;
use crate::redis::{breaker::BreakerSettings, events::ServiceLoggingEvents, hosts};
pyo3::import_exception!(redis.exceptions, RedisClusterException);
/// The kwargs `_merge_redis_kwargs` returns for a `RedisCache`, which the resolver reads the way
/// redis-py's constructors do. A key whose value is a Python object, such as `retry`, reaches the
/// resolver by name only, so it can name the setting it cannot honor. A `credential_provider` object
/// is also kept, so the connector can ask it for credentials, and reaches the resolver as a handle.
pub(crate) struct RedisSpec {
kwargs: Kwargs,
credential_provider: Option<Py<PyAny>>,
}
impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec {
type Error = PyErr;
fn extract(kwargs: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let py = kwargs.py();
let source = PyDict::new(py);
let mapping = kwargs.cast::<PyMapping>()?;
source.update(&mapping)?;
let plain = PyDict::new(py);
let mut host_objects = Vec::new();
for (name, value) in source.iter() {
match value.is_none() || is_plain(&value) {
true => plain.set_item(&name, value)?,
false => host_objects.push(name.extract::<String>()?),
}
}
let credential_provider = source
.get_item("credential_provider")?
.filter(|provider| !provider.is_none() && !is_plain(provider))
.map(Bound::unbind);
let fields: BTreeMap<String, serde_json::Value> =
litellm_host_python::from_py_argument(plain.as_any())?;
let objects = host_objects
.into_iter()
.map(|name| (name, serde_json::Value::Bool(true)));
Ok(Self {
kwargs: litellm_redis_resolve::kwargs(fields.into_iter().chain(objects)),
credential_provider,
})
}
}
/// `litellm.redis.store_settings.RedisStoreSettings`, which a `RedisCache` captures when it is built.
#[derive(FromPyObject)]
pub(crate) struct StoreSettings {
redis_kwargs: RedisSpec,
namespace: Option<String>,
default_ttl: f64,
breaker: Option<BreakerSettings>,
}
fn is_plain(value: &Bound<'_, PyAny>) -> bool {
value.is_instance_of::<PyString>()
|| value.is_instance_of::<PyBool>()
|| value.is_instance_of::<PyInt>()
|| value.is_instance_of::<PyFloat>()
|| value.is_instance_of::<PyList>()
|| value.is_instance_of::<PyTuple>()
|| value.is_instance_of::<PyDict>()
}
/// The one client for a resolved Redis config and credential provider, so every family in this process shares its pool
/// and its breaker. A config redis-py refuses building its client raises the same class here. One it only fails on
/// once a command connects, where `RedisCache` logs every failed call, leaves the family on memory with an error.
fn shared_client(
py: Python<'_>,
spec: RedisSpec,
namespace: Option<String>,
breaker: Option<BreakerConfig>,
) -> PyResult<Option<RedisClient>> {
let environment: SharedEnvironment = Arc::new(ProcessEnvironment);
let logger = crate::logger::capture(py);
let resolved = match logger.scope(|| resolve_spec(py, &spec))? {
Ok(resolution) => {
logger.scope(|| {
for ignored in &resolution.ignored {
litellm_tracing::warn!(
"Redis setting {} is not supported by the Rust client and is ignored: {}",
ignored.setting,
ignored.consequence
);
}
});
Resolution {
namespace,
..resolution
}
}
Err(error) => return unresolved(py, &error),
};
let runtime = litellm_host_python::runtime_handle()?;
let events: Arc<dyn RedisEvents> = Arc::new(ServiceLoggingEvents::new(py, runtime.clone())?);
let host = HostServices {
events: Some(events),
http: crate::http::host_client(py, litellm_http::ClientVariant::Provider)?,
environment,
credential_provider: spec
.credential_provider
.as_ref()
.map(|provider| -> PyResult<Arc<dyn CredentialSource>> {
Ok(Arc::new(PythonCredentialSource::new(
provider.clone_ref(py),
PythonContext::capture(py)?,
)))
})
.transpose()?,
};
let connected = py.detach(|| {
let _runtime = runtime.enter();
hosts().client(resolved, host, breaker)
});
match connected {
Ok(client) => Ok(Some(client)),
Err(error) => Ok(unconnected(py, &error)),
}
}
/// The Rust client for one `RedisCache`, built once and handed to every native store, so a store
/// never sees the redis-py kwargs or the namespace.
#[pyclass(frozen, name = "NativeRedisClient")]
pub(crate) struct NativeRedisClient {
pub(crate) client: RedisClient,
/// `RedisCache.default_ttl`, the TTL a store's write takes when its caller gives none.
pub(crate) default_ttl: Duration,
}
#[pymethods]
impl NativeRedisClient {
/// `None` when redis-py would fail only once a command connects, which leaves every store given it on memory.
#[staticmethod]
fn connect(py: Python<'_>, settings: StoreSettings) -> PyResult<Option<Self>> {
let StoreSettings {
redis_kwargs,
namespace,
default_ttl,
breaker,
} = settings;
let default_ttl = number::seconds(default_ttl).map_err(seconds_error)?;
let breaker = breaker.as_ref().map(BreakerSettings::config).transpose()?;
Ok(
shared_client(py, redis_kwargs, namespace, breaker)?.map(|client| Self {
client,
default_ttl,
}),
)
}
#[getter]
fn default_ttl_seconds(&self) -> f64 {
self.default_ttl.as_secs_f64()
}
}
/// The Redis server `spec` names, whatever its db, credentials or namespace, or none when the Rust client cannot reach it.
pub(crate) fn resolved_endpoint(py: Python<'_>, spec: &RedisSpec) -> PyResult<Option<Endpoint>> {
Ok(resolve_spec(py, spec)?
.ok()
.and_then(|resolution| resolution.target.endpoint()))
}
// A live registry entry's client keeps its provider alive, so no live entry shares its address
fn resolve_spec(
py: Python<'_>,
spec: &RedisSpec,
) -> PyResult<Result<Resolution, litellm_redis_resolve::Error>> {
let handle = spec
.credential_provider
.as_ref()
.map(|provider| ProviderHandle::new(provider.as_ptr() as usize as u64));
resolve(&spec.kwargs, handle, &py.import("os")?.getattr("environ")?)
}
/// What redis-py makes of the kwargs, with `REDIS_USERNAME` read from `environ` the way Python's Azure AD auth
/// reads `os.environ`.
fn resolve(
kwargs: &Kwargs,
provider: Option<ProviderHandle>,
environ: &Bound<'_, PyAny>,
) -> PyResult<Result<Resolution, litellm_redis_resolve::Error>> {
let redis_username: Option<String> = environ
.call_method1("get", ("REDIS_USERNAME",))?
.extract()?;
Ok(litellm_redis_resolve::interpret(
kwargs,
provider,
redis_username.as_deref(),
))
}
fn unresolved(
py: Python<'_>,
error: &litellm_redis_resolve::Error,
) -> PyResult<Option<RedisClient>> {
let message = || litellm_error::Report::new(error).to_string();
match error.raised_building_the_python_client() {
Some(PythonException::ValueError) => Err(PyValueError::new_err(message())),
Some(PythonException::TypeError) => Err(PyTypeError::new_err(message())),
Some(PythonException::RedisClusterException) => {
Err(RedisClusterException::new_err(message()))
}
None => Ok(memory_only(py, error)),
}
}
// redis-py reads TLS files and fetches credentials only once a command connects, so none of these fails building
// its client
fn unconnected(py: Python<'_>, error: &litellm_redis_connect::Error) -> Option<RedisClient> {
match error {
litellm_redis_connect::Error::Connect(_)
| litellm_redis_connect::Error::NoHostCredentialProvider
| litellm_redis_connect::Error::ProcessEnvironmentRequired { .. }
| litellm_redis_connect::Error::Credential { .. }
| litellm_redis_connect::Error::CredentialTimeout { .. }
| litellm_redis_connect::Error::ExpiredCredential
| litellm_redis_connect::Error::MissingCredential { .. } => memory_only(py, error),
}
}
fn memory_only(py: Python<'_>, error: &(dyn std::error::Error + 'static)) -> Option<RedisClient> {
let report = litellm_error::Report::new(error).to_string();
crate::logger::capture(py).scope(|| {
litellm_tracing::error!(
"Redis is not usable from the Rust stores, their state stays in memory: {report}"
)
});
None
}
fn seconds_error(error: SecondsError) -> PyErr {
match error {
SecondsError::Overflow => PyOverflowError::new_err(error.to_string()),
SecondsError::NotANumber => PyValueError::new_err(error.to_string()),
}
}
#[cfg(test)]
mod tests {
use std::ffi::CStr;
use litellm_redis_types::Auth;
use rstest::rstest;
use serde_json::json;
use super::*;
fn kwargs(fields: serde_json::Value) -> Kwargs {
let fields: BTreeMap<String, serde_json::Value> = serde_json::from_value(fields).unwrap();
litellm_redis_resolve::kwargs(fields)
}
#[rstest]
#[case::set(c"environ = {'REDIS_USERNAME': 'azure-user'}", Some("azure-user"))]
#[case::empty_is_unset(c"environ = {'REDIS_USERNAME': ''}", None)]
fn azure_auth_takes_its_username_from_the_environ_like_python(
#[case] source: &CStr,
#[case] expected: Option<&str>,
) {
Python::initialize();
Python::attach(|py| {
let locals = PyDict::new(py);
py.run(source, Some(&locals), Some(&locals)).unwrap();
let environ = locals.get_item("environ").unwrap().unwrap();
let kwargs = kwargs(json!({"host": "redis", "azure_redis_ad_token": "true"}));
let resolved = resolve(&kwargs, None, &environ).unwrap().unwrap();
assert!(
matches!(&resolved.auth, Auth::AzureAd(azure) if azure.username.as_deref() == expected),
"{:?}",
resolved.auth
);
});
}
#[rstest]
#[case::no_target(json!({"port": 6379}), Some("ValueError"))]
#[case::unenforceable(json!({"host": "redis", "ssl": true, "ssl_validate_ocsp": true}), Some("TypeError"))]
#[case::fails_only_once_a_command_connects(
json!({"host": "redis", "ssl": true, "ssl_cert_reqs": "None"}),
None
)]
fn a_config_raises_where_redis_py_raises_building_its_client(
#[case] fields: serde_json::Value,
#[case] raised: Option<&str>,
) {
Python::initialize();
Python::attach(|py| {
let error = litellm_redis_resolve::interpret(&kwargs(fields), None, None).unwrap_err();
let outcome = unresolved(py, &error)
.map(|client| client.is_some())
.map_err(|error| error.get_type(py).name().unwrap().to_string());
assert_eq!(
outcome,
raised.map_or(Ok(false), |class| Err(class.to_owned()))
);
});
}
#[rstest]
#[case::overflow(f64::INFINITY, "OverflowError")]
#[case::not_a_number(f64::NAN, "ValueError")]
fn a_default_ttl_python_cannot_make_an_int_raises_like_int(
#[case] seconds: f64,
#[case] raised: &str,
) {
Python::initialize();
Python::attach(|py| {
let error = number::seconds(seconds).map_err(seconds_error).unwrap_err();
assert_eq!(error.get_type(py).name().unwrap().to_string(), raised);
});
}
}

403
litellm-rust/crates/python-bridge/src/redis/events.rs Normal file
View file

@ -0,0 +1,403 @@
use std::{
sync::Arc,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_circuit_breaker::{State, Transition};
use litellm_host_python::{PythonContext, attach_blocking};
use litellm_redis_client::{
CallEvent, CallTag, FailureKind, Operation, Outcome, RedisEvents, current_tag, tagged,
};
use pyo3::{prelude::*, types::PyModule};
use tokio::runtime::Handle;
/// The asyncio loop a Python caller ran on. `RedisCache` scheduled its service events there with
/// `asyncio.create_task`, so the native events go there too instead of a new loop per event.
#[derive(Clone)]
pub(crate) struct EventLoop(Arc<Py<PyAny>>);
impl EventLoop {
pub(crate) fn running(py: Python<'_>) -> Option<Self> {
pyo3_async_runtimes::get_running_loop(py)
.ok()
.map(|event_loop| Self(Arc::new(event_loop.unbind())))
}
}
/// The Python caller of one store operation. Its context carries `with_service_target` and the
/// ambient OTel span, and its `call_type` is the label `RedisCache` reports, so dashboards keep
/// their series.
pub(crate) struct RedisCall {
pub(crate) context: PythonContext,
pub(crate) event_loop: Option<EventLoop>,
pub(crate) call_type: &'static str,
}
impl RedisCall {
pub(crate) fn tag(self) -> CallTag {
Arc::new(self)
}
pub(crate) async fn scope<F: Future>(self, operation: F) -> F::Output {
tagged(self.tag(), operation).await
}
}
/// Reports every native Redis call and breaker transition to `ServiceLogging`.
pub(crate) struct ServiceLoggingEvents {
record_call: Arc<Py<PyAny>>,
record_breaker: Arc<Py<PyAny>>,
context: PythonContext,
/// The loop the client was built on, for calls whose caller had none running.
event_loop: Option<EventLoop>,
runtime: Handle,
}
impl ServiceLoggingEvents {
const MODULE: &str = "litellm.rust_bridge.host.redis_events";
pub(crate) fn new(py: Python<'_>, runtime: Handle) -> PyResult<Self> {
Self::from_module(
py,
&py.import(Self::MODULE)?,
runtime,
EventLoop::running(py),
)
}
/// For a breaker, whose transitions never run on the caller's loop, so building it never looks one up.
pub(crate) fn transitions_only(py: Python<'_>, runtime: Handle) -> PyResult<Self> {
Self::from_module(py, &py.import(Self::MODULE)?, runtime, None)
}
fn from_module(
py: Python<'_>,
module: &Bound<'_, PyModule>,
runtime: Handle,
event_loop: Option<EventLoop>,
) -> PyResult<Self> {
Ok(Self {
record_call: Arc::new(module.getattr("record_redis_call")?.unbind()),
record_breaker: Arc::new(module.getattr("record_redis_breaker")?.unbind()),
context: PythonContext::capture(py)?,
event_loop,
runtime,
})
}
fn report(&self, context: PythonContext, record: Arc<Py<PyAny>>, args: ReportArgs) {
self.runtime.spawn(async move {
let reported = attach_blocking(context, move |py| match args {
ReportArgs::Call(call) => record
.bind(py)
.call1((
call.event_loop
.as_ref()
.map(|event_loop| event_loop.0.bind(py)),
call.call_type,
call.start_time,
call.end_time,
call.failure,
))
.map(drop),
ReportArgs::Breaker(state) => record.bind(py).call1((state,)).map(drop),
})
.await
.and_then(|reported| reported);
if let Err(error) = reported {
tracing::warn!(error = %error, "reporting a Redis event to ServiceLogging failed");
}
});
}
}
/// One call as `record_redis_call` takes it, its times in `time.time()` seconds.
struct CallReport {
event_loop: Option<EventLoop>,
call_type: &'static str,
start_time: f64,
end_time: f64,
failure: Option<&'static str>,
}
enum ReportArgs {
Call(CallReport),
Breaker(&'static str),
}
/// The wall-clock start and end of a call that just finished after `duration`.
fn wall_clock_span(duration: Duration) -> (f64, f64) {
let end = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default();
(
end.saturating_sub(duration).as_secs_f64(),
end.as_secs_f64(),
)
}
fn operation_call_type(operation: Operation) -> &'static str {
match operation {
Operation::Command(name) => name,
Operation::Script => "eval",
Operation::Pipeline { .. } => "pipeline",
}
}
fn failure_label(outcome: Outcome) -> Option<&'static str> {
match outcome {
Outcome::Succeeded => None,
Outcome::Failed(FailureKind::Timeout) => Some("timeout"),
Outcome::Failed(FailureKind::Transport) => Some("transport"),
Outcome::Failed(FailureKind::Server) => Some("server"),
Outcome::Failed(FailureKind::Decode) => Some("decode"),
Outcome::Failed(FailureKind::BreakerOpen) => Some("breaker_open"),
}
}
pub(crate) fn state_label(state: &State) -> &'static str {
match state {
State::Closed { .. } => "closed",
State::Open { .. } => "open",
State::HalfOpen => "half_open",
}
}
impl RedisEvents for ServiceLoggingEvents {
fn call(&self, event: &CallEvent) {
let scoped = current_tag().and_then(|tag| tag.downcast::<RedisCall>().ok());
let (context, event_loop, call_type) = match scoped {
Some(call) => (
call.context.clone(),
call.event_loop.clone().or_else(|| self.event_loop.clone()),
call.call_type,
),
None => (
self.context.clone(),
self.event_loop.clone(),
operation_call_type(event.operation),
),
};
let (start_time, end_time) = wall_clock_span(event.duration);
self.report(
context,
self.record_call.clone(),
ReportArgs::Call(CallReport {
event_loop,
call_type,
start_time,
end_time,
failure: failure_label(event.outcome),
}),
);
}
fn breaker(&self, transition: &Transition<State>) {
self.report(
self.context.clone(),
self.record_breaker.clone(),
ReportArgs::Breaker(state_label(&transition.to)),
);
}
}
#[cfg(test)]
mod tests {
use std::{sync::Arc, time::Duration};
use litellm_host_python::PythonContext;
use litellm_redis_client::{CallEvent, FailureKind, Operation, Outcome, RedisEvents};
use pyo3::{
prelude::*,
types::{PyCFunction, PyModule, PyString},
};
use rstest::rstest;
use tokio::runtime::Handle;
use super::{EventLoop, RedisCall, ServiceLoggingEvents};
const RECORDER: &str = "
import contextvars, time
target = contextvars.ContextVar('target', default=None)
seen = []
def record_redis_call(event_loop, call_type, start_time, end_time, failure):
just_now = abs(time.time() - end_time) < 60
seen.append((event_loop, call_type, round(end_time - start_time, 3), just_now, failure, target.get()))
def record_redis_breaker(state):
seen.append(('breaker', state))
";
fn stand_in_loop(py: Python<'_>, name: &str) -> EventLoop {
EventLoop(Arc::new(PyString::new(py, name).into_any().unbind()))
}
fn recorder(name: &str, client_loop: Option<&str>) -> (ServiceLoggingEvents, Py<PyModule>) {
Python::initialize();
Python::attach(|py| {
let name = std::ffi::CString::new(format!("redis_events_recorder_{name}")).unwrap();
let module =
PyModule::from_code(py, &std::ffi::CString::new(RECORDER).unwrap(), &name, &name)
.unwrap();
let events = ServiceLoggingEvents::from_module(
py,
&module,
Handle::current(),
client_loop.map(|client_loop| stand_in_loop(py, client_loop)),
)
.unwrap();
(events, module.unbind())
})
}
async fn seen(module: &Py<PyModule>, count: usize) -> String {
let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
loop {
let seen: (usize, String) = Python::attach(|py| {
let seen = module.bind(py).getattr("seen").unwrap();
(seen.len().unwrap(), seen.repr().unwrap().to_string())
});
if seen.0 >= count || tokio::time::Instant::now() > deadline {
return seen.1;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
}
fn scoped_call(namespace: &Py<PyModule>, caller_loop: Option<&str>) -> RedisCall {
Python::attach(|py| {
namespace
.bind(py)
.getattr("target")
.unwrap()
.call_method1("set", ("spend_counter",))
.unwrap();
RedisCall {
context: PythonContext::capture(py).unwrap(),
event_loop: caller_loop.map(|caller_loop| stand_in_loop(py, caller_loop)),
call_type: "async_increment",
}
})
}
#[rstest]
#[case::on_the_callers_loop(Some("caller loop"), "caller loop")]
#[case::on_the_clients_loop_when_the_caller_had_none(None, "client loop")]
#[tokio::test(flavor = "multi_thread")]
async fn a_scoped_call_reports_the_store_labels_and_its_times_in_the_callers_context(
#[case] caller_loop: Option<&str>,
#[case] expected_loop: &str,
) {
let (events, namespace) = recorder(
expected_loop.replace(' ', "_").as_str(),
Some("client loop"),
);
let call = scoped_call(&namespace, caller_loop);
let event = CallEvent {
operation: Operation::Command("INCRBYFLOAT"),
duration: Duration::from_millis(250),
outcome: Outcome::Failed(FailureKind::Timeout),
};
call.scope(async { events.call(&event) }).await;
assert_eq!(
seen(&namespace, 1).await,
format!(
"[('{expected_loop}', 'async_increment', 0.25, True, 'timeout', 'spend_counter')]"
)
);
}
#[rstest]
#[case::opened(litellm_circuit_breaker::State::Open { since: std::time::Instant::now() }, "open")]
#[case::probing(litellm_circuit_breaker::State::HalfOpen, "half_open")]
#[tokio::test(flavor = "multi_thread")]
async fn breaker_transitions_report_python_state_names(
#[case] to: litellm_circuit_breaker::State,
#[case] expected: &str,
) {
let (events, namespace) = recorder(expected, None);
events.breaker(&litellm_circuit_breaker::Transition {
from: litellm_circuit_breaker::State::HalfOpen,
to,
generation: 1,
});
assert_eq!(
seen(&namespace, 1).await,
format!("[('breaker', '{expected}')]")
);
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn a_transition_on_a_thread_without_a_runtime_is_still_reported() {
let (events, namespace) = recorder("plain_thread", None);
std::thread::spawn(move || {
events.breaker(&litellm_circuit_breaker::Transition {
from: litellm_circuit_breaker::State::HalfOpen,
to: litellm_circuit_breaker::State::HalfOpen,
generation: 1,
});
})
.join()
.unwrap();
assert_eq!(seen(&namespace, 1).await, "[('breaker', 'half_open')]");
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn an_unscoped_call_falls_back_to_the_command_name_and_the_clients_loop() {
let (events, namespace) = recorder("unscoped", None);
events.call(&CallEvent {
operation: Operation::Pipeline { ops: 3 },
duration: Duration::from_millis(5),
outcome: Outcome::Succeeded,
});
assert_eq!(
seen(&namespace, 1).await,
"[(None, 'pipeline', 0.005, True, None, None)]"
);
}
#[rstest]
fn running_is_the_loop_the_caller_runs_on_and_none_outside_one() {
Python::initialize();
Python::attach(|py| {
let module = PyModule::from_code(
py,
c"import asyncio\nasync def probe_inside(probe):\n return probe() is asyncio.get_running_loop()\n",
c"event_loop_probe",
c"event_loop_probe",
)
.unwrap();
let probe = PyCFunction::new_closure(py, None, None, |_, _| {
Python::attach(|py| -> PyResult<Option<Py<PyAny>>> {
Ok(EventLoop::running(py).map(|event_loop| event_loop.0.clone_ref(py)))
})
})
.unwrap();
let coroutine = module
.getattr("probe_inside")
.unwrap()
.call1((probe,))
.unwrap();
let inside: bool = py
.import("asyncio")
.unwrap()
.call_method1("run", (coroutine,))
.unwrap()
.extract()
.unwrap();
assert!(inside);
assert!(EventLoop::running(py).is_none());
});
}
}

15
litellm-rust/crates/python-bridge/src/redis/mod.rs Normal file
View file

@ -0,0 +1,15 @@
use std::sync::LazyLock;
use litellm_redis_connect::RedisHosts;
pub(crate) mod breaker;
pub(crate) mod client;
mod events;
pub(crate) mod store;
static HOSTS: LazyLock<RedisHosts> = LazyLock::new(RedisHosts::default);
/// The Redis servers this process talks to, shared by every `RedisCache` and native store in it.
pub(crate) fn hosts() -> &'static RedisHosts {
&HOSTS
}

265
litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs Normal file
View file

@ -0,0 +1,265 @@
use std::sync::Arc;
use litellm_python_compat::Value;
use litellm_redis_store::auth_cache::{AuthCache, Settings};
use pyo3::prelude::*;
use super::{
NativePending, PyNumber, kind_error, pending, redis_async, redis_sync, remote,
request::{NativeRequestRedis, NativeSend},
ttl,
value::PyValue,
};
use crate::redis::client::NativeRedisClient;
fn values(values: Vec<Option<Value>>) -> Vec<Option<PyValue>> {
values.into_iter().map(|value| value.map(PyValue)).collect()
}
/// `UserApiKeyCache` with `enable_redis_auth_cache` on. The sync methods back `DualCache`'s sync
/// API and block with the GIL released, like redis-py's sync client.
#[pyclass(frozen, name = "NativeAuthCache")]
pub(crate) struct NativeAuthCache(Arc<AuthCache>);
#[pymethods]
impl NativeAuthCache {
#[new]
#[pyo3(signature = (client, *, memory_entries, memory_ttl_seconds = None))]
fn new(
client: Option<PyRef<'_, NativeRedisClient>>,
memory_entries: usize,
memory_ttl_seconds: Option<f64>,
) -> PyResult<Self> {
Ok(Self(Arc::new(AuthCache::new(
client.map(|client| remote(&client)),
Settings {
memory_entries,
memory_ttl: ttl(memory_ttl_seconds)?,
},
))))
}
#[getter]
fn has_redis(&self) -> bool {
self.0.has_redis()
}
/// `request` answers a key its pre-call reads already found absent without another `GET`.
#[pyo3(signature = (key, request = None))]
fn get<'py>(
&self,
py: Python<'py>,
key: String,
request: Option<&NativeRequestRedis>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let request = request.map(|request| request.0.clone());
redis_async(py, "async_get_cache", kind_error, async move {
store
.get(request.as_ref(), &key)
.await
.map(|value| value.map(PyValue))
})
}
fn get_sync(&self, py: Python<'_>, key: String) -> PyResult<Option<PyValue>> {
let store = self.0.clone();
redis_sync(py, "get_cache", kind_error, async move {
store.get(None, &key).await.map(|value| value.map(PyValue))
})
}
/// `update_cache_ttl` and `update_in_memory_max_size`: both partitions keep what they hold.
#[pyo3(signature = (*, memory_entries, memory_ttl_seconds = None))]
fn reconfigure(&self, memory_entries: usize, memory_ttl_seconds: Option<f64>) -> PyResult<()> {
self.0.reconfigure(Settings {
memory_entries,
memory_ttl: ttl(memory_ttl_seconds)?,
});
Ok(())
}
#[pyo3(signature = (key, value, ttl_seconds = None))]
fn set<'py>(
&self,
py: Python<'py>,
key: String,
value: PyValue,
ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let ttl = ttl(ttl_seconds)?;
redis_async(py, "async_set_cache", kind_error, async move {
store.set(&key, value.0, ttl).await
})
}
#[pyo3(signature = (key, value, ttl_seconds = None))]
fn set_sync(
&self,
py: Python<'_>,
key: String,
value: PyValue,
ttl_seconds: Option<f64>,
) -> PyResult<()> {
let store = self.0.clone();
let ttl = ttl(ttl_seconds)?;
redis_sync(py, "set_cache", kind_error, async move {
store.set_str(&key, value.0, ttl).await
})
}
/// Memory now and the `SET` on the batch `send` picks.
#[pyo3(signature = (key, value, ttl_seconds, send))]
fn declare_set(
&self,
key: String,
value: PyValue,
ttl_seconds: Option<f64>,
send: &NativeSend,
) -> PyResult<NativePending> {
let written = self
.0
.declare_set(send.send(), &key, value.0, ttl(ttl_seconds)?)
.map_err(kind_error)?;
Ok(pending(
"async_set_cache",
kind_error,
written,
std::convert::identity,
))
}
fn declare_delete(&self, key: String, send: &NativeSend) -> PyResult<NativePending> {
let deleted = self
.0
.declare_delete(send.send(), &key)
.map_err(kind_error)?;
Ok(pending(
"async_delete_cache",
kind_error,
deleted,
std::convert::identity,
))
}
#[pyo3(signature = (entries, ttl_seconds = None))]
fn set_many<'py>(
&self,
py: Python<'py>,
entries: Vec<(String, PyValue)>,
ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let entries = entries
.into_iter()
.map(|(key, value)| (key, value.0))
.collect();
let ttl = ttl(ttl_seconds)?;
redis_async(py, "async_set_cache_pipeline", kind_error, async move {
store.set_many(entries, ttl).await
})
}
/// On a request batch it returns before anything is sent, and `Now` is a pipeline of its own.
fn write_back<'py>(
&self,
py: Python<'py>,
entries: Vec<(String, PyValue, Option<f64>)>,
send: &NativeSend,
) -> PyResult<Bound<'py, PyAny>> {
let entries = entries
.into_iter()
.map(|(key, value, seconds)| Ok((key, value.0, ttl(seconds)?)))
.collect::<PyResult<_>>()?;
let written = self
.0
.write_back(send.send(), entries)
.map_err(kind_error)?;
redis_async(
py,
"async_set_cache_pipeline_with_ttls",
kind_error,
written,
)
}
fn delete<'py>(&self, py: Python<'py>, keys: Vec<String>) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_delete_cache", kind_error, async move {
store.delete(&keys).await
})
}
fn delete_sync(&self, py: Python<'_>, keys: Vec<String>) -> PyResult<()> {
let store = self.0.clone();
redis_sync(py, "delete_cache", kind_error, async move {
store.delete(&keys).await
})
}
/// One `MGET` of keys the caller found missing in memory, on the batch `send` picks; `None`
/// without Redis.
fn get_many(&self, keys: Vec<String>, send: &NativeSend) -> Option<NativePending> {
let read = self.0.get_many(send.send(), &keys)?;
Some(pending("async_batch_get_cache", kind_error, read, values))
}
fn increment_if_present<'py>(
&self,
py: Python<'py>,
key: String,
amount: PyNumber,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_increment_if_present", kind_error, async move {
store.increment_if_present(&key, amount.0).await
})
}
fn get_local(&self, key: &str) -> PyResult<Option<PyValue>> {
Ok(self.0.get_local(key).map_err(kind_error)?.map(PyValue))
}
#[pyo3(signature = (key, value, ttl_seconds = None))]
fn set_local(&self, key: &str, value: PyValue, ttl_seconds: Option<f64>) -> PyResult<()> {
self.0
.set_local(key, value.0, ttl(ttl_seconds)?)
.map_err(kind_error)
}
fn delete_local(&self, key: &str) -> PyResult<()> {
self.0.delete_local(key).map_err(kind_error)
}
fn flush_local(&self) -> PyResult<()> {
self.0.flush_local().map_err(kind_error)
}
fn len_local(&self) -> PyResult<PyLocalSizes> {
let (objects, key_objects) = self.0.len_local().map_err(kind_error)?;
Ok(PyLocalSizes {
objects,
key_objects,
})
}
}
/// Entries the management-object and the user-key partitions hold in memory.
#[pyclass(frozen, eq, get_all, name = "LocalSizes")]
#[derive(PartialEq)]
pub(crate) struct PyLocalSizes {
objects: usize,
key_objects: usize,
}
#[pymethods]
impl PyLocalSizes {
#[new]
fn new(objects: usize, key_objects: usize) -> Self {
Self {
objects,
key_objects,
}
}
}

271
litellm-rust/crates/python-bridge/src/redis/store/mod.rs Normal file
View file

@ -0,0 +1,271 @@
pub(crate) mod auth_cache;
mod policy;
pub(crate) mod rate_limits;
pub(crate) mod request;
pub(crate) mod router_cooldowns;
pub(crate) mod router_usage;
pub(crate) mod sensitive_route_pins;
pub(crate) mod spend_counters;
mod value;
use value::PyValue;
use std::{
pin::Pin,
sync::{Mutex, PoisonError},
time::Duration,
};
use litellm_host_python::PythonContext;
use litellm_python_compat::number::{SecondsError, seconds};
use litellm_redis_store::{Error, FromValue, IntoValue, Number, Pending, Remote};
use pyo3::{
IntoPyObjectExt,
exceptions::{PyOverflowError, PyRuntimeError, PyTypeError, PyValueError},
prelude::*,
types::{PyBool, PyFloat, PyInt},
};
use crate::execution::{run_async_value, run_sync_value};
use crate::redis::{
client::NativeRedisClient,
events::{EventLoop, RedisCall},
};
/// The labels `RedisCache` reports for this operation, in the caller's context so its
/// `service_target` reaches the service event, dispatched on the loop the caller runs on.
fn call(py: Python<'_>, call_type: &'static str) -> PyResult<RedisCall> {
Ok(RedisCall {
context: PythonContext::capture(py)?,
event_loop: EventLoop::running(py),
call_type,
})
}
/// Runs one Redis operation as a Python awaitable, reported to `ServiceLogging` as `call_type`.
fn redis_async<'py, T, E: 'static>(
py: Python<'py>,
call_type: &'static str,
error: fn(E) -> PyErr,
operation: impl Future<Output = Result<T, E>> + Send + 'static,
) -> PyResult<Bound<'py, PyAny>>
where
T: for<'a> IntoPyObject<'a> + Send + 'static,
{
let scope = call(py, call_type)?;
run_async_value(
py,
async move { scope.scope(operation).await.map_err(error) },
)
}
/// Runs one Redis operation with the GIL released, for the router's sync paths.
fn redis_sync<T, E: 'static>(
py: Python<'_>,
call_type: &'static str,
error: fn(E) -> PyErr,
operation: impl Future<Output = Result<T, E>> + Send + 'static,
) -> PyResult<T>
where
T: Send + 'static,
{
let scope = call(py, call_type)?;
run_sync_value(
py,
async move { scope.scope(operation).await.map_err(error) },
)
}
/// A TTL in seconds, raising what `int(seconds)` raises for infinity and NaN.
fn duration(value: f64) -> PyResult<Duration> {
seconds(value).map_err(|error| match error {
SecondsError::Overflow => PyOverflowError::new_err(error.to_string()),
SecondsError::NotANumber => PyValueError::new_err(error.to_string()),
})
}
fn ttl(value: Option<f64>) -> PyResult<Option<Duration>> {
value.map(duration).transpose()
}
fn remote(client: &NativeRedisClient) -> Remote {
Remote {
client: client.client.clone(),
default_ttl: client.default_ttl,
}
}
type Resolve = Pin<Box<dyn Future<Output = PyResult<Py<PyAny>>> + Send>>;
/// A Redis call declared on a request's batch. `resolve` flushes that batch if it has not gone
/// out and answers the call, once.
#[pyclass(frozen, name = "NativePending")]
pub(crate) struct NativePending {
resolve: Mutex<Option<Resolve>>,
call_type: &'static str,
}
impl NativePending {
pub(crate) fn new<T, E: 'static>(
call_type: &'static str,
error: fn(E) -> PyErr,
reply: impl Future<Output = Result<T, E>> + Send + 'static,
) -> Self
where
T: for<'a> IntoPyObject<'a> + Send + 'static,
{
let resolve = async move {
let reply = reply.await.map_err(error)?;
Python::attach(|py| reply.into_py_any(py))
};
Self {
resolve: Mutex::new(Some(Box::pin(resolve))),
call_type,
}
}
}
#[pymethods]
impl NativePending {
fn resolve<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let resolve = self
.resolve
.lock()
.unwrap_or_else(PoisonError::into_inner)
.take()
.ok_or_else(|| PyRuntimeError::new_err("this Redis call was already resolved"))?;
redis_async(py, self.call_type, std::convert::identity, resolve)
}
}
pyo3::create_exception!(
_native,
NativeRedisStoreError,
pyo3::exceptions::PyException,
"A native Redis store operation failed. Args are (kind, message), kind spelled as litellm-redis-store's ErrorKind."
);
/// What every native store raises, its kind telling a caller which failures it may drop.
fn kind_error(error: Error) -> PyErr {
let kind: &'static str = error.kind().into();
NativeRedisStoreError::new_err((kind, litellm_error::Report::new(&error).to_string()))
}
/// A store op declared on the batch the caller's `NativeSend` picked, answered as `reply` shapes it.
fn pending<T, R>(
call_type: &'static str,
error: fn(Error) -> PyErr,
declared: Pending<T>,
reply: fn(T) -> R,
) -> NativePending
where
T: Send + 'static,
R: for<'a> IntoPyObject<'a> + Send + 'static,
{
NativePending::new(call_type, error, async move { declared.await.map(reply) })
}
/// A Python `int` or `float` across the bridge, keeping which one it is.
#[derive(Clone, Copy, PartialEq)]
pub(crate) struct PyNumber(Number);
impl<'a, 'py> FromPyObject<'a, 'py> for PyNumber {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
if value.is_instance_of::<PyInt>() && !value.is_instance_of::<PyBool>() {
return Ok(Self(Number::Int(value.extract()?)));
}
if value.is_instance_of::<PyFloat>() {
return Ok(Self(Number::Float(value.extract()?)));
}
Err(PyTypeError::new_err("expected an int or a float"))
}
}
impl<'py> IntoPyObject<'py> for PyNumber {
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
Ok(match self.0 {
Number::Int(value) => value.into_pyobject(py)?.into_any(),
Number::Float(value) => value.into_pyobject(py)?.into_any(),
})
}
}
/// A store enum read from its wire spelling, so the bridge names no variant.
pub(crate) struct Spelled<T>(pub(crate) T);
impl<'a, 'py, T: std::str::FromStr> FromPyObject<'a, 'py> for Spelled<T> {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let spelling: String = value.extract()?;
spelling
.parse()
.map(Self)
.map_err(|_| PyValueError::new_err(format!("unknown spelling {spelling:?}")))
}
}
/// A store type whose Python side is plain data, converted through its `FromValue` and `IntoValue`
/// so the bridge never names a field. A shape the store refuses is `ValueError`.
pub(crate) struct PyData<T>(pub(crate) T);
impl<'a, 'py, T: FromValue> FromPyObject<'a, 'py> for PyData<T> {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let PyValue(value) = value.extract()?;
T::from_value(&value)
.map(Self)
.map_err(|error| PyValueError::new_err(error.to_string()))
}
}
impl<'py, T: IntoValue> IntoPyObject<'py> for PyData<T> {
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
PyValue(self.0.into_value()).into_pyobject(py)
}
}
fn unwrap_all<T>(values: Vec<PyData<T>>) -> Vec<T> {
values.into_iter().map(|value| value.0).collect()
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use pyo3::prelude::*;
use rstest::rstest;
use super::ttl;
#[rstest]
#[case::infinity(f64::INFINITY, "OverflowError")]
#[case::nan(f64::NAN, "ValueError")]
fn a_ttl_int_refuses_raises_what_int_raises(#[case] seconds: f64, #[case] raised: &str) {
Python::initialize();
Python::attach(|py| {
let error = ttl(Some(seconds)).unwrap_err();
assert_eq!(error.get_type(py).name().unwrap().to_string(), raised);
});
}
#[rstest]
#[case::fraction(Some(1.5), Some(Duration::from_millis(1500)))]
#[case::negative_clamps_to_zero(Some(-1.0), Some(Duration::ZERO))]
#[case::absent(None, None)]
fn a_finite_ttl_converts(#[case] seconds: Option<f64>, #[case] expected: Option<Duration>) {
assert_eq!(ttl(seconds).unwrap(), expected);
}
}

151
litellm-rust/crates/python-bridge/src/redis/store/policy.rs Normal file
View file

@ -0,0 +1,151 @@
use std::time::Duration;
use litellm_redis_store::TierPolicy;
use pyo3::prelude::*;
use crate::{
coercion::{Field, FieldSpec, ProjectionError},
python_settings::{PythonSettings, Snapshot},
};
const COOLDOWN_REFRESH: FieldSpec<Duration> =
FieldSpec::new("cooldown_redis_read_interval", seconds);
const USAGE_REFRESH: FieldSpec<Duration> = FieldSpec::new("usage_redis_read_interval", seconds);
const TRACKED_READS: FieldSpec<usize> = FieldSpec::new("max_redis_batch_cache_size", count);
const MEMORY_TTL: FieldSpec<Option<Duration>> =
FieldSpec::new("default_in_memory_ttl", optional_seconds);
/// The tier values this Python process runs with, from its environment and `litellm` globals.
#[derive(Debug)]
pub(crate) struct RouterTierPolicies {
pub(crate) cooldowns: TierPolicy,
pub(crate) usage: TierPolicy,
}
pub(crate) fn router_tier_policies(py: Python<'_>) -> PyResult<RouterTierPolicies> {
Ok(project(&PythonSettings::RouterTiers.read(py)?)?)
}
fn project(snapshot: &Snapshot<'_>) -> Result<RouterTierPolicies, ProjectionError> {
let tracked_reads = snapshot.read(&TRACKED_READS)?;
let default_in_memory_ttl = snapshot.read(&MEMORY_TTL)?;
Ok(RouterTierPolicies {
cooldowns: TierPolicy {
refresh_after: snapshot.read(&COOLDOWN_REFRESH)?,
tracked_reads,
default_in_memory_ttl,
},
usage: TierPolicy {
refresh_after: snapshot.read(&USAGE_REFRESH)?,
tracked_reads,
default_in_memory_ttl,
},
})
}
fn seconds(field: &Field<'_>) -> Result<Duration, ProjectionError> {
match field.value().extract::<f64>() {
Ok(seconds) if seconds.is_finite() && seconds >= 0.0 => {
Ok(Duration::from_secs_f64(seconds))
}
_ => Err(field.invalid("a non-negative number of seconds")),
}
}
fn optional_seconds(field: &Field<'_>) -> Result<Option<Duration>, ProjectionError> {
match field.value().is_none() {
true => Ok(None),
false => seconds(field).map(Some),
}
}
fn count(field: &Field<'_>) -> Result<usize, ProjectionError> {
field
.value()
.extract::<usize>()
.map_err(|_| field.invalid("a non-negative integer"))
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use pyo3::{prelude::*, types::PyDict};
use super::project;
use crate::python_settings::PythonSettings;
fn projected(py: Python<'_>, fields: &str) -> Result<super::RouterTierPolicies, String> {
let locals = PyDict::new(py);
let source = std::ffi::CString::new(format!(
"import types\nsettings = types.SimpleNamespace({fields})"
))
.unwrap();
py.run(&source, Some(&locals), Some(&locals)).unwrap();
let snapshot =
PythonSettings::RouterTiers.snapshot(locals.get_item("settings").unwrap().unwrap());
project(&snapshot).map_err(|error| PyErr::from(error).to_string())
}
#[test]
fn the_host_values_are_each_familys_refresh_and_tracked_reads() {
Python::initialize();
Python::attach(|py| {
let policies = projected(
py,
"cooldown_redis_read_interval=2.5, usage_redis_read_interval=30, max_redis_batch_cache_size=42, default_in_memory_ttl=None",
)
.unwrap();
assert_eq!(
policies.cooldowns.refresh_after,
Duration::from_millis(2500)
);
assert_eq!(policies.usage.refresh_after, Duration::from_secs(30));
assert_eq!(policies.cooldowns.tracked_reads, 42);
assert_eq!(policies.usage.tracked_reads, 42);
assert_eq!(policies.cooldowns.default_in_memory_ttl, None);
});
}
#[test]
fn a_set_default_in_memory_ttl_reaches_both_familys_backfill() {
Python::initialize();
Python::attach(|py| {
let policies = projected(
py,
"cooldown_redis_read_interval=1, usage_redis_read_interval=10, max_redis_batch_cache_size=1000, default_in_memory_ttl=7.5",
)
.unwrap();
assert_eq!(
policies.cooldowns.default_in_memory_ttl,
Some(Duration::from_millis(7500))
);
assert_eq!(
policies.usage.default_in_memory_ttl,
Some(Duration::from_millis(7500))
);
});
}
#[test]
fn a_value_python_would_not_run_with_is_a_configuration_error() {
Python::initialize();
Python::attach(|py| {
let negative = projected(
py,
"cooldown_redis_read_interval=-1, usage_redis_read_interval=10, max_redis_batch_cache_size=1000, default_in_memory_ttl=None",
)
.unwrap_err();
let text = projected(
py,
"cooldown_redis_read_interval=1, usage_redis_read_interval=10, max_redis_batch_cache_size='many', default_in_memory_ttl=None",
)
.unwrap_err();
assert!(
negative.contains("cooldown_redis_read_interval"),
"{negative}"
);
assert!(text.contains("max_redis_batch_cache_size"), "{text}");
});
}
}

577
litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs Normal file
View file

@ -0,0 +1,577 @@
use std::{sync::Arc, time::Duration};
use indexmap::IndexMap;
use litellm_redis_store::{
Pending,
rate_limits::{
AcquireReply, Applied, AtomicCounter, AtomicReply, Counter, Entry, Gauge, GuardedIncrement,
GuardedReply, Increment, Kind, Local, RateLimits, Scope, WindowCount,
},
};
use pyo3::{exceptions::PyValueError, prelude::*};
use super::{
NativePending, PyNumber, Spelled, kind_error, redis_async, remote,
request::{NativeRequestRedis, NativeSend},
ttl,
};
use crate::redis::client::NativeRedisClient;
/// The v3 rate limiter's keyspace. Memory reads and writes are synchronous like `InMemoryCache`;
/// every Redis call is a future.
#[pyclass(frozen, name = "NativeRateLimits")]
pub(crate) struct NativeRateLimits(Arc<RateLimits>);
/// An `int` or `float` counter, a `str` window start, or a `{slot_id: acquired_at}` registry.
#[derive(FromPyObject)]
enum PyLocal {
Number(PyNumber),
Text(String),
Slots(IndexMap<String, f64>),
}
impl From<PyLocal> for Local {
fn from(local: PyLocal) -> Self {
match local {
PyLocal::Number(number) => Self::Number(number.0),
PyLocal::Text(text) => Self::Text(text),
PyLocal::Slots(slots) => Self::Slots(slots.into_iter().collect()),
}
}
}
/// What the memory tier holds, back as the Python value it was written as.
struct LocalValue(Local);
impl<'py> IntoPyObject<'py> for LocalValue {
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
Ok(match self.0 {
Local::Number(number) => PyNumber(number).into_pyobject(py)?,
Local::Text(text) => text.into_pyobject(py)?.into_any(),
Local::Slots(slots) => slots
.into_iter()
.collect::<IndexMap<_, _>>()
.into_pyobject(py)?
.into_any(),
})
}
}
/// `litellm.redis.stores.rate_limits.RateLimitScope`.
#[derive(FromPyObject)]
struct PyScope {
descriptor_key: String,
descriptor_value: String,
}
impl From<PyScope> for Scope {
fn from(scope: PyScope) -> Self {
Self {
descriptor: scope.descriptor_key,
value: scope.descriptor_value,
}
}
}
/// `RateLimitCounter`.
#[derive(FromPyObject)]
struct PyCounter {
scope: PyScope,
kind: Spelled<Kind>,
}
impl From<PyCounter> for Counter {
fn from(counter: PyCounter) -> Self {
Self {
scope: counter.scope.into(),
kind: counter.kind.0,
}
}
}
/// `RateLimitWindow`, `RateLimitCounter` or `RateLimitGauge`, told apart by the key suffix each names.
#[derive(FromPyObject)]
struct PyEntry {
suffix: String,
scope: PyScope,
}
impl TryFrom<PyEntry> for Entry {
type Error = PyErr;
fn try_from(entry: PyEntry) -> PyResult<Self> {
Self::new(&entry.suffix, entry.scope.into())
.map_err(|_| PyValueError::new_err(format!("unknown key suffix {:?}", entry.suffix)))
}
}
/// `CounterIncrement`.
#[derive(FromPyObject)]
struct PyIncrement {
counter: PyCounter,
amount: PyNumber,
ttl: Option<u64>,
}
impl From<PyIncrement> for Increment {
fn from(increment: PyIncrement) -> Self {
Self {
counter: increment.counter.into(),
amount: increment.amount.0,
ttl: increment.ttl.map(Duration::from_secs),
}
}
}
/// `AtomicCounter`.
#[derive(FromPyObject)]
struct PyAtomicCounter {
counter: PyCounter,
limit: i64,
increment: i64,
ttl: i64,
window_size: i64,
}
impl From<PyAtomicCounter> for AtomicCounter {
fn from(counter: PyAtomicCounter) -> Self {
Self {
counter: counter.counter.into(),
limit: counter.limit,
increment: counter.increment,
ttl_seconds: counter.ttl,
window_size_seconds: counter.window_size,
}
}
}
fn all<P, T: From<P>>(records: Vec<P>) -> Vec<T> {
records.into_iter().map(T::from).collect()
}
#[pyclass(frozen, eq, get_all, name = "WindowCount")]
#[derive(PartialEq)]
pub(crate) struct PyWindowCount {
window_start: String,
count: i64,
}
#[pymethods]
impl PyWindowCount {
#[new]
fn new(window_start: String, count: i64) -> Self {
Self {
window_start,
count,
}
}
}
/// Per counter, the count after the increment and the start of the window it landed in.
#[pyclass(frozen, eq, get_all, name = "AtomicApplied")]
#[derive(PartialEq)]
pub(crate) struct PyAtomicApplied {
counters: Vec<(i64, i64)>,
}
#[pymethods]
impl PyAtomicApplied {
#[new]
fn new(counters: Vec<(i64, i64)>) -> Self {
Self { counters }
}
}
#[pyclass(frozen, eq, get_all, name = "AtomicOverLimit")]
#[derive(PartialEq)]
pub(crate) struct PyAtomicOverLimit {
index: usize,
current: i64,
limit: i64,
}
#[pymethods]
impl PyAtomicOverLimit {
#[new]
fn new(index: usize, current: i64, limit: i64) -> Self {
Self {
index,
current,
limit,
}
}
}
#[derive(IntoPyObject)]
enum PyAtomicReply {
Applied(PyAtomicApplied),
OverLimit(PyAtomicOverLimit),
}
#[pyclass(frozen, eq, get_all, name = "GuardedReply")]
#[derive(PartialEq)]
pub(crate) struct PyGuardedReply {
applied: bool,
count: i64,
}
#[pymethods]
impl PyGuardedReply {
#[new]
fn new(applied: bool, count: i64) -> Self {
Self { applied, count }
}
}
#[pyclass(frozen, eq, get_all, name = "GaugesAcquired")]
#[derive(PartialEq)]
pub(crate) struct PyGaugesAcquired {
in_flight: Vec<i64>,
}
#[pymethods]
impl PyGaugesAcquired {
#[new]
fn new(in_flight: Vec<i64>) -> Self {
Self { in_flight }
}
}
#[pyclass(frozen, eq, get_all, name = "GaugeOverLimit")]
#[derive(PartialEq)]
pub(crate) struct PyGaugeOverLimit {
index: usize,
in_flight: i64,
limit: i64,
}
#[pymethods]
impl PyGaugeOverLimit {
#[new]
fn new(index: usize, in_flight: i64, limit: i64) -> Self {
Self {
index,
in_flight,
limit,
}
}
}
#[derive(IntoPyObject)]
enum PyAcquireReply {
Acquired(PyGaugesAcquired),
OverLimit(PyGaugeOverLimit),
}
fn window_counts(counts: Vec<WindowCount>) -> Vec<PyWindowCount> {
counts
.into_iter()
.map(
|WindowCount {
window_start,
count,
}| PyWindowCount {
window_start,
count,
},
)
.collect()
}
fn atomic(reply: AtomicReply) -> PyAtomicReply {
match reply {
AtomicReply::Applied(applied) => PyAtomicReply::Applied(PyAtomicApplied {
counters: applied
.into_iter()
.map(
|Applied {
count,
window_start,
}| (count, window_start),
)
.collect(),
}),
AtomicReply::OverLimit {
index,
current,
limit,
} => PyAtomicReply::OverLimit(PyAtomicOverLimit {
index,
current,
limit,
}),
}
}
fn guarded(GuardedReply { applied, count }: GuardedReply) -> PyGuardedReply {
PyGuardedReply { applied, count }
}
fn acquire(reply: AcquireReply) -> PyAcquireReply {
match reply {
AcquireReply::Acquired(in_flight) => {
PyAcquireReply::Acquired(PyGaugesAcquired { in_flight })
}
AcquireReply::OverLimit {
index,
in_flight,
limit,
} => PyAcquireReply::OverLimit(PyGaugeOverLimit {
index,
in_flight,
limit,
}),
}
}
/// One script call declared on the batch the caller's `NativeSend` picked; `None` without Redis.
fn script<T, R>(
call_type: &'static str,
declared: Option<Pending<T>>,
reply: fn(T) -> R,
) -> Option<NativePending>
where
T: Send + 'static,
R: for<'a> IntoPyObject<'a> + Send + 'static,
{
Some(super::pending(call_type, kind_error, declared?, reply))
}
#[pymethods]
impl NativeRateLimits {
#[new]
#[pyo3(signature = (client, backfill_ttl_seconds = None))]
fn new(
client: Option<PyRef<'_, NativeRedisClient>>,
backfill_ttl_seconds: Option<f64>,
) -> PyResult<Self> {
Ok(Self(Arc::new(RateLimits::new(
client.map(|client| remote(&client)),
ttl(backfill_ttl_seconds)?,
))))
}
#[getter]
fn has_redis(&self) -> bool {
self.0.has_redis()
}
fn slot_groups(&self, scopes: Vec<PyScope>, force: bool) -> Vec<Vec<usize>> {
self.0.slot_groups(&all(scopes), force)
}
fn get_local(&self, entry: PyEntry) -> PyResult<Option<LocalValue>> {
Ok(self
.0
.get_local(&entry.try_into()?)
.map_err(kind_error)?
.map(LocalValue))
}
#[pyo3(signature = (entry, value, ttl_seconds))]
fn set_local(&self, entry: PyEntry, value: PyLocal, ttl_seconds: Option<f64>) -> PyResult<()> {
self.0
.set_local(&entry.try_into()?, value.into(), ttl(ttl_seconds)?)
.map_err(kind_error)
}
/// `request` answers a key its pre-call reads already found absent without another `GET`.
#[pyo3(signature = (entry, *, backfill_ttl_seconds = None, request = None))]
fn get<'py>(
&self,
py: Python<'py>,
entry: PyEntry,
backfill_ttl_seconds: Option<f64>,
request: Option<&NativeRequestRedis>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let entry = Entry::try_from(entry)?;
let backfill_ttl = ttl(backfill_ttl_seconds)?;
let request = request.map(|request| request.0.clone());
redis_async(py, "async_get_cache", kind_error, async move {
store
.get(request.as_ref(), &entry, backfill_ttl)
.await
.map(|local| local.map(LocalValue))
})
}
fn sliding_window(
&self,
counters: Vec<PyCounter>,
now: i64,
window_size: i64,
send: &NativeSend,
) -> Option<NativePending> {
script(
"batch_rate_limiter_script",
self.0
.sliding_window(send.send(), &all(counters), now, window_size),
window_counts,
)
}
fn read_counters(&self, counters: Vec<PyCounter>, send: &NativeSend) -> Option<NativePending> {
script(
"batch_counter_read_script",
self.0.read_counters(send.send(), &all(counters)),
std::convert::identity,
)
}
fn check_and_increment(
&self,
counters: Vec<PyAtomicCounter>,
send: &NativeSend,
) -> Option<NativePending> {
script(
"check_and_increment_by_n_script",
self.0.check_and_increment(send.send(), &all(counters)),
atomic,
)
}
fn guarded_increment(
&self,
counter: PyCounter,
expected_window_start: String,
amount: PyNumber,
ttl_seconds: Option<u64>,
send: &NativeSend,
) -> Option<NativePending> {
let increment = GuardedIncrement {
counter: counter.into(),
expected_window_start,
amount: amount.0,
ttl: ttl_seconds.map(Duration::from_secs),
};
script(
"window_guarded_token_increment_script",
self.0.guarded_increment(send.send(), &increment),
guarded,
)
}
fn parallel_acquire(
&self,
gauges: Vec<(PyScope, i64)>,
slot_ttl_seconds: u64,
slot_id: &str,
send: &NativeSend,
) -> Option<NativePending> {
let gauges: Vec<Gauge> = gauges
.into_iter()
.map(|(scope, limit)| Gauge {
scope: scope.into(),
limit,
})
.collect();
script(
"parallel_acquire_script",
self.0.parallel_acquire(
send.send(),
&gauges,
Duration::from_secs(slot_ttl_seconds),
slot_id,
),
acquire,
)
}
fn parallel_release(
&self,
gauges: Vec<PyScope>,
slot_id: &str,
send: &NativeSend,
) -> Option<NativePending> {
script(
"parallel_release_script",
self.0.parallel_release(send.send(), &all(gauges), slot_id),
std::convert::identity,
)
}
fn parallel_count(
&self,
gauges: Vec<PyScope>,
slot_ttl_seconds: u64,
send: &NativeSend,
) -> Option<NativePending> {
script(
"parallel_count_script",
self.0.parallel_count(
send.send(),
&all(gauges),
Duration::from_secs(slot_ttl_seconds),
),
std::convert::identity,
)
}
fn token_increment(
&self,
increments: Vec<PyIncrement>,
send: &NativeSend,
) -> Option<NativePending> {
script(
"token_increment_script",
self.0.token_increment(send.send(), &all(increments)),
std::convert::identity,
)
}
/// `False` without Redis.
fn token_increment_post_call(
&self,
request: &NativeRequestRedis,
increments: Vec<PyIncrement>,
) -> bool {
self.0
.token_increment_post_call(&request.0, all(increments))
}
/// `False` without Redis.
fn parallel_release_post_call(
&self,
request: &NativeRequestRedis,
gauges: Vec<PyScope>,
slot_id: &str,
) -> bool {
self.0
.parallel_release_post_call(&request.0, &all(gauges), slot_id)
}
fn increment<'py>(
&self,
py: Python<'py>,
increment: PyIncrement,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let increment = Increment::from(increment);
redis_async(py, "async_increment", kind_error, async move {
Ok(store.increment(&increment).await.map(PyNumber))
})
}
fn increment_many<'py>(
&self,
py: Python<'py>,
increments: Vec<PyIncrement>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let increments: Vec<Increment> = all(increments);
redis_async(py, "async_increment_pipeline", kind_error, async move {
store.increment_many(&increments).await;
Ok(())
})
}
fn increment_many_post_call(&self, request: &NativeRequestRedis, increments: Vec<PyIncrement>) {
self.0
.increment_many_post_call(&request.0, &all(increments));
}
}

60
litellm-rust/crates/python-bridge/src/redis/store/request.rs Normal file
View file

@ -0,0 +1,60 @@
use litellm_redis_store::{RequestRedis, Send};
use pyo3::prelude::*;
use super::{call, duration};
use crate::execution::run_async_value;
/// One request's native Redis batches. Python builds one per request scope and passes it to each
/// request-aware store call.
#[pyclass(frozen, name = "NativeRequestRedis")]
pub(crate) struct NativeRequestRedis(pub(crate) RequestRedis);
#[pymethods]
impl NativeRequestRedis {
#[new]
fn new(py: Python<'_>, post_call_deadline_seconds: f64) -> PyResult<Self> {
let runtime = litellm_host_python::runtime_handle()?;
let deadline = duration(post_call_deadline_seconds)?;
let deadline_tag = call(py, "post_call_redis_batch")?.tag();
Ok(Self(crate::logger::capture(py).scope(|| {
RequestRedis::with_deadline_tag(runtime, deadline, Some(deadline_tag))
})))
}
fn flush_pre<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let request = self.0.clone();
let scope = call(py, "request_redis_batch")?;
run_async_value(py, async move {
scope.scope(request.flush_pre()).await;
Ok(())
})
}
fn flush_post<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let request = self.0.clone();
let scope = call(py, "post_call_redis_batch")?;
run_async_value(py, async move {
scope.scope(request.flush_post()).await;
Ok(())
})
}
}
/// The Python face of `Send`: which batch a store op rides. Each variant carries exactly what it
/// needs, so there is no phase string and no phase without a request.
#[pyclass(frozen, name = "NativeSend")]
pub(crate) enum NativeSend {
Now(),
PreCall(Py<NativeRequestRedis>),
PostCall(Py<NativeRequestRedis>),
}
impl NativeSend {
pub(crate) fn send(&self) -> Send<'_> {
match self {
Self::Now() => Send::Now,
Self::PreCall(request) => Send::PreCall(&request.get().0),
Self::PostCall(request) => Send::PostCall(&request.get().0),
}
}
}

182
litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs Normal file
View file

@ -0,0 +1,182 @@
use std::{collections::HashMap, sync::Arc};
use litellm_redis_store::{
Send,
router_cooldowns::{Cooldown, RouterCooldowns},
};
use pyo3::{prelude::*, types::PyTuple};
use super::{
NativePending, PyData, PyNumber, duration, kind_error, pending, policy::router_tier_policies,
redis_async, redis_sync, remote, request::NativeSend,
};
use crate::redis::client::NativeRedisClient;
/// One router's cooldown store. The router's sync completion path calls `put`, `active_sync`,
/// `minimum` and `increment_allowed_fails`, so those block with the GIL released like redis-py; the
/// `async_` twins return awaitables for async callers.
#[pyclass(frozen, name = "NativeRouterCooldowns")]
pub(crate) struct NativeRouterCooldowns(Arc<RouterCooldowns>);
/// `[(model_id, cooldown), ...]` the way `RedisCache` hands the router its batch read.
struct PyActive(Vec<(String, Cooldown)>);
impl<'py> IntoPyObject<'py> for PyActive {
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
active_to_py(py, self.0)?
.into_pyobject(py)
.map(Bound::into_any)
}
}
fn active_to_py(
py: Python<'_>,
active: Vec<(String, Cooldown)>,
) -> PyResult<Vec<Bound<'_, PyTuple>>> {
active
.into_iter()
.map(|(model_id, cooldown)| {
PyTuple::new(
py,
[
model_id.into_pyobject(py)?.into_any(),
PyData(cooldown).into_pyobject(py)?,
],
)
})
.collect()
}
#[pymethods]
impl NativeRouterCooldowns {
#[new]
fn new(py: Python<'_>, client: Option<PyRef<'_, NativeRedisClient>>) -> PyResult<Self> {
let remote = client.map(|client| remote(&client));
let policy = router_tier_policies(py)?.cooldowns;
Ok(Self(Arc::new(RouterCooldowns::new(remote, policy))))
}
/// Adds the Redis tier and keeps memory; `false` when one was already attached.
fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool {
self.0.attach(remote(&client))
}
#[getter]
fn has_redis(&self) -> bool {
self.0.has_redis()
}
fn put(
&self,
py: Python<'_>,
model_id: String,
cooldown: PyData<Cooldown>,
ttl_seconds: f64,
) -> PyResult<()> {
let store = self.0.clone();
let ttl = duration(ttl_seconds)?;
redis_sync(py, "set_cache", kind_error, async move {
store.put_str(&model_id, &cooldown.0, ttl).await
})
}
fn increment_allowed_fails(
&self,
py: Python<'_>,
model_id: String,
suffix: Option<String>,
ttl_seconds: f64,
) -> PyResult<i64> {
let store = self.0.clone();
let ttl = duration(ttl_seconds)?;
redis_sync(py, "increment_cache", kind_error, async move {
store
.increment_allowed_fails_int(&model_id, suffix.as_deref(), ttl)
.await
})
}
fn async_put<'py>(
&self,
py: Python<'py>,
model_id: String,
cooldown: PyData<Cooldown>,
ttl_seconds: f64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let ttl = duration(ttl_seconds)?;
redis_async(py, "async_set_cache", kind_error, async move {
store.put(&model_id, &cooldown.0, ttl).await
})
}
fn async_increment_allowed_fails<'py>(
&self,
py: Python<'py>,
model_id: String,
suffix: Option<String>,
ttl_seconds: f64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let ttl = duration(ttl_seconds)?;
redis_async(py, "async_increment", kind_error, async move {
store
.increment_allowed_fails(&model_id, suffix.as_deref(), ttl)
.await
})
}
/// The deployments still cooling down, read on the batch `send` picks.
fn active(&self, model_ids: Vec<String>, send: &NativeSend) -> PyResult<NativePending> {
let active = self.0.active(send.send(), &model_ids).map_err(kind_error)?;
Ok(pending(
"async_batch_get_cache",
kind_error,
active,
PyActive,
))
}
/// `active` on its own with the GIL released, for the router's sync completion path.
fn active_sync<'py>(
&self,
py: Python<'py>,
model_ids: Vec<String>,
) -> PyResult<Vec<Bound<'py, PyTuple>>> {
let active = self.0.active(Send::Now, &model_ids).map_err(kind_error)?;
let active = redis_sync(py, "batch_get_cache", kind_error, active)?;
active_to_py(py, active)
}
fn minimum(&self, py: Python<'_>, model_ids: Vec<String>) -> PyResult<Option<PyNumber>> {
let store = self.0.clone();
redis_sync(py, "batch_get_cache", kind_error, async move {
store
.minimum(&model_ids)
.await
.map(|minimum| minimum.map(PyNumber))
})
}
fn delete_local(&self, model_id: &str) -> PyResult<()> {
self.0.delete_local(model_id).map_err(kind_error)
}
fn flush_local_counters(&self) -> PyResult<()> {
self.0.flush_local_counters().map_err(kind_error)
}
fn snapshot(&self) -> PyResult<HashMap<String, PyData<Cooldown>>> {
Ok(self
.0
.snapshot()
.map_err(kind_error)?
.into_iter()
.map(|(key, cooldown)| (key, PyData(cooldown)))
.collect())
}
}

182
litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs Normal file
View file

@ -0,0 +1,182 @@
use std::{collections::HashMap, sync::Arc, time::Duration};
use litellm_redis_store::router_usage::{Counter, Kind, RouterUsage};
use pyo3::prelude::*;
use super::{
PyNumber, Spelled, kind_error, policy::router_tier_policies, redis_async, redis_sync, remote,
};
use crate::redis::client::NativeRedisClient;
/// `litellm.redis.stores.router_usage.UsageCounter`.
#[derive(FromPyObject)]
struct PyCounter {
deployment_id: String,
model: Option<String>,
minute: String,
kind: Spelled<Kind>,
}
impl From<PyCounter> for Counter {
fn from(counter: PyCounter) -> Self {
Self {
deployment_id: counter.deployment_id,
model: counter.model,
minute: counter.minute,
kind: counter.kind.0,
}
}
}
fn counters(counters: Vec<PyCounter>) -> Vec<Counter> {
counters.into_iter().map(Counter::from).collect()
}
/// One router's global TPM/RPM counters. The `_sync` methods block with the GIL released, for the
/// Rust router's sync calls, which have no event loop to await on.
#[pyclass(frozen, name = "NativeRouterUsage")]
pub(crate) struct NativeRouterUsage(Arc<RouterUsage>);
#[pymethods]
impl NativeRouterUsage {
#[new]
fn new(py: Python<'_>, client: Option<PyRef<'_, NativeRedisClient>>) -> PyResult<Self> {
let remote = client.map(|client| remote(&client));
let policy = router_tier_policies(py)?.usage;
Ok(Self(Arc::new(RouterUsage::new(remote, policy))))
}
/// Adds the Redis tier and keeps memory; `false` when one was already attached.
fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool {
self.0.attach(remote(&client))
}
#[getter]
fn has_redis(&self) -> bool {
self.0.has_redis()
}
fn increment<'py>(
&self,
py: Python<'py>,
counter: PyCounter,
value: PyNumber,
ttl_seconds: u64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let counter = Counter::from(counter);
redis_async(py, "async_increment", kind_error, async move {
Ok(store
.increment(&counter, value.0, Duration::from_secs(ttl_seconds))
.await
.map(PyNumber))
})
}
/// `_increment_deployment`'s `async_increment_pipeline`, always on its own round trip.
fn increment_deployment<'py>(
&self,
py: Python<'py>,
deployment_id: String,
model: Option<String>,
minute: String,
tokens: PyNumber,
requests: i64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_increment_pipeline", kind_error, async move {
store
.increment_deployment(
&deployment_id,
model.as_deref(),
&minute,
tokens.0,
requests,
)
.await
})
}
fn get_many<'py>(
&self,
py: Python<'py>,
counters: Vec<PyCounter>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let counters = self::counters(counters);
redis_async(py, "async_batch_get_cache", kind_error, async move {
let read = store.get_many(&counters).await?;
Ok(read
.into_iter()
.map(|count| count.map(PyNumber))
.collect::<Vec<_>>())
})
}
fn increment_sync(
&self,
py: Python<'_>,
counter: PyCounter,
value: PyNumber,
ttl_seconds: u64,
) -> PyResult<Option<PyNumber>> {
let store = self.0.clone();
let counter = Counter::from(counter);
redis_sync(py, "increment_cache", kind_error, async move {
Ok(store
.increment(&counter, value.0, Duration::from_secs(ttl_seconds))
.await
.map(PyNumber))
})
}
fn increment_deployment_sync(
&self,
py: Python<'_>,
deployment_id: String,
model: Option<String>,
minute: String,
tokens: PyNumber,
requests: i64,
) -> PyResult<()> {
let store = self.0.clone();
redis_sync(py, "increment_pipeline", kind_error, async move {
store
.increment_deployment(
&deployment_id,
model.as_deref(),
&minute,
tokens.0,
requests,
)
.await
})
}
fn get_many_sync(
&self,
py: Python<'_>,
counters: Vec<PyCounter>,
) -> PyResult<Vec<Option<PyNumber>>> {
let store = self.0.clone();
let counters = self::counters(counters);
redis_sync(py, "batch_get_cache", kind_error, async move {
let read = store.get_many(&counters).await?;
Ok(read.into_iter().map(|count| count.map(PyNumber)).collect())
})
}
fn flush_local_counters(&self) -> PyResult<()> {
self.0.flush_local_counters().map_err(kind_error)
}
fn snapshot(&self) -> PyResult<HashMap<String, PyNumber>> {
Ok(self
.0
.snapshot()
.map_err(kind_error)?
.into_iter()
.map(|(key, count)| (key, PyNumber(count)))
.collect())
}
}

49
litellm-rust/crates/python-bridge/src/redis/store/sensitive_route_pins.rs Normal file
View file

@ -0,0 +1,49 @@
use std::sync::Arc;
use litellm_redis_store::sensitive_route_pins::SensitiveRoutePins;
use pyo3::prelude::*;
use super::{duration, kind_error, redis_async, remote};
use crate::redis::client::NativeRedisClient;
#[pyclass(frozen, name = "NativeSensitiveRoutePins")]
pub(crate) struct NativeSensitiveRoutePins(Arc<SensitiveRoutePins>);
#[pymethods]
impl NativeSensitiveRoutePins {
#[new]
fn new(client: Option<PyRef<'_, NativeRedisClient>>) -> Self {
Self(Arc::new(SensitiveRoutePins::new(
client.map(|client| remote(&client)),
)))
}
fn get<'py>(
&self,
py: Python<'py>,
tenant: String,
session_id: String,
default_ttl_seconds: f64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let default_ttl = duration(default_ttl_seconds)?;
redis_async(py, "async_get_cache", kind_error, async move {
store.get(&tenant, &session_id, default_ttl).await
})
}
fn set<'py>(
&self,
py: Python<'py>,
tenant: String,
session_id: String,
model: String,
ttl_seconds: f64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let ttl = duration(ttl_seconds)?;
redis_async(py, "async_set_cache", kind_error, async move {
store.set(&tenant, &session_id, &model, ttl).await
})
}
}

258
litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs Normal file
View file

@ -0,0 +1,258 @@
use std::{collections::HashMap, sync::Arc};
use litellm_redis_store::{
Number,
spend_counters::{LostSeed, Reading, SpendCounter, SpendCounters},
};
use pyo3::prelude::*;
use super::{
NativePending, PyData, PyNumber, kind_error, pending, redis_async, remote,
request::{NativeRequestRedis, NativeSend},
ttl, unwrap_all,
};
use crate::redis::client::NativeRedisClient;
/// A counter's value and whether Redis answered it, a clean miss included.
#[pyclass(frozen, eq, get_all, name = "SpendRead")]
#[derive(PartialEq)]
pub(crate) struct PySpendRead {
value: Option<PyNumber>,
from_redis: bool,
}
#[pymethods]
impl PySpendRead {
#[new]
fn new(value: Option<PyNumber>, from_redis: bool) -> Self {
Self { value, from_redis }
}
}
fn increments(increments: Vec<(PyData<SpendCounter>, PyNumber)>) -> Vec<(SpendCounter, Number)> {
increments
.into_iter()
.map(|(counter, amount)| (counter.0, amount.0))
.collect()
}
fn numbers(values: Option<Vec<Option<Number>>>) -> Option<Vec<Option<PyNumber>>> {
values.map(|values| {
values
.into_iter()
.map(|value| value.map(PyNumber))
.collect()
})
}
/// The proxy's spend counters, Redis first with this worker's memory copy behind it.
#[pyclass(frozen, name = "NativeSpendCounters")]
pub(crate) struct NativeSpendCounters(Arc<SpendCounters>);
#[pymethods]
impl NativeSpendCounters {
#[new]
fn new(client: Option<PyRef<'_, NativeRedisClient>>) -> Self {
Self(Arc::new(SpendCounters::new(
client.map(|client| remote(&client)),
)))
}
/// Adds the Redis tier and keeps memory; `false` when one was already attached.
fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool {
self.0.attach(remote(&client))
}
#[getter]
fn has_redis(&self) -> bool {
self.0.has_redis()
}
#[getter]
fn redis_ttl_seconds(&self) -> Option<f64> {
self.0.redis_ttl().map(|ttl| ttl.as_secs_f64())
}
/// `("redis" | "memory", value)`: where the answer came from, a Redis miss included.
fn read<'py>(
&self,
py: Python<'py>,
counter: PyData<SpendCounter>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_get_cache", kind_error, async move {
let reading = store.read(&counter.0).await?;
Ok(PySpendRead {
value: reading.value().map(PyNumber),
from_redis: matches!(reading, Reading::Redis(_)),
})
})
}
/// One `MGET` on the batch `send` picks; `None` without Redis. It resolves to `None` when the
/// read failed, so each caller reads on its own.
fn read_many(
&self,
counters: Vec<PyData<SpendCounter>>,
send: &NativeSend,
) -> Option<NativePending> {
let read = self.0.read_many(send.send(), &unwrap_all(counters))?;
Some(pending("async_batch_get_cache", kind_error, read, numbers))
}
fn increment<'py>(
&self,
py: Python<'py>,
counter: PyData<SpendCounter>,
amount: PyNumber,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_increment", kind_error, async move {
store.increment(&counter.0, amount.0).await.map(PyNumber)
})
}
fn increment_many<'py>(
&self,
py: Python<'py>,
increments: Vec<(PyData<SpendCounter>, PyNumber)>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let increments = self::increments(increments);
redis_async(py, "async_increment_pipeline", kind_error, async move {
let totals = store.increment_many(&increments).await?;
Ok(totals.into_iter().map(PyNumber).collect::<Vec<_>>())
})
}
/// Declares the increments on the request's post-call batch; `False` when there is no Redis.
fn increment_many_post_call(
&self,
request: &NativeRequestRedis,
increments: Vec<(PyData<SpendCounter>, PyNumber)>,
) -> bool {
self.0
.increment_many_post_call(&request.0, &self::increments(increments))
}
fn set_max<'py>(
&self,
py: Python<'py>,
counter: PyData<SpendCounter>,
spend: PyNumber,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_set_max", kind_error, async move {
store
.set_max(&counter.0, spend.0)
.await
.map(|stored| stored.map(PyNumber))
})
}
#[pyo3(signature = (counter, spend, *, increment_if_lost))]
fn seed<'py>(
&self,
py: Python<'py>,
counter: PyData<SpendCounter>,
spend: PyNumber,
increment_if_lost: bool,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let lost = if increment_if_lost {
LostSeed::Increment
} else {
LostSeed::Answer
};
redis_async(py, "async_set_cache", kind_error, async move {
store.seed(&counter.0, spend.0, lost).await.map(PyNumber)
})
}
#[pyo3(signature = (counter, spend, ttl_seconds = None))]
fn set<'py>(
&self,
py: Python<'py>,
counter: PyData<SpendCounter>,
spend: PyNumber,
ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let ttl = ttl(ttl_seconds)?;
redis_async(py, "async_set_cache", kind_error, async move {
store.set(&counter.0, spend.0, ttl).await
})
}
fn delete<'py>(
&self,
py: Python<'py>,
counters: Vec<PyData<SpendCounter>>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let counters = unwrap_all(counters);
redis_async(py, "async_delete_cache", kind_error, async move {
store.delete(&counters).await
})
}
fn refresh_ttl<'py>(
&self,
py: Python<'py>,
counter: PyData<SpendCounter>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_refresh_ttl", kind_error, async move {
store.refresh_ttl(&counter.0).await
})
}
fn get_local(&self, counter: PyData<SpendCounter>) -> PyResult<Option<PyNumber>> {
Ok(self
.0
.get_local(&counter.0)
.map_err(kind_error)?
.map(PyNumber))
}
#[pyo3(signature = (counter, spend, ttl_seconds = None))]
fn set_local(
&self,
counter: PyData<SpendCounter>,
spend: PyNumber,
ttl_seconds: Option<f64>,
) -> PyResult<()> {
self.0
.set_local(&counter.0, spend.0, ttl(ttl_seconds)?)
.map_err(kind_error)
}
fn delete_local(&self, counter: PyData<SpendCounter>) -> PyResult<()> {
self.0.delete_local(&counter.0).map_err(kind_error)
}
fn increment_local(
&self,
counter: PyData<SpendCounter>,
amount: PyNumber,
) -> PyResult<PyNumber> {
self.0
.increment_local(&counter.0, amount.0)
.map(PyNumber)
.map_err(kind_error)
}
fn flush_local(&self) -> PyResult<()> {
self.0.flush_local().map_err(kind_error)
}
fn snapshot(&self) -> PyResult<HashMap<String, PyNumber>> {
Ok(self
.0
.snapshot()
.map_err(kind_error)?
.into_iter()
.map(|(key, spend)| (key, PyNumber(spend)))
.collect())
}
}

125
litellm-rust/crates/python-bridge/src/redis/store/value.rs Normal file
View file

@ -0,0 +1,125 @@
use litellm_python_compat::{BigInt, MAX_DEPTH, Value};
use pyo3::{
exceptions::{PyTypeError, PyValueError},
prelude::*,
types::{
PyBool, PyBytes, PyComplex, PyDict, PyFloat, PyFrozenSet, PyInt, PyList, PySet, PyString,
PyTuple,
},
};
/// A plain Python value (what `json.loads` or `ast.literal_eval` returns) across the bridge.
/// Anything else, a live object included, is refused with `TypeError`.
pub(crate) struct PyValue(pub(crate) Value);
impl<'a, 'py> FromPyObject<'a, 'py> for PyValue {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
to_value(&value, 0).map(Self)
}
}
impl<'py> IntoPyObject<'py> for PyValue {
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
to_python(py, self.0)
}
}
fn to_value(value: &Bound<'_, PyAny>, depth: usize) -> PyResult<Value> {
if depth > MAX_DEPTH {
return Err(PyValueError::new_err("value is nested too deeply"));
}
let nested = |items: Bound<'_, PyAny>| -> PyResult<Vec<Value>> {
items
.try_iter()?
.map(|item| to_value(&item?, depth + 1))
.collect()
};
if value.is_none() {
return Ok(Value::None);
}
if let Ok(flag) = value.cast::<PyBool>() {
return Ok(Value::Bool(flag.is_true()));
}
if value.is_instance_of::<PyInt>() {
return Ok(Value::Int(match value.extract::<i64>() {
Ok(small) => small.into(),
Err(_) => value
.str()?
.to_str()?
.parse::<BigInt>()
.map_err(|error| PyValueError::new_err(error.to_string()))?,
}));
}
if let Ok(number) = value.cast::<PyFloat>() {
return Ok(Value::Float(number.value()));
}
if let Ok(text) = value.cast::<PyString>() {
return Ok(Value::Str(text.to_str()?.to_owned()));
}
if let Ok(bytes) = value.cast::<PyBytes>() {
return Ok(Value::Bytes(bytes.as_bytes().to_vec()));
}
if let Ok(complex) = value.cast::<PyComplex>() {
return Ok(Value::Complex {
re: complex.real(),
im: complex.imag(),
});
}
if let Ok(entries) = value.cast::<PyDict>() {
return entries
.iter()
.map(|(key, item)| Ok((to_value(&key, depth + 1)?, to_value(&item, depth + 1)?)))
.collect::<PyResult<_>>()
.map(Value::Dict);
}
if value.is_instance_of::<PyList>() {
return nested(value.clone()).map(Value::List);
}
if value.is_instance_of::<PyTuple>() {
return nested(value.clone()).map(Value::Tuple);
}
if value.is_instance_of::<PySet>() || value.is_instance_of::<PyFrozenSet>() {
return nested(value.clone()).map(Value::Set);
}
Err(PyTypeError::new_err(format!(
"{} is not a plain data value",
value.get_type().name()?
)))
}
fn to_python(py: Python<'_>, value: Value) -> PyResult<Bound<'_, PyAny>> {
let all = |items: Vec<Value>| -> PyResult<Vec<Bound<'_, PyAny>>> {
items.into_iter().map(|item| to_python(py, item)).collect()
};
Ok(match value {
Value::None => py.None().into_bound(py),
Value::Bool(flag) => PyBool::new(py, flag).to_owned().into_any(),
Value::Int(number) => match i64::try_from(&number) {
Ok(small) => small.into_pyobject(py)?.into_any(),
Err(_) => py
.get_type::<PyInt>()
.call1((number.to_string(),))?
.into_any(),
},
Value::Float(number) => number.into_pyobject(py)?.into_any(),
Value::Complex { re, im } => PyComplex::from_doubles(py, re, im).into_any(),
Value::Str(text) => text.into_pyobject(py)?.into_any(),
Value::Bytes(bytes) => PyBytes::new(py, &bytes).into_any(),
Value::Tuple(items) => PyTuple::new(py, all(items)?)?.into_any(),
Value::List(items) => PyList::new(py, all(items)?)?.into_any(),
Value::Set(items) => PySet::new(py, all(items)?)?.into_any(),
Value::Dict(entries) => {
let dict = PyDict::new(py);
for (key, item) in entries {
dict.set_item(to_python(py, key)?, to_python(py, item)?)?;
}
dict.into_any()
}
})
}

4
litellm-rust/crates/python-bridge/src/router/config.rs
View file

@ -188,6 +188,8 @@ struct RequestSpec {
web_search: Option<bool>,
route_handler: Option<String>,
decisions_media: Option<bool>,
batch_create: Option<bool>,
embedding_batch: Option<bool>,
wildcard_request: Option<String>,
unrouted: Option<bool>,
passthrough_on_no_deployment: Option<bool>,
@ -255,6 +257,8 @@ pub(crate) fn parse_request(value: Value, model: Option<&str>) -> Result<RouterR
web_search: spec.web_search.unwrap_or(false),
route_handler: spec.route_handler,
decisions_media: spec.decisions_media.unwrap_or(false),
batch_create: spec.batch_create.unwrap_or(false),
embedding_batch: spec.embedding_batch.unwrap_or(false),
wildcard_request: spec.wildcard_request,
unrouted: spec.unrouted.unwrap_or(false),
passthrough_on_no_deployment: spec.passthrough_on_no_deployment.unwrap_or(false),

148
litellm-rust/crates/python-bridge/src/router/host.rs
View file

@ -20,10 +20,10 @@ use std::{
use async_trait::async_trait;
use litellm_router_core::{
Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, Deployment, Error, Host, Output,
Payload, PreRoute, PreRouted, Provided, Randomness, RouterCache, RouterEvent, RouterRequest,
RouterResponse, RoutingContext, SecretRef, SelectStage, ThreadRandomness,
build::BuiltDeployment, host::ModelInfoLookup, model_info::MergedModelInfo,
Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, CallbackFiltered, Deployment,
Error, Host, Output, Payload, PreRoute, PreRouted, Provided, Randomness, RouterCache,
RouterEvent, RouterRequest, RouterResponse, RoutingContext, SecretRef, SelectStage,
ThreadRandomness, build::BuiltDeployment, host::ModelInfoLookup, model_info::MergedModelInfo,
precall::PreCallFacts,
};
use pyo3::{PyResult, Python, types::PyAnyMethods};
@ -68,7 +68,12 @@ pub(crate) enum Stage {
/// Pre-call checks and sub-routers: return the deployments that remain.
Plugins,
/// A strategy router-core does not implement: return the candidates with the pick first.
Strategy(String),
/// When candidates share an id, `records` holds each candidate's deployment, in order, and the
/// host answers with the pick's position, so a pick splits over every entry as Python's does.
Strategy {
strategy: String,
records: Vec<Value>,
},
/// Tag routing: return the deployments that remain.
Tags,
/// `enable_pre_call_checks`: return the facts the filter decides by.
@ -79,6 +84,10 @@ pub(crate) enum Stage {
RouteIncapable { handler: String },
/// `_decisions_deployment_refuses_media`: return the candidates that refuse the call's media.
MediaRefusing,
/// `deployments_admitting_embedding_batch`: return the candidates the embedding pick keeps.
EmbeddingBatch,
/// The override strategy's pre-call check on a deployment named by id.
OverrideCheck(String),
/// The routing plugins: run them when `run`, then return the candidate models they left.
RoutingPlugins { run: bool },
/// The host's `CustomLogger` deployment filters: return the ids that remain, or a failure.
@ -163,6 +172,8 @@ pub(crate) enum HostCall {
/// The name the attempt resolved: the requested name or the fallback target.
model: String,
model_group: String,
/// The group a pre-routing hook selected for `model`, when it differs.
routed: Option<String>,
/// The `order` level of a same-group order fallback.
target_order: Option<i64>,
/// Every deployment `model` resolved to for this caller, the picked one included.
@ -279,14 +290,17 @@ pub(crate) enum HostCallResponse {
RefuseFallbacks(bool),
/// The candidate ids that remain, pick first.
Select(Vec<String>),
/// A strategy's pick among candidates sharing an id: positions in the candidates, pick first.
Positions(Vec<usize>),
/// The model group to resolve instead, a custom strategy's deployment, or `None`.
PreRoute(Option<PreRouted>),
/// The facts a `pre_call` select step asked for.
PreCall(PreCallFacts),
/// The candidate models a `routing_plugins` select step reported.
Models(Option<Vec<String>>),
/// The ids a `callback_filters` select step left, or the error a filter raised.
Filtered(Result<Vec<String>, AttemptError>),
/// The ids a `callback_filters` select step left with the one a filter prefers, or the error
/// a filter raised.
Filtered(Result<(Vec<String>, Option<String>), AttemptError>),
/// What a pick's select or pre-route step, or a fallback check, raised on the host instead of
/// answering: the walk takes it as the try's failure.
Raised(AttemptError),
@ -309,9 +323,11 @@ impl HostCallResponse {
Self::Slept => StepKind::Sleep,
Self::AllowFallback(_) => StepKind::AllowFallback,
Self::RefuseFallbacks(_) => StepKind::RefuseFallbacks,
Self::Select(_) | Self::PreCall(_) | Self::Models(_) | Self::Filtered(_) => {
StepKind::Select
}
Self::Select(_)
| Self::Positions(_)
| Self::PreCall(_)
| Self::Models(_)
| Self::Filtered(_) => StepKind::Select,
Self::PreRoute(_) => StepKind::PreRoute,
Self::Cache(_) => StepKind::Cache,
Self::Secrets(_) => StepKind::Secrets,
@ -337,12 +353,21 @@ impl HostCallResponse {
got: self.kind(),
});
}
if let (Self::Select(ids) | Self::Filtered(Ok(ids)), HostCall::Select { candidates, .. }) =
(self, call)
if let (
Self::Select(ids) | Self::Filtered(Ok((ids, _))),
HostCall::Select { candidates, .. },
) = (self, call)
&& let Some(id) = ids.iter().find(|id| !candidates.contains(id))
{
return Err(DriveError::UnknownCandidate { id: id.clone() });
}
if let (Self::Positions(positions), HostCall::Select { candidates, .. }) = (self, call)
&& let Some(position) = positions.iter().find(|p| **p >= candidates.len())
{
return Err(DriveError::UnknownCandidate {
id: format!("#{position}"),
});
}
if let (Self::Secrets(values), HostCall::Secrets { refs, .. }) = (self, call)
&& values.len() != refs.len()
{
@ -654,6 +679,28 @@ impl RouterCache for StepCache {
}
/// A deployment as Python spells a `model_list` entry, with its id in `model_info`.
/// Whether two candidates share an id, which a pick answered by id can't tell apart.
fn shares_ids(candidates: &[Arc<Deployment>]) -> bool {
let mut seen = std::collections::HashSet::new();
!candidates.iter().all(|d| seen.insert(d.id.as_str()))
}
/// The candidates a host answer names, in its order. An id names every candidate that carries
/// it, so deployments sharing an id stay together through a filter, as Python's list keeps them.
fn answered_by_id(candidates: &[Arc<Deployment>], ids: &[String]) -> Vec<Arc<Deployment>> {
let mut taken = vec![false; candidates.len()];
let mut kept = Vec::new();
for id in ids {
for (position, deployment) in candidates.iter().enumerate() {
if !taken[position] && deployment.id == *id {
taken[position] = true;
kept.push(deployment.clone());
}
}
}
kept
}
pub(crate) fn deployment_record(deployment: &Deployment) -> Value {
deployment.record()
}
@ -693,6 +740,7 @@ impl Host for StepHost {
deployment: deployment_record(&attempt.deployment),
model: attempt.name.clone(),
model_group: attempt.model_group.clone(),
routed: attempt.routed.clone(),
target_order: attempt.target_order,
candidate_ids: attempt.candidate_ids.clone(),
retry: attempt.retry,
@ -736,13 +784,22 @@ impl Host for StepHost {
let call = HostCall::Select {
stage: match stage {
SelectStage::Plugins => Stage::Plugins,
SelectStage::Strategy(strategy) => Stage::Strategy(strategy.to_string()),
SelectStage::Strategy(strategy) => Stage::Strategy {
strategy: strategy.to_string(),
records: if shares_ids(&candidates) {
candidates.iter().map(|d| deployment_record(d)).collect()
} else {
Vec::new()
},
},
SelectStage::Tags => Stage::Tags,
SelectStage::TextOnly => Stage::TextOnly,
SelectStage::RouteIncapable { handler } => Stage::RouteIncapable {
handler: handler.clone(),
},
SelectStage::MediaRefusing => Stage::MediaRefusing,
SelectStage::EmbeddingBatch => Stage::EmbeddingBatch,
SelectStage::OverrideCheck(strategy) => Stage::OverrideCheck(strategy.to_string()),
},
pool: ctx.pool.iter().map(|d| d.id.clone()).collect(),
model_group: ctx.model_group.to_owned(),
@ -751,13 +808,14 @@ impl Host for StepHost {
group: group.map(str::to_owned),
candidates: candidates.iter().map(|d| d.id.clone()).collect(),
};
let HostCallResponse::Select(ids) = answered(call_host(call).await)? else {
unreachable!("{CHECKED}");
};
Ok(ids
.iter()
.filter_map(|id| candidates.iter().find(|d| d.id == *id).cloned())
.collect())
match answered(call_host(call).await)? {
HostCallResponse::Select(ids) => Ok(answered_by_id(&candidates, &ids)),
HostCallResponse::Positions(positions) => Ok(positions
.iter()
.map(|position| candidates[*position].clone())
.collect()),
_ => unreachable!("{CHECKED}"),
}
}
async fn pre_route(&self, pre: &PreRoute<'_>) -> Result<Option<PreRouted>, Error> {
@ -835,7 +893,7 @@ impl Host for StepHost {
&self,
ctx: &RoutingContext<'_>,
candidates: Vec<Arc<Deployment>>,
) -> Result<Vec<Arc<Deployment>>, Error> {
) -> Result<CallbackFiltered, Error> {
let (scope, group) = ctx.scope.wire();
let joined = |own: &[String], walked: &[String]| [own, walked].concat();
let call = HostCall::Select {
@ -854,11 +912,11 @@ impl Host for StepHost {
let HostCallResponse::Filtered(result) = call_host(call).await else {
unreachable!("{CHECKED}");
};
let ids = result.map_err(Error::Attempt)?;
Ok(ids
.iter()
.filter_map(|id| candidates.iter().find(|d| d.id == *id).cloned())
.collect())
let (ids, preferred) = result.map_err(Error::Attempt)?;
Ok(CallbackFiltered {
kept: answered_by_id(&candidates, &ids),
preferred,
})
}
async fn allow_fallback(
@ -922,6 +980,7 @@ impl Host for StepHost {
| RouterEvent::Blocked { .. }
| RouterEvent::CooldownSet { .. }
| RouterEvent::Selected { .. }
| RouterEvent::PickStarted { .. }
) {
return;
}
@ -1461,7 +1520,7 @@ mod tests {
let mut driver = start(&router, "primary");
let step = store.resume(&mut driver, None);
let Step::HostCall(HostCall::Select {
stage: Stage::Strategy(strategy),
stage: Stage::Strategy { strategy, records },
candidates,
..
}) = step
@ -1470,6 +1529,7 @@ mod tests {
};
assert_eq!(strategy, "least-busy");
assert_eq!(candidates, ["a1", "a2"]);
assert!(records.is_empty(), "distinct ids send no records");
let step = store.resume(
&mut driver,
Some(HostCallResponse::Select(vec!["a2".to_owned()])),
@ -1480,6 +1540,40 @@ mod tests {
);
}
#[rstest]
fn a_pick_among_deployments_sharing_an_id_is_answered_by_position() {
let mut second = deployment("primary", "a1");
second["litellm_params"]["api_base"] = json!("https://second.example");
let router = router(
json!({
"model_list": [deployment("primary", "a1"), second],
"routing_strategy": "least-busy",
}),
json!({}),
);
let mut store = Store::default();
let mut driver = start(&router, "primary");
let step = store.resume(&mut driver, None);
let Step::HostCall(HostCall::Select {
stage: Stage::Strategy { records, .. },
..
}) = step
else {
panic!("a host strategy should ask for a selection, got {step:?}");
};
assert_eq!(
records.len(),
2,
"each candidate sharing an id sends its record"
);
let step = store.resume(&mut driver, Some(HostCallResponse::Positions(vec![1])));
assert!(
matches!(&step, Step::HostCall(HostCall::Attempt { deployment, .. })
if deployment["litellm_params"]["api_base"] == "https://second.example"),
"expected the second entry, got {step:?}"
);
}
#[rstest]
fn a_selection_is_driven_through_cache_steps(pair: Router) {
let request = RouterRequest {

48
litellm-rust/crates/python-bridge/src/router/mod.rs
View file

@ -28,7 +28,7 @@ use litellm_router_core::{
use pyo3::{
exceptions::{PyRuntimeError, PyValueError},
prelude::*,
types::{PyBool, PyDict, PyList, PyString},
types::{PyBool, PyDict, PyList, PyString, PyTuple},
};
use serde_json::Value;
@ -457,7 +457,7 @@ impl RouterCoreHandle {
}
/// Record a success a logging callback saw, driven like `report_failure`; it ends with
/// `done` (`key`: the success counter's key, or None for an unknown deployment).
/// `done` (`key`: the tokens counter's key when usage was added, else None).
/// `usage_model` is the model its usage is keyed by, and `tracks_usage` the host's answer to
/// whether it counts any (see `Output::tracks_usage`). `callbacks`: see `report_failure`.
#[pyo3(signature = (id, total_tokens, requests, usage_model=None, tracks_usage=None, callbacks=None))]
@ -1181,7 +1181,8 @@ fn host_call_response(
stage: Stage::CallbackFilters { .. },
..
} => {
// The remaining ids, or a failure mapping as an attempt reports one.
// The remaining ids, alone or paired with the id a filter prefers, or a failure
// mapping as an attempt reports one.
let outcome = outcome.ok_or_else(missing)?;
if let Ok(failure) = outcome.cast::<PyDict>() {
return match attempt_outcome(failure)? {
@ -1191,8 +1192,17 @@ fn host_call_response(
)),
};
}
Ok(HostCallResponse::Filtered(Ok(outcome.extract()?)))
if let Ok(answer) = outcome.cast::<PyTuple>() {
return Ok(HostCallResponse::Filtered(Ok(answer.extract()?)));
}
Ok(HostCallResponse::Filtered(Ok((outcome.extract()?, None))))
}
HostCall::Select {
stage: Stage::Strategy { records, .. },
..
} if !records.is_empty() => Ok(HostCallResponse::Positions(
outcome.ok_or_else(missing)?.extract()?,
)),
HostCall::Select { .. } => Ok(HostCallResponse::Select(
outcome.ok_or_else(missing)?.extract()?,
)),
@ -1285,6 +1295,15 @@ fn pre_route_outcome(outcome: Option<Bound<'_, PyAny>>) -> PyResult<Option<PreRo
if let Ok(model) = outcome.cast::<PyString>() {
return Ok(Some(PreRouted::Model(model.to_str()?.to_owned())));
}
// `(model, group)`: the model to resolve, served as the group a session binding stamped.
if let Ok(pair) = outcome.cast::<PyTuple>()
&& pair.len() == 2
{
return Ok(Some(PreRouted::Bound {
model: pair.get_item(0)?.extract()?,
group: pair.get_item(1)?.extract()?,
}));
}
let spec: DeploymentSpec = serde_json::from_value(document(&outcome)?)
.map_err(|error| PyValueError::new_err(error.to_string()))?;
let deployment = Deployment::from_spec(spec)
@ -1401,7 +1420,7 @@ fn cause_items(dict: &Bound<'_, PyDict>, cause: RetryCause) -> PyResult<()> {
fn pick_items(dict: &Bound<'_, PyDict>, error: &Error) -> PyResult<()> {
error_items(dict, error)?;
match error {
Error::NoHealthyDeployments { model } | Error::UnknownModel { model } => {
Error::NoHealthyDeployments { model, .. } | Error::UnknownModel { model } => {
dict.set_item("failed_model", model)
}
_ => Ok(()),
@ -1443,6 +1462,10 @@ fn error_items(dict: &Bound<'_, PyDict>, error: &Error) -> PyResult<()> {
dict.set_item("error_model", model)?;
dict.set_item("error_detail", format!("{until} {timezone}"))
}
Error::NoHealthyDeployments {
request_ids: Some(ids),
..
} => dict.set_item("error_request_ids", PyList::new(dict.py(), ids)?),
Error::NonNumericWeight { metric } => dict.set_item("error_detail", metric),
Error::MissingRoutingGroup { group } => dict.set_item("error_model", group),
Error::ModelInfo { exception, detail } => {
@ -1637,6 +1660,10 @@ fn events_list(py: Python<'_>, events: Vec<RouterEvent>) -> PyResult<Bound<'_, P
dict.set_item("deployment_id", deployment_id)?;
dict.set_item("cooldown_time", cooldown.as_secs_f64())?;
}
RouterEvent::PickStarted { model } => {
dict.set_item("kind", "pick_started")?;
dict.set_item("model", model)?;
}
RouterEvent::Selected {
deployment_id,
pick,
@ -1685,6 +1712,7 @@ fn host_call_items(py: Python<'_>, dict: &Bound<'_, PyDict>, call: HostCall) ->
deployment,
model,
model_group,
routed,
target_order,
candidate_ids,
retry,
@ -1696,6 +1724,7 @@ fn host_call_items(py: Python<'_>, dict: &Bound<'_, PyDict>, call: HostCall) ->
max_retries,
model_group_size,
} => {
dict.set_item("routed", routed)?;
dict.set_item("scope", scope)?;
dict.set_item("max_retries", max_retries)?;
dict.set_item("model_group_size", model_group_size)?;
@ -1729,7 +1758,12 @@ fn host_call_items(py: Python<'_>, dict: &Bound<'_, PyDict>, call: HostCall) ->
} => {
let (stage, strategy) = match stage {
Stage::Plugins => ("plugins", None),
Stage::Strategy(strategy) => ("strategy", Some(strategy)),
Stage::Strategy { strategy, records } => {
if !records.is_empty() {
dict.set_item("records", to_py(py, &Value::Array(records))?)?;
}
("strategy", Some(strategy))
}
Stage::Tags => ("tags", None),
Stage::PreCall => ("pre_call", None),
Stage::TextOnly => ("text_only", None),
@ -1738,6 +1772,8 @@ fn host_call_items(py: Python<'_>, dict: &Bound<'_, PyDict>, call: HostCall) ->
("route_incapable", None)
}
Stage::MediaRefusing => ("media_refusing", None),
Stage::EmbeddingBatch => ("embedding_batch", None),
Stage::OverrideCheck(strategy) => ("override_check", Some(strategy)),
Stage::ProviderRequest => ("provider_request", None),
Stage::CallbackFilters {
excluded_ids,

11
litellm-rust/crates/python-compat/benches/formats.rs
View file

@ -10,7 +10,7 @@
use std::{hint::black_box, time::Duration};
use criterion::{
BatchSize, BenchmarkGroup, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main,
BenchmarkGroup, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main,
measurement::WallTime,
};
use litellm_python_compat::{Value, json, literal::literal_eval, pickle, repr::repr};
@ -72,14 +72,9 @@ fn formats(c: &mut Criterion) {
let mut group = c.benchmark_group("cached_completion");
text_formats(&mut group, &text, &value);
binary_formats(&mut group, &value, &pickled);
// `from_json` consumes its input, so each iteration gets a freshly parsed one.
group.throughput(Throughput::Bytes(dumped.len() as u64));
group.bench_function("from_json", |bencher| {
bencher.iter_batched(
|| serde_json::from_str::<serde_json::Value>(&dumped).expect("dumps output parses"),
json::from_json,
BatchSize::SmallInput,
)
group.bench_function("json_loads", |bencher| {
bencher.iter(|| json::loads(black_box(&dumped)))
});
group.finish();
}

2
litellm-rust/crates/python-compat/src/error.rs
View file

@ -6,6 +6,8 @@ use crate::MAX_DEPTH;
pub enum Error {
#[error("malformed Python literal at byte {0}")]
InvalidLiteral(usize),
#[error("malformed JSON at byte {0}")]
InvalidJson(usize),
#[error("unhashable type: '{0}'")]
Unhashable(&'static str),
#[error("value nests deeper than {MAX_DEPTH} levels")]

Some files were not shown because too many files have changed in this diff Show more