Merge remote-tracking branch 'origin/main' into litellm_feat_955c0d
Some checks are pending
Cost map pins recheck / rerun-cost-map-guard (push) Waiting to run
LiteLLM Rust DB / db-queries (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run
Terraform Modules / fmt, validate, test (aws) (push) Waiting to run
Terraform Modules / fmt, validate, test (gcp) (push) Waiting to run
Terraform Provider / gofmt, vet, build, test (push) Waiting to run
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Waiting to run

# Conflicts:
#	litellm/repositories/verification_token_repository.py
This commit is contained in:
Yuneng Jiang 2026-11-06 13:38:49 -08:00
commit b6567f5f71
No known key found for this signature in database
2314 changed files with 232708 additions and 91253 deletions

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>)

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:

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

View file

@ -6,6 +6,8 @@ on:
- main
- "litellm_**"
- "rc/**"
- "stable/**"
- "lts/**"
schedule:
- cron: "23 6 * * *"
workflow_dispatch:

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

View file

@ -5,6 +5,9 @@ on:
branches:
- main
- "litellm_**"
- "rc/**"
- "stable/**"
- "lts/**"
push:
branches:
- main
@ -42,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
@ -95,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
@ -167,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
@ -280,6 +296,7 @@ jobs:
tests/unit/llms/e2b
tests/unit/llms/edenai
tests/unit/llms/elevenlabs
tests/unit/llms/empiriolabs
tests/unit/llms/exa_ai
tests/unit/llms/fal_ai
tests/unit/llms/fastcrw
@ -415,6 +432,7 @@ jobs:
- shard: router
test-path: >-
tests/unit/router_rust
tests/unit/test_router
python-version: "3.12"
workers: 4
@ -973,10 +991,10 @@ jobs:
dist: loadscope
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
timeout-minutes: ${{ matrix.shard == 'db-and-spend' && 10 || 3 }}
timeout-minutes: 3
with:
persist-credentials: false
fetch-depth: ${{ matrix.shard == 'db-and-spend' && '0' || '1' }}
fetch-depth: 1
- name: Detect relevant changes
id: changes
@ -1009,27 +1027,37 @@ jobs:
with:
artifact: ${{ needs.rust-bridge.outputs.artifact }}
- name: Install dependencies and check schema
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 8
timeout-minutes: 5
env:
RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }}
UNIT_SHARD: ${{ matrix.shard }}
SCHEMA_BASE_REF: origin/${{ github.event.repository.default_branch }}
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'))"
fi
uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]'
if [ "$UNIT_SHARD" = "db-and-spend" ]; then
uv run --no-sync python scripts/check_schema_compatibility.py --base-ref "$SCHEMA_BASE_REF"
fi
- name: Check schema.prisma against the merge base
if: steps.changes.outputs.decision != 'skip' && matrix.shard == 'db-and-spend' && github.event_name == 'pull_request'
timeout-minutes: 3
env:
GH_TOKEN: ${{ github.token }}
PR_NUMBER: ${{ github.event.pull_request.number }}
EVENT_BASE_REF: ${{ github.event.pull_request.base.ref }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
run: |
base_ref=$(gh api "repos/${GITHUB_REPOSITORY}/branches/${EVENT_BASE_REF}" --jq .name) || base_ref=$(gh api "repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}" --jq .base.ref)
merge_base=$(gh api "repos/${GITHUB_REPOSITORY}/compare/${base_ref}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha')
test -n "$merge_base"
git fetch --no-tags --depth=1 origin "$merge_base"
uv run --no-sync python scripts/check_schema_compatibility.py --base-commit "$merge_base"
- name: Cache Prisma binaries
if: steps.changes.outputs.decision != 'skip'
@ -1266,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'
@ -1469,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

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

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

View file

@ -398,6 +398,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
| [Deepseek (`deepseek`)](https://docs.litellm.ai/docs/providers/deepseek) | ✅ | ✅ | ✅ | | | | | | | |
| [Eden AI (`edenai`)](https://docs.litellm.ai/docs/providers/edenai) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | | |
| [ElevenLabs (`elevenlabs`)](https://docs.litellm.ai/docs/providers/elevenlabs) | ✅ | ✅ | ✅ | | | ✅ | ✅ | | | |
| [EmpirioLabs AI (`empiriolabs`)](https://docs.litellm.ai/docs/providers/empiriolabs) | ✅ | | ✅ | | | | | | | |
| [Empower (`empower`)](https://docs.litellm.ai/docs/providers/empower) | ✅ | ✅ | ✅ | | | | | | | |
| [eu/jev Decisions API (`eu_jev`)](https://docs.litellm.ai/docs/pass_through/eu_jev) | | | | | | | | | | |
| [Fal AI (`fal_ai`)](https://docs.litellm.ai/docs/providers/fal_ai) | ✅ | ✅ | ✅ | | ✅ | | | | | |

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.

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

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")))

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

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

View file

@ -200,7 +200,6 @@ async def log_billed_background_response(billed: BilledBackgroundResponse) -> Bi
unpriced: Final = _without_read_price(billed.response)
cost: Final = logging_obj.response_cost_calculator(result=unpriced)
end_time: Final = datetime.now()
logging_obj.capture_usage_before_logging(unpriced, logging_obj.start_time, end_time)
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
async_coroutine=logging_obj.async_success_handler(
result=unpriced, start_time=logging_obj.start_time, end_time=end_time

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(

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()

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==",

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.

View file

@ -18,7 +18,7 @@ type: application
# This is the chart version. This version number should be incremented each time you make changes
# to the chart and its templates, including the app version.
# Versions are expected to follow Semantic Versioning (https://semver.org/)
version: 1.1.3
version: 1.1.4
# This is the version number of the application being deployed. This version number should be
# incremented each time you make changes to the application. Versions are not expected to

View file

@ -56,8 +56,8 @@ If `db.useStackgresOperator` is used (not yet implemented):
| `billingMetrics.enabled` | Enable enterprise billable-request metering. Requires an enterprise license. | `false` |
| `billingMetrics.endpoint` | Collector that the billable-request counter is pushed to. | `https://telemetry.litellm.ai` |
| `billingMetrics.secretName` | Name of an existing Secret holding the mTLS client certificate, under the keys `tls.crt` and `tls.key`. | `litellm-billing-metrics-mtls` |
| `billingMetrics.caSecretName` | Name of an existing Secret holding a CA bundle under the key `ca.crt`. Only needed for a private or test collector whose server certificate is not on the public web PKI. | `""` |
| `billingMetrics.tokenUrl` | Cognito OAuth2 client-credentials token endpoint. | `https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token` |
| `billingMetrics.secretName` | Name of an existing Secret with `client-id` and `client-secret` keys. | `litellm-billing-metrics` |
| `billingMetrics.exportIntervalMs` | How often the counter is pushed, in milliseconds. The proxy defaults to `60000` when unset. | `""` |
#### Example `proxy_config` ConfigMap from values (default):
@ -109,10 +109,10 @@ type: Opaque
#### Enterprise billable-request metering
Enterprise licenses meter billable requests by pushing a counter to LiteLLM's collector over mutual TLS. The chart does not create the client certificate; it mounts one you already hold, read-only, so the private key is never exposed through the environment. Create the Secret under the name the chart expects, then turn the block on:
Enterprise licenses meter billable requests by pushing a counter to LiteLLM's collector with Cognito bearer authentication. Create a Secret with `client-id` and `client-secret` keys, then enable the chart values:
```
kubectl create secret tls litellm-billing-metrics-mtls --cert=client.crt --key=client.key
kubectl create secret generic litellm-billing-metrics --from-literal=client-id=... --from-literal=client-secret=...
```
```
@ -120,7 +120,8 @@ billingMetrics:
enabled: true
```
Set `billingMetrics.caSecretName` only when the collector is a private or test one whose server certificate is not on the public web PKI; the production collector needs no CA override. The chart fails the render rather than deploying a proxy that silently never exports, so a missing `secretName` or an emptied `endpoint` surfaces at `helm install` time.
The endpoint and token URL must use HTTPS. The chart requires both URLs and
the Secret name when billing metrics are enabled
### Database Settings

View file

@ -50,53 +50,27 @@ app.kubernetes.io/name: {{ include "litellm.name" . }}
app.kubernetes.io/instance: {{ .Release.Name }}
{{- end }}
{{/*
Enterprise billable-request metering. The client certificate identifies the
deployment to LiteLLM's collector, so it is mounted read-only from an existing
Secret rather than passed through the environment.
*/}}
{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}}
{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}}
{{- define "litellm.billingMetricsEnv" -}}
- name: LITELLM_BILLING_METRICS_ENDPOINT
value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }}
- name: LITELLM_BILLING_METRICS_CLIENT_CERT
value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }}
- name: LITELLM_BILLING_METRICS_CLIENT_KEY
value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }}
{{- if .Values.billingMetrics.caSecretName }}
- name: LITELLM_BILLING_METRICS_CA_CERT
value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }}
{{- end }}
- name: LITELLM_BILLING_METRICS_TOKEN_URL
value: {{ required "billingMetrics.tokenUrl is required when billingMetrics.enabled is true" .Values.billingMetrics.tokenUrl | quote }}
- name: LITELLM_BILLING_METRICS_CLIENT_ID
valueFrom:
secretKeyRef:
name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }}
key: client-id
- name: LITELLM_BILLING_METRICS_CLIENT_SECRET
valueFrom:
secretKeyRef:
name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }}
key: client-secret
{{- with .Values.billingMetrics.exportIntervalMs }}
- name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
value: {{ . | quote }}
{{- end }}
{{- end -}}
{{- define "litellm.billingMetricsVolumes" -}}
- name: billing-metrics-mtls
secret:
secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }}
{{- if .Values.billingMetrics.caSecretName }}
- name: billing-metrics-mtls-ca
secret:
secretName: {{ .Values.billingMetrics.caSecretName }}
{{- end }}
{{- end -}}
{{- define "litellm.billingMetricsVolumeMounts" -}}
- name: billing-metrics-mtls
mountPath: {{ include "litellm.billingMetrics.certDir" . }}
readOnly: true
{{- if .Values.billingMetrics.caSecretName }}
- name: billing-metrics-mtls-ca
mountPath: {{ include "litellm.billingMetrics.caDir" . }}
readOnly: true
{{- end }}
{{- end -}}
{{/*
Create the name of the service account to use
*/}}

View file

@ -148,9 +148,6 @@ spec:
- name: npm
mountPath: /.npm
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
{{- if include "litellm.collector.socketDir" . }}
- name: collector-socket
mountPath: {{ include "litellm.collector.socketDir" . }}
@ -234,9 +231,6 @@ spec:
items:
- key: {{ .Values.proxyConfigMap.key | default "config.yaml" }}
path: "config.yaml"
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}
{{- if include "litellm.collector.socketDir" . }}
- name: collector-socket
emptyDir:

View file

@ -1,31 +1,21 @@
suite: test billingMetrics wiring on the proxy deployment
suite: test bearer-only billingMetrics wiring
templates:
- deployment.yaml
- configmap-litellm.yaml
- migrations-job.yaml
tests:
- it: is off by default, adding no env, volume, or mount
- it: is disabled by default
template: deployment.yaml
asserts:
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: litellm-billing-metrics-mtls
- notContains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls
mountPath: /etc/litellm/billing-mtls
readOnly: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics
- it: renders the endpoint and the mounted cert paths when enabled
- it: renders required URLs and Cognito secret references when enabled
template: deployment.yaml
set:
billingMetrics:
@ -39,119 +29,34 @@ tests:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CLIENT_CERT
value: /etc/litellm/billing-mtls/tls.crt
name: LITELLM_BILLING_METRICS_TOKEN_URL
value: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CLIENT_KEY
value: /etc/litellm/billing-mtls/tls.key
# The conventional Secret name is the default, so enabling the block is enough.
- it: mounts the default cert secret read-only alongside the config volume
template: deployment.yaml
set:
billingMetrics:
enabled: true
asserts:
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: litellm-billing-metrics-mtls
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls
mountPath: /etc/litellm/billing-mtls
readOnly: true
- it: honours a secretName override
template: deployment.yaml
set:
billingMetrics:
enabled: true
secretName: my-billing-mtls
asserts:
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: my-billing-mtls
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: litellm-billing-metrics-mtls
- it: honours an endpoint override
template: deployment.yaml
set:
billingMetrics:
enabled: true
endpoint: https://collector.internal:4318
asserts:
name: LITELLM_BILLING_METRICS_CLIENT_ID
valueFrom:
secretKeyRef:
name: litellm-billing-metrics
key: client-id
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://collector.internal:4318
# The production collector presents a public web-PKI certificate, so the CA
# override must stay absent unless a private collector is configured.
- it: omits the CA env, volume, and mount when no caSecretName is set
template: deployment.yaml
set:
billingMetrics:
enabled: true
asserts:
name: LITELLM_BILLING_METRICS_CLIENT_SECRET
valueFrom:
secretKeyRef:
name: litellm-billing-metrics
key: client-secret
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls-ca
secret:
secretName: billing-ca
name: billing-metrics
- notContains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls-ca
mountPath: /etc/litellm/billing-mtls-ca
readOnly: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CA_CERT
value: /etc/litellm/billing-mtls-ca/ca.crt
name: billing-metrics
- it: mounts the CA secret when caSecretName is set
template: deployment.yaml
set:
billingMetrics:
enabled: true
caSecretName: billing-ca
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CA_CERT
value: /etc/litellm/billing-mtls-ca/ca.crt
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls-ca
secret:
secretName: billing-ca
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls-ca
mountPath: /etc/litellm/billing-mtls-ca
readOnly: true
- it: passes the export interval through only when set
- it: renders the export interval only when set
template: deployment.yaml
set:
billingMetrics:
@ -164,129 +69,7 @@ tests:
name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
value: "5000"
- it: omits the export interval when unset
template: deployment.yaml
set:
billingMetrics:
enabled: true
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
value: "60000"
# Kubernetes resolves duplicate env names last-wins, so the chart-owned billing
# entries must render after .Values.envVars or a user could silently redirect
# the metering export. The three billing entries are the last ones emitted here
# (migrationJob, which appends DISABLE_SCHEMA_UPDATE, is off for this case).
- it: renders the billing endpoint after envVars so it cannot be shadowed
template: deployment.yaml
set:
migrationJob:
enabled: false
billingMetrics:
enabled: true
envVars:
LITELLM_BILLING_METRICS_ENDPOINT: https://shadowed.example
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://shadowed.example
- equal:
path: spec.template.spec.containers[0].env[-3]
value:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- equal:
path: spec.template.spec.containers[0].env[-2].name
value: LITELLM_BILLING_METRICS_CLIENT_CERT
- equal:
path: spec.template.spec.containers[0].env[-1].name
value: LITELLM_BILLING_METRICS_CLIENT_KEY
- it: keeps user-supplied volumes and mounts alongside the billing secret
template: deployment.yaml
set:
billingMetrics:
enabled: true
volumes:
- name: custom-callbacks
configMap:
name: my-callbacks
volumeMounts:
- name: custom-callbacks
mountPath: /app/callbacks
asserts:
- contains:
path: spec.template.spec.volumes
content:
name: custom-callbacks
configMap:
name: my-callbacks
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: litellm-billing-metrics-mtls
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: custom-callbacks
mountPath: /app/callbacks
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls
mountPath: /etc/litellm/billing-mtls
readOnly: true
- it: still mounts the proxy config when enabled
template: deployment.yaml
set:
billingMetrics:
enabled: true
asserts:
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: litellm-config
mountPath: /etc/litellm/config.yaml
subPath: config.yaml
# Only the proxy serves billable traffic. The migrations Job must never mount
# the client certificate, and it renders its own env and volumes, so nothing
# stops a future edit from wiring the billing include into it by mistake.
- it: does not touch the migrations job when enabled
template: migrations-job.yaml
set:
billingMetrics:
enabled: true
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- notExists:
path: spec.template.spec.containers[0].volumeMounts
- notExists:
path: spec.template.spec.volumes
- it: fails loudly when enabled with an emptied secretName
template: deployment.yaml
set:
billingMetrics:
enabled: true
secretName: ""
asserts:
- failedTemplate:
errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)
- it: fails loudly when enabled without an endpoint
- it: requires the endpoint when enabled
template: deployment.yaml
set:
billingMetrics:
@ -295,3 +78,23 @@ tests:
asserts:
- failedTemplate:
errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true
- it: requires the token URL when enabled
template: deployment.yaml
set:
billingMetrics:
enabled: true
tokenUrl: ""
asserts:
- failedTemplate:
errorMessage: billingMetrics.tokenUrl is required when billingMetrics.enabled is true
- it: requires the secret name when enabled
template: deployment.yaml
set:
billingMetrics:
enabled: true
secretName: ""
asserts:
- failedTemplate:
errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true

View file

@ -191,7 +191,7 @@ tests:
metricsServer.port: 9090
billingMetrics.enabled: true
billingMetrics.endpoint: https://metering.example.com
billingMetrics.secretName: billing-mtls
billingMetrics.secretName: billing-metrics
asserts:
- contains:
path: spec.template.spec.containers[0].env
@ -214,9 +214,9 @@ tests:
name: LITELLM_BILLING_METRICS_ENDPOINT
any: true
- notContains:
path: spec.template.spec.containers[1].volumeMounts
path: spec.template.spec.containers[1].env
content:
name: billing-metrics-mtls
name: LITELLM_BILLING_METRICS_CLIENT_SECRET
any: true
- it: should give the sidecar the same scratch mounts as the proxy on a read-only root

View file

@ -154,19 +154,13 @@ masterkeySecretName: ""
# if set, use this secret key for the master key; otherwise, use the default key
masterkeySecretKey: ""
# Optional: enterprise billable-request metering. When enabled, the proxy counts
# successful requests to inference, MCP, and A2A endpoints and pushes them to
# LiteLLM's collector over mutual TLS. Requires an enterprise license.
# The client certificate identifies the deployment, so it is mounted read-only
# from an existing Secret and never passed through the environment.
# Optional: enterprise billable-request metering. Requires an enterprise license.
billingMetrics:
enabled: false
endpoint: https://telemetry.litellm.ai # collector to push the counter to
secretName: litellm-billing-metrics-mtls # existing Secret holding tls.crt and tls.key
# Only for private or test collectors whose server certificate is not on the
# public web PKI. The production collector needs no CA override.
caSecretName: "" # existing Secret holding ca.crt
exportIntervalMs: "" # push cadence; the proxy defaults to 60000
endpoint: https://telemetry.litellm.ai
tokenUrl: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token
secretName: litellm-billing-metrics
exportIntervalMs: ""
proxyConfigMap:
# when true, creates a new configmap

View file

@ -2,7 +2,7 @@ apiVersion: v2
name: litellm
description: LiteLLM componentized — gateway, UI backend, and UI as separate services
type: application
version: 0.1.0
version: 0.1.1
appVersion: "0.1.0"
annotations:
org.opencontainers.image.source: "https://github.com/BerriAI/litellm"

View file

@ -22,6 +22,23 @@ backend:
gatewayUrl: https://gateway.example.com
```
## Enterprise billing metrics
Billing metrics use Cognito bearer authentication. Create the Secret with
`client-id` and `client-secret` keys, then enable billing metrics:
```bash
kubectl create secret generic litellm-billing-metrics --from-literal=client-id=... --from-literal=client-secret=...
```
```yaml
billingMetrics:
enabled: true
```
The chart requires an HTTPS endpoint, an HTTPS token URL, and the Secret name
when billing metrics are enabled
## Routes to another Service
For a shared hostname, `ingress.extraPaths` can route to a Service in the same namespace using `serviceName` and `servicePort`. The port can be a number or a named Service port. Omit `service` when using these fields

View file

@ -58,56 +58,27 @@ app.kubernetes.io/managed-by: {{ .Release.Service }}
helm.sh/chart: {{ printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" }}
{{- end -}}
{{/*
Enterprise billable-request metering. Wired into gateway and backend, not the
migrations job. The gateway serves nearly all billable traffic, but the backend
keeps the named-server MCP transport (/{mcp_server_name}/mcp), which writes a
SpendLogs row, so metering only the gateway would silently drop that traffic.
The client certificate identifies the deployment to LiteLLM's collector, so it is
mounted read-only from an existing Secret rather than passed through the
environment.
*/}}
{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}}
{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}}
{{- define "litellm.billingMetricsEnv" -}}
- name: LITELLM_BILLING_METRICS_ENDPOINT
value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }}
- name: LITELLM_BILLING_METRICS_CLIENT_CERT
value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }}
- name: LITELLM_BILLING_METRICS_CLIENT_KEY
value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }}
{{- if .Values.billingMetrics.caSecretName }}
- name: LITELLM_BILLING_METRICS_CA_CERT
value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }}
{{- end }}
- name: LITELLM_BILLING_METRICS_TOKEN_URL
value: {{ required "billingMetrics.tokenUrl is required when billingMetrics.enabled is true" .Values.billingMetrics.tokenUrl | quote }}
- name: LITELLM_BILLING_METRICS_CLIENT_ID
valueFrom:
secretKeyRef:
name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }}
key: client-id
- name: LITELLM_BILLING_METRICS_CLIENT_SECRET
valueFrom:
secretKeyRef:
name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }}
key: client-secret
{{- with .Values.billingMetrics.exportIntervalMs }}
- name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
value: {{ . | quote }}
{{- end }}
{{- end -}}
{{- define "litellm.billingMetricsVolumes" -}}
- name: billing-metrics-mtls
secret:
secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }}
{{- if .Values.billingMetrics.caSecretName }}
- name: billing-metrics-mtls-ca
secret:
secretName: {{ .Values.billingMetrics.caSecretName }}
{{- end }}
{{- end -}}
{{- define "litellm.billingMetricsVolumeMounts" -}}
- name: billing-metrics-mtls
mountPath: {{ include "litellm.billingMetrics.certDir" . }}
readOnly: true
{{- if .Values.billingMetrics.caSecretName }}
- name: billing-metrics-mtls-ca
mountPath: {{ include "litellm.billingMetrics.caDir" . }}
readOnly: true
{{- end }}
{{- end -}}
{{/*
Per-component selector labels — used in both Service selectors and Deployment matchLabels.

View file

@ -79,16 +79,13 @@ spec:
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- include "litellm.envFrom" .Values.backend | nindent 10 }}
{{- if or .Values.gateway.config.create .Values.backend.volumeMounts .Values.billingMetrics.enabled }}
{{- if or .Values.gateway.config.create .Values.backend.volumeMounts }}
volumeMounts:
{{- if .Values.gateway.config.create }}
- name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
{{- with .Values.backend.volumeMounts }}
{{- toYaml . | nindent 12 }}
{{- end }}
@ -114,16 +111,13 @@ spec:
{{- with .Values.backend.extraContainers }}
{{- tpl (toYaml .) $ | nindent 8 }}
{{- end }}
{{- if or .Values.gateway.config.create .Values.backend.volumes .Values.billingMetrics.enabled }}
{{- if or .Values.gateway.config.create .Values.backend.volumes }}
volumes:
{{- if .Values.gateway.config.create }}
- name: gateway-config
configMap:
name: {{ include "litellm.gateway.fullname" . }}-config
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}
{{- with .Values.backend.volumes }}
{{- toYaml . | nindent 8 }}
{{- end }}

View file

@ -86,7 +86,7 @@ spec:
{{- include "litellm.gateway.collectorEnv" . | nindent 12 }}
{{- end }}
{{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
volumeMounts:
{{- if .Values.gateway.config.create }}
- name: gateway-config
@ -101,9 +101,6 @@ spec:
- name: collector-socket
mountPath: {{ include "litellm.gateway.collectorSocketDir" . }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
{{- with .Values.gateway.volumeMounts }}
{{- toYaml . | nindent 12 }}
{{- end }}
@ -206,7 +203,7 @@ spec:
{{- with .Values.gateway.extraContainers }}
{{- tpl (toYaml .) $ | nindent 8 }}
{{- end }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
volumes:
{{- if .Values.gateway.config.create }}
- name: gateway-config
@ -222,9 +219,6 @@ spec:
emptyDir:
sizeLimit: 1Mi
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}
{{- with .Values.gateway.volumes }}
{{- toYaml . | nindent 8 }}
{{- end }}

View file

@ -1,39 +1,30 @@
suite: test billingMetrics wiring on gateway and backend
suite: test bearer-only billingMetrics wiring
templates:
- gateway/deployment.yaml
- gateway/configmap.yaml
- backend/deployment.yaml
- backend/configmap.yaml
- migrations-job.yaml
values:
- ./values/required.yaml
tests:
- it: is off by default, adding no env, volume, or mount
- it: is disabled by default
template: gateway/deployment.yaml
asserts:
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: billing-mtls
- equal:
path: spec.template.spec.containers[0].volumeMounts
value:
- name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics
- it: renders the endpoint and the mounted cert paths when enabled
- it: wires gateway URLs and bearer credentials from Secret keys
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
asserts:
- contains:
path: spec.template.spec.containers[0].env
@ -43,87 +34,68 @@ tests:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CLIENT_CERT
value: /etc/litellm/billing-mtls/tls.crt
name: LITELLM_BILLING_METRICS_TOKEN_URL
value: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CLIENT_KEY
value: /etc/litellm/billing-mtls/tls.key
- it: mounts the cert secret read-only alongside the config volume
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
asserts:
name: LITELLM_BILLING_METRICS_CLIENT_ID
valueFrom:
secretKeyRef:
name: litellm-billing-metrics
key: client-id
- contains:
path: spec.template.spec.volumes
path: spec.template.spec.containers[0].env
content:
name: billing-metrics-mtls
secret:
secretName: billing-mtls
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls
mountPath: /etc/litellm/billing-mtls
readOnly: true
# The production collector presents a public web-PKI certificate, so the CA
# override must stay absent unless a private collector is configured.
- it: omits the CA env, volume, and mount when no caSecretName is set
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
asserts:
name: LITELLM_BILLING_METRICS_CLIENT_SECRET
valueFrom:
secretKeyRef:
name: litellm-billing-metrics
key: client-secret
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls-ca
secret:
secretName: billing-ca
name: billing-metrics
- notContains:
path: spec.template.spec.containers[0].env
path: spec.template.spec.containers[0].volumeMounts
content:
name: LITELLM_BILLING_METRICS_CA_CERT
value: /etc/litellm/billing-mtls-ca/ca.crt
name: billing-metrics
- it: mounts the CA secret when caSecretName is set
template: gateway/deployment.yaml
- it: wires billing metrics on the backend without volumes
template: backend/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
caSecretName: billing-ca
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CA_CERT
value: /etc/litellm/billing-mtls-ca/ca.crt
- contains:
name: LITELLM_BILLING_METRICS_CLIENT_SECRET
valueFrom:
secretKeyRef:
name: litellm-billing-metrics
key: client-secret
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls-ca
secret:
secretName: billing-ca
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls-ca
mountPath: /etc/litellm/billing-mtls-ca
readOnly: true
name: billing-metrics
- it: passes the export interval through only when set
- it: does not add billing metrics to migrations
template: migrations-job.yaml
set:
billingMetrics:
enabled: true
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
- it: renders an optional export interval
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
exportIntervalMs: 5000
asserts:
- contains:
@ -132,102 +104,27 @@ tests:
name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
value: "5000"
- it: keeps user-supplied gateway volumes alongside the billing secret
- it: requires the endpoint when enabled
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
gateway.volumes:
- name: custom-callbacks
configMap:
name: my-callbacks
gateway.volumeMounts:
- name: custom-callbacks
mountPath: /app/callbacks
endpoint: ""
asserts:
- contains:
path: spec.template.spec.volumes
content:
name: custom-callbacks
configMap:
name: my-callbacks
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: billing-mtls
- failedTemplate:
errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true
# The backend keeps the named-server MCP transport (/{mcp_server_name}/mcp),
# which writes a SpendLogs row, so it must meter too or that traffic is lost.
- it: meters the backend as well, since it serves the MCP transport
template: backend/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls
mountPath: /etc/litellm/billing-mtls
readOnly: true
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: billing-mtls
- it: leaves the backend alone when metering is off
template: backend/deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
# The migrations job runs prisma and serves no traffic; it must never receive
# the client key.
- it: never mounts the billing cert on the migrations job
template: migrations-job.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- isNull:
path: spec.template.spec.volumes
# The conventional Secret name is the default, so enabling metering needs no
# secretName at all; the guard below only fires on an explicitly blanked one.
- it: uses the conventional secret name by default
- it: requires the token URL when enabled
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
tokenUrl: ""
asserts:
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: litellm-billing-metrics-mtls
- failedTemplate:
errorMessage: billingMetrics.tokenUrl is required when billingMetrics.enabled is true
- it: fails loudly when the secretName is explicitly blanked
- it: requires the secret name when enabled
template: gateway/deployment.yaml
set:
billingMetrics:
@ -235,15 +132,4 @@ tests:
secretName: ""
asserts:
- failedTemplate:
errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)
- it: fails loudly when enabled without an endpoint
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
endpoint: ""
secretName: billing-mtls
asserts:
- failedTemplate:
errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true
errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true

View file

@ -91,7 +91,7 @@ tests:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/managed-by: Helm
helm.sh/chart: litellm-0.1.0
helm.sh/chart: litellm-0.1.1
app.kubernetes.io/component: migrations
- it: renders pod-level and container-level securityContext in their own scopes

View file

@ -180,23 +180,14 @@ masterKey:
secretKey: master-key
# Optional: enterprise billable-request metering. When enabled, the gateway and
# backend count successful requests to inference, MCP, and A2A endpoints and push
# them to LiteLLM's collector over mutual TLS. Both components serve billable
# routes: the backend keeps the named-server MCP transport. Requires an
# enterprise license. The client certificate identifies the deployment, so it is
# mounted read-only from an existing Secret and never passed through the env.
# backend count successful requests to inference, MCP, and A2A endpoints. Requires
# an enterprise license.
billingMetrics:
enabled: false
endpoint: https://telemetry.litellm.ai # collector to push the counter to
# An existing Secret holding the client certificate under tls.crt and its key
# under tls.key, usually created from the onboarding artifact. The default is
# the conventional name, so the common path is to create that Secret and set
# enabled: true. Override only if yours is named differently.
secretName: litellm-billing-metrics-mtls
# Only for private or test collectors whose server certificate is not on the
# public web PKI. The production collector needs no CA override.
caSecretName: "" # existing Secret holding ca.crt
exportIntervalMs: "" # push cadence; the proxy defaults to 60000
endpoint: https://telemetry.litellm.ai
tokenUrl: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token
secretName: litellm-billing-metrics
exportIntervalMs: ""
# External Postgres connection.
database:

View file

@ -0,0 +1,19 @@
SET LOCAL lock_timeout = '10s';
CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogDelivery" (
"delivery_id" TEXT PRIMARY KEY DEFAULT gen_random_uuid()::text,
"request_id" TEXT NOT NULL,
"state" "AccountingDeliveryState" NOT NULL DEFAULT 'pending',
"attempts" INTEGER NOT NULL DEFAULT 0,
"next_attempt_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"applied_at" TIMESTAMP(3),
"last_error" TEXT,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"accounting_payload" JSONB NOT NULL,
"is_batch" BOOLEAN NOT NULL
);
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_state_next_attempt_at_idx"
ON "LiteLLM_SpendLogDelivery" ("state", "next_attempt_at");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_state_applied_at_idx"
ON "LiteLLM_SpendLogDelivery" ("state", "applied_at");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_request_id_idx"
ON "LiteLLM_SpendLogDelivery" ("request_id");

View file

@ -0,0 +1,5 @@
SET LOCAL lock_timeout = '10s';
DROP TABLE IF EXISTS "LiteLLM_AccountingDelivery";
DROP TABLE IF EXISTS "LiteLLM_AccountingEvent";
DROP TYPE IF EXISTS "AccountingDeliveryTarget";
DROP TYPE IF EXISTS "AccountingEventKind";

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");

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);

View file

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

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 $$;

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
@ -2071,74 +2130,27 @@ model LiteLLM_TelemetryReport {
@@index([created_at])
}
enum AccountingEventKind {
inbound_usage
outbound_usage
cost
}
enum AccountingDeliveryState {
pending
applied
parked
}
enum AccountingDeliveryTarget {
LiteLLM_VerificationToken
LiteLLM_UserTable
LiteLLM_TeamTable
LiteLLM_TeamMembership
LiteLLM_OrganizationTable
LiteLLM_OrganizationMembership
LiteLLM_EndUserTable
LiteLLM_AgentsTable
LiteLLM_TagTable
LiteLLM_ProjectTable
LiteLLM_ModelAccessGroupBudgetTable
LiteLLM_BudgetTable
LiteLLM_BudgetWindowSpend
LiteLLM_DailyUserSpend
LiteLLM_DailyTeamSpend
LiteLLM_DailyOrganizationSpend
LiteLLM_DailyEndUserSpend
LiteLLM_DailyAgentSpend
LiteLLM_DailyTagSpend
LiteLLM_DailyModelUsage
LiteLLM_DailyToolSpend
LiteLLM_AutoRouterDailySpend
}
model LiteLLM_SpendLogDelivery {
delivery_id String @id @default(dbgenerated("(gen_random_uuid())::text"))
request_id String
state AccountingDeliveryState @default(pending)
accounting_payload Json
is_batch Boolean
attempts Int @default(0)
next_attempt_at DateTime @default(now())
applied_at DateTime?
last_error String?
created_at DateTime @default(now())
// Durable accounting ledger. One row per usage or cost event; payload holds the
// private event envelope versioned by schema_version. Usage rows retain request-scoped
// pricing context; cost rows include resolved rates and the finalized pricing snapshot.
model LiteLLM_AccountingEvent {
event_id String @id
kind AccountingEventKind
source_event_id String?
schema_version Int
payload_hash String
occurred_at DateTime
recorded_at DateTime @default(now())
payload Json
@@index([recorded_at])
@@index([kind, source_event_id])
}
// One row per (event, target table, version): pending work for the projector and,
// once applied in the same transaction as the target write, the receipt that stops a retry.
model LiteLLM_AccountingDelivery {
event_id String
target AccountingDeliveryTarget
version Int
state AccountingDeliveryState @default(pending)
attempts Int @default(0)
next_attempt_at DateTime @default(now())
applied_at DateTime?
last_error String?
@@id([event_id, target, version])
@@index([target, state, next_attempt_at])
@@index([state, next_attempt_at])
@@index([state, applied_at])
@@index([request_id])
}
model LiteLLM_SubscriptionAccountTable {

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==",

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

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

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`

297
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"
@ -3389,15 +3447,19 @@ version = "0.1.0"
dependencies = [
"aws-config",
"aws-credential-types",
"aws-runtime",
"aws-sdk-sts",
"aws-sigv4",
"aws-smithy-runtime-api",
"aws-smithy-types",
"aws-types",
"bytes",
"futures-util",
"litellm-auth-types",
"litellm-error",
"litellm-http",
"litellm-secret-value",
"litellm-secrets",
"moka",
"reqwest",
"rstest",
@ -3406,6 +3468,7 @@ dependencies = [
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"url",
"veil",
]
@ -3418,6 +3481,7 @@ dependencies = [
"litellm-auth-types",
"litellm-error",
"litellm-secret-value",
"litellm-secrets",
"moka",
"rstest",
"serde_json",
@ -3431,6 +3495,7 @@ dependencies = [
name = "litellm-auth-gcp"
version = "0.1.0"
dependencies = [
"axum",
"base64 0.22.1",
"gcp_auth",
"google-cloud-auth",
@ -3440,10 +3505,12 @@ dependencies = [
"litellm-secret-value",
"moka",
"rstest",
"rustls 0.23.42",
"serde",
"serde_json",
"sha2 0.10.9",
"tempfile",
"time",
"tokio",
"veil",
]
@ -3654,8 +3721,7 @@ dependencies = [
"strum",
"tempfile",
"thiserror 2.0.19",
"tracing",
"tracing-subscriber",
"veil",
]
[[package]]
@ -3758,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",
@ -3767,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",
@ -3782,6 +3854,7 @@ dependencies = [
"tower-sessions-moka-store",
"tracing",
"uuid",
"veil",
]
[[package]]
@ -3806,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"
@ -4279,16 +4375,19 @@ dependencies = [
name = "litellm-python-bridge"
version = "0.1.0"
dependencies = [
"async-trait",
"aws-sdk-secretsmanager",
"bytes",
"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",
@ -4303,6 +4402,13 @@ 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",
"litellm-secrets",
@ -4352,23 +4458,108 @@ 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",
]
[[package]]
name = "litellm-router"
version = "0.1.0"
@ -4384,6 +4575,24 @@ dependencies = [
[[package]]
name = "litellm-router-core"
version = "0.1.0"
dependencies = [
"arc-swap",
"async-trait",
"jiff",
"litellm-cache",
"litellm-router-types",
"parking_lot",
"rand 0.8.7",
"rstest",
"serde",
"serde_json",
"serde_with",
"sha2 0.10.9",
"strum",
"thiserror 2.0.19",
"tokio",
"tracing",
]
[[package]]
name = "litellm-router-types"
@ -5233,6 +5442,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"
@ -5343,6 +5562,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"
@ -5350,6 +5585,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7"
dependencies = [
"der",
"pkcs5",
"rand_core 0.6.4",
"spki",
]
@ -5537,6 +5774,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",
@ -6002,7 +6240,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",
@ -6358,6 +6596,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"
@ -6421,6 +6668,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"
@ -7483,6 +7741,7 @@ dependencies = [
"bytes",
"futures-core",
"futures-sink",
"futures-util",
"pin-project-lite",
"tokio",
]
@ -7562,7 +7821,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",
@ -8138,9 +8397,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",
]
@ -8522,7 +8781,7 @@ dependencies = [
"futures",
"http 1.4.2",
"http-body-util",
"hyper 1.10.1",
"hyper 1.12.0",
"hyper-util",
"log",
"once_cell",

View file

@ -13,6 +13,7 @@ litellm-config = { path = "crates/config" }
litellm-error = { path = "crates/error" }
litellm-router = { path = "crates/router" }
litellm-router-types = { path = "crates/router-types" }
litellm-router-core = { path = "crates/router-core" }
litellm-tracing = { path = "crates/tracing" }
litellm-spend-clickhouse = { path = "crates/spend-clickhouse" }
litellm-storage-clickhouse = { path = "crates/storage-clickhouse" }
@ -25,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" }
@ -71,11 +73,17 @@ 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" }
arc-swap = "1.9.2"
async-trait = "0.1.91"
askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] }
tracing = "0.1"
axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "multipart"] }

View file

@ -11,14 +11,16 @@ litellm-error.workspace = true
litellm-http.workspace = true
reqwest.workspace = true
moka = { workspace = true, features = ["sync"] }
moka.workspace = true
litellm-secret-value.workspace = true
serde.workspace = true
serde_json = { workspace = true, features = ["preserve_order"] }
sha2.workspace = true
thiserror.workspace = true
url.workspace = true
veil.workspace = true
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"] }
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio", "credentials-process"] }
aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"] }
aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"] }
aws-sigv4 = "1.5.1"
@ -26,6 +28,9 @@ aws-types = "1.4.0"
aws-smithy-runtime-api = "1.13.0"
[dev-dependencies]
litellm-secrets.workspace = true
aws-smithy-types = "1.6.4"
aws-runtime = { version = "1.9.4", default-features = false }
rstest.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
bytes.workspace = true

File diff suppressed because it is too large Load diff

View file

@ -11,8 +11,8 @@ pub const AWS_SESSION_NAME: &str = AwsParams::SESSION_NAME.env[0];
pub const AWS_PROFILE_NAME: &str = AwsParams::PROFILE_NAME.env[0];
pub const AWS_ROLE_NAME: &str = AwsParams::ROLE_NAME.env[0];
pub const AWS_WEB_IDENTITY_TOKEN: &str = AwsParams::WEB_IDENTITY_TOKEN.env[0];
pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN";
pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
pub const AWS_ROLE_ARN: &str = AwsParams::ROLE_ARN_ENV;
pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = AwsParams::WEB_IDENTITY_TOKEN_FILE_ENV;
pub const AWS_STS_ENDPOINT: &str = AwsParams::STS_ENDPOINT.env[0];
pub const AWS_EXTERNAL_ID: &str = AwsParams::EXTERNAL_ID.env[0];
pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK";

View file

@ -2,6 +2,8 @@ use litellm_error::Source;
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
pub enum Error {
#[error("AWS credential provider returned expired credentials")]
ExpiredCredentials,
#[error("could not load AWS {flow} credentials")]
Credentials {
flow: CredentialFlow,

View file

@ -75,7 +75,8 @@ impl RequestSigner for SigV4Signer {
)
.map_err(|error| match error {
Error::NonAsciiHeader { .. } => litellm_http::Error::unsignable_header(error),
Error::Credentials { .. }
Error::ExpiredCredentials
| Error::Credentials { .. }
| Error::MissingWebIdentityCredentials
| Error::WebIdentityExpiration(_)
| Error::Signing { .. } => litellm_http::Error::signature(error),

View file

@ -8,17 +8,19 @@ repository.workspace = true
[dependencies]
litellm-secret-value.workspace = true
litellm-auth-types.workspace = true
litellm-error.workspace = true
moka.workspace = true
serde_json.workspace = true
sha2.workspace = true
strum.workspace = true
tokio.workspace = true
url.workspace = true
azure_core = "1.0.0"
azure_identity = { version = "1.0.0", features = ["tokio"] }
[dev-dependencies]
litellm-error.workspace = true
litellm-secrets.workspace = true
rstest.workspace = true
tokio.workspace = true
tokio = { workspace = true, features = ["test-util"] }

View file

@ -3,5 +3,5 @@ mod native;
mod resolve;
mod types;
pub use resolve::{AzureAuthService, SECRET_NAMES};
pub use resolve::{AzureAuthService, AzureIdentity, SECRET_NAMES};
pub use types::{AzureAuthInputs, ConfigValue};

View file

@ -59,6 +59,7 @@ pub(crate) enum NativeAzureRequest {
pub(crate) struct ValidatedAzureRequest {
request: NativeAzureRequest,
credential_source: InputSource,
probe_imds: bool,
}
impl ValidatedAzureRequest {
@ -68,9 +69,18 @@ impl ValidatedAzureRequest {
Ok(Self {
request,
credential_source,
probe_imds: false,
})
}
pub(crate) fn with_imds_probe(self, enabled: bool) -> Self {
Self {
probe_imds: enabled
&& matches!(self.request, NativeAzureRequest::ManagedIdentity { .. }),
..self
}
}
pub(crate) fn credential_source(&self) -> InputSource {
self.credential_source
}
@ -87,9 +97,34 @@ impl ValidatedAzureRequest {
}
}
pub(crate) enum CredentialAttempt {
Resolved(ResolvedCredential),
Unavailable(Error),
Failed(Error),
}
impl CredentialAttempt {
pub(crate) fn into_result(self) -> Result<ResolvedCredential, Error> {
match self {
Self::Resolved(credential) => Ok(credential),
Self::Unavailable(error) | Self::Failed(error) => Err(error),
}
}
}
pub(crate) struct NativeAzureTokenAcquirer {
cache: AzureCredentialProviderCache,
transport: Option<azure_core::http::Transport>,
imds_available: tokio::sync::Mutex<ImdsAvailability>,
}
enum ImdsAvailability {
Unknown,
Available,
Unavailable {
retry_at: tokio::time::Instant,
error: Error,
},
}
impl Default for NativeAzureTokenAcquirer {
@ -103,6 +138,7 @@ impl NativeAzureTokenAcquirer {
Self {
cache: AzureCredentialProviderCache::new(cache_capacity),
transport: None,
imds_available: tokio::sync::Mutex::new(ImdsAvailability::Unknown),
}
}
@ -114,10 +150,51 @@ impl NativeAzureTokenAcquirer {
Self {
cache: AzureCredentialProviderCache::new(cache_capacity),
transport: Some(transport),
imds_available: tokio::sync::Mutex::new(ImdsAvailability::Unknown),
}
}
pub(crate) async fn acquire(
pub(crate) async fn acquire(&self, request: ValidatedAzureRequest) -> CredentialAttempt {
let discover_imds = request.probe_imds;
if discover_imds && let Err(error) = self.discover_imds().await {
return CredentialAttempt::Unavailable(error);
}
match self.acquire_token(request).await {
Ok(credential) => CredentialAttempt::Resolved(credential),
Err(error) => {
let missing_identity = discover_imds && unassigned_identity(&error);
if missing_identity {
CredentialAttempt::Unavailable(error)
} else {
CredentialAttempt::Failed(error)
}
}
}
}
async fn discover_imds(&self) -> Result<(), Error> {
let mut availability = self.imds_available.lock().await;
match &*availability {
ImdsAvailability::Available => return Ok(()),
ImdsAvailability::Unavailable { retry_at, error }
if tokio::time::Instant::now() < *retry_at =>
{
return Err(error.clone());
}
ImdsAvailability::Unknown | ImdsAvailability::Unavailable { .. } => {}
}
let result = probe_imds(self.transport.clone().unwrap_or_default()).await;
*availability = match &result {
Ok(()) => ImdsAvailability::Available,
Err(error) => ImdsAvailability::Unavailable {
retry_at: tokio::time::Instant::now() + Duration::from_secs(30),
error: error.clone(),
},
};
result
}
async fn acquire_token(
&self,
request: ValidatedAzureRequest,
) -> Result<ResolvedCredential, Error> {
@ -148,6 +225,54 @@ impl NativeAzureTokenAcquirer {
}
}
fn unassigned_identity(error: &(dyn std::error::Error + 'static)) -> bool {
std::iter::successors(Some(error), |cause| cause.source()).any(|cause| {
cause
.downcast_ref::<azure_core::Error>()
.or_else(|| {
cause
.downcast_ref::<litellm_error::Source>()?
.downcast_ref::<azure_core::Error>()
})
.is_some_and(|error| {
error.http_status() == Some(azure_core::http::StatusCode::BadRequest)
})
})
}
async fn probe_imds(transport: azure_core::http::Transport) -> Result<(), Error> {
use azure_core::http::{Context, Method, Request, Url};
let endpoint =
Url::parse("http://169.254.169.254/metadata/identity/oauth2/token?api-version=2019-08-01")
.map_err(|error| {
Error::InvalidConfiguration(ErrorDetail::failed("Azure IMDS probe endpoint", error))
})?;
let mut request = Request::new(endpoint, Method::Get);
let response = tokio::time::timeout(
Duration::from_secs(1),
transport.send(&Context::new(), &mut request),
)
.await
.map_err(|error| {
Error::CredentialAcquisition(ErrorDetail::failed("Azure IMDS availability probe", error))
})?
.map_err(|error| {
Error::CredentialAcquisition(ErrorDetail::failed("Azure IMDS availability probe", error))
})?;
// Headerless IMDS discovery expects 400, checked 2026-11-01
// https://learn.microsoft.com/en-us/azure/virtual-machines/instance-metadata-service#errors-and-debugging
if response.status() != azure_core::http::StatusCode::BadRequest {
return Err(Error::CredentialAcquisition(
ErrorDetail::UnexpectedStatus {
operation: "Azure IMDS availability probe",
status: u16::from(response.status()),
},
));
}
Ok(())
}
impl NativeAzureRequest {
fn scope(&self) -> &str {
match self {
@ -493,6 +618,7 @@ impl ClientAssertion for StaticAssertion {
#[cfg(test)]
mod tests {
use std::sync::{Arc, Mutex};
use std::time::Duration;
use azure_core::http::headers::Headers;
use azure_core::http::{AsyncRawResponse, HttpClient, Request, StatusCode, Transport};
@ -569,6 +695,296 @@ mod tests {
}
}
#[derive(Debug)]
struct AvailableImdsClient {
probe_status: StatusCode,
probe_delay: Duration,
recover_after: Option<usize>,
response_status: Option<StatusCode>,
probes: std::sync::atomic::AtomicUsize,
}
impl Default for AvailableImdsClient {
fn default() -> Self {
Self {
probe_status: StatusCode::BadRequest,
probe_delay: Duration::ZERO,
recover_after: None,
response_status: None,
probes: std::sync::atomic::AtomicUsize::new(0),
}
}
}
impl HttpClient for AvailableImdsClient {
fn execute_request<'life0, 'life1, 'async_trait>(
&'life0 self,
request: &'life1 Request,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<AsyncRawResponse>> + Send + 'async_trait>,
>
where
'life0: 'async_trait,
'life1: 'async_trait,
Self: 'async_trait,
{
Box::pin(async move {
if request
.headers()
.get_optional_str(&azure_core::http::headers::HeaderName::from_static(
"metadata",
))
.is_none()
{
let probe = self
.probes
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if self.recover_after.is_some_and(|failures| probe >= failures) {
return Ok(AsyncRawResponse::from_bytes(
StatusCode::BadRequest,
Headers::new(),
"metadata header required",
));
}
tokio::time::sleep(self.probe_delay).await;
return Ok(AsyncRawResponse::from_bytes(
self.probe_status,
Headers::new(),
"metadata header required",
));
}
if let Some(status) = self.response_status {
return Ok(AsyncRawResponse::from_bytes(
status,
Headers::new(),
"identity endpoint failure",
));
}
tokio::time::sleep(std::time::Duration::from_secs(2)).await;
let expires = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs()
+ 3600;
Ok(AsyncRawResponse::from_bytes(
StatusCode::Ok,
Headers::new(),
serde_json::json!({
"access_token": "managed-token", "expires_on": expires.to_string(),
"token_type": "Bearer", "resource": "https://service.test",
})
.to_string(),
))
})
}
}
#[rstest::rstest]
#[tokio::test(start_paused = true)]
async fn available_imds_retains_full_sdk_token_acquisition() {
let acquirer = NativeAzureTokenAcquirer::with_transport(
4,
Transport::new(Arc::new(AvailableImdsClient {
probe_status: StatusCode::BadRequest,
response_status: None,
..Default::default()
})),
);
let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
client_id: None,
scope: deployment("https://service.test/.default".into()),
selection_source: InputSource::Deployment,
})
.unwrap()
.with_imds_probe(true);
let credential = acquirer.acquire(request).await.into_result().unwrap();
assert_eq!(credential.secret().expose(), "managed-token");
}
#[rstest::rstest]
#[case::another_metadata_service(StatusCode::NotFound)]
#[case::token_required(StatusCode::Unauthorized)]
#[case::proxy_rejection(StatusCode::Forbidden)]
#[case::unavailable(StatusCode::ServiceUnavailable)]
#[tokio::test(start_paused = true)]
async fn unrelated_metadata_responses_do_not_start_managed_identity_acquisition(
#[case] status: StatusCode,
) {
let client = Arc::new(AvailableImdsClient {
probe_status: status,
..Default::default()
});
let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone()));
let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
client_id: None,
scope: deployment("https://service.test/.default".into()),
selection_source: InputSource::Deployment,
})
.unwrap()
.with_imds_probe(true);
assert!(matches!(
acquirer.acquire(request.clone()).await,
super::CredentialAttempt::Unavailable(_)
));
assert!(matches!(
acquirer.acquire(request).await,
super::CredentialAttempt::Unavailable(_)
));
assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1);
}
#[rstest::rstest]
#[tokio::test(start_paused = true)]
async fn an_imds_timeout_is_not_repeated_on_the_next_request() {
let client = Arc::new(AvailableImdsClient {
probe_delay: Duration::from_secs(10),
..Default::default()
});
let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone()));
let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
client_id: None,
scope: deployment("https://service.test/.default".into()),
selection_source: InputSource::Deployment,
})
.unwrap()
.with_imds_probe(true);
assert!(matches!(
acquirer.acquire(request.clone()).await,
super::CredentialAttempt::Unavailable(_)
));
let second = tokio::time::timeout(Duration::from_millis(10), acquirer.acquire(request))
.await
.unwrap();
assert!(matches!(second, super::CredentialAttempt::Unavailable(_)));
assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1);
}
#[rstest::fixture]
fn imds_request() -> ValidatedAzureRequest {
ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
client_id: None,
scope: deployment("https://service.test/.default".into()),
selection_source: InputSource::Deployment,
})
.unwrap()
.with_imds_probe(true)
}
#[rstest::rstest]
#[case::timeout(StatusCode::BadRequest, Duration::from_secs(10))]
#[case::service_unavailable(StatusCode::ServiceUnavailable, Duration::ZERO)]
#[tokio::test(start_paused = true)]
async fn failed_imds_discovery_recovers_without_restarting_the_acquirer(
imds_request: ValidatedAzureRequest,
#[case] status: StatusCode,
#[case] delay: Duration,
) {
let client = Arc::new(AvailableImdsClient {
probe_status: status,
probe_delay: delay,
recover_after: Some(1),
..Default::default()
});
let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone()));
let (first, second) = tokio::join!(
acquirer.acquire(imds_request.clone()),
acquirer.acquire(imds_request.clone()),
);
assert!(matches!(first, super::CredentialAttempt::Unavailable(_)));
assert!(matches!(second, super::CredentialAttempt::Unavailable(_)));
assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1);
tokio::time::advance(Duration::from_secs(30)).await;
let (first, second) = tokio::join!(
acquirer.acquire(imds_request.clone()),
acquirer.acquire(imds_request.clone()),
);
assert_eq!(
first.into_result().unwrap().secret().expose(),
"managed-token"
);
assert_eq!(
second.into_result().unwrap().secret().expose(),
"managed-token"
);
assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2);
tokio::time::advance(Duration::from_secs(60)).await;
assert_eq!(
acquirer
.acquire(imds_request)
.await
.into_result()
.unwrap()
.secret()
.expose(),
"managed-token",
);
assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2);
}
#[rstest::rstest]
#[tokio::test(start_paused = true)]
async fn a_persistent_imds_outage_reuses_each_failed_probe(
imds_request: ValidatedAzureRequest,
) {
let client = Arc::new(AvailableImdsClient {
probe_delay: Duration::from_secs(10),
..Default::default()
});
let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone()));
assert!(matches!(
acquirer.acquire(imds_request.clone()).await,
super::CredentialAttempt::Unavailable(_),
));
tokio::time::advance(Duration::from_secs(30)).await;
let (first, second) = tokio::join!(
acquirer.acquire(imds_request.clone()),
acquirer.acquire(imds_request),
);
assert!(matches!(first, super::CredentialAttempt::Unavailable(_)));
assert!(matches!(second, super::CredentialAttempt::Unavailable(_)));
assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2);
}
#[rstest::rstest]
#[case::default_identity_not_assigned(true, StatusCode::BadRequest, true)]
#[case::explicit_identity_not_assigned(false, StatusCode::BadRequest, false)]
#[case::default_identity_forbidden(true, StatusCode::Forbidden, false)]
#[case::default_identity_service_failure(true, StatusCode::InternalServerError, false)]
#[tokio::test(start_paused = true)]
async fn only_unavailable_default_identities_allow_fallback(
#[case] discovery: bool,
#[case] status: StatusCode,
#[case] unavailable: bool,
) {
let acquirer = NativeAzureTokenAcquirer::with_transport(
4,
Transport::new(Arc::new(AvailableImdsClient {
probe_status: StatusCode::BadRequest,
response_status: Some(status),
..Default::default()
})),
);
let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
client_id: None,
scope: deployment("https://service.test/.default".into()),
selection_source: InputSource::Deployment,
})
.unwrap()
.with_imds_probe(discovery);
let outcome = acquirer.acquire(request).await;
let is_unavailable = matches!(outcome, super::CredentialAttempt::Unavailable(_));
let error = outcome.into_result().unwrap_err();
assert_eq!(
is_unavailable,
unavailable,
"{}",
litellm_error::Report::new(&error)
);
}
#[rstest::rstest]
#[tokio::test]
async fn client_secret_uses_sdk_protocol_and_reuses_cached_credential() {
let transport = Arc::new(RecordingTokenClient::default());
@ -582,8 +998,12 @@ mod tests {
"https://login.test",
);
let first = acquirer.acquire(request.clone()).await.unwrap();
let second = acquirer.acquire(request).await.unwrap();
let first = acquirer
.acquire(request.clone())
.await
.into_result()
.unwrap();
let second = acquirer.acquire(request).await.into_result().unwrap();
assert_eq!(first.secret().expose(), "native-token");
assert_eq!(second.secret().expose(), "native-token");
@ -599,6 +1019,7 @@ mod tests {
);
}
#[rstest::rstest]
#[tokio::test]
async fn credential_provider_cache_isolates_every_client_secret_identity_field() {
let transport = Arc::new(RecordingTokenClient::default());
@ -645,10 +1066,10 @@ mod tests {
),
];
acquirer.acquire(base.clone()).await.unwrap();
acquirer.acquire(base).await.unwrap();
acquirer.acquire(base.clone()).await.into_result().unwrap();
acquirer.acquire(base).await.into_result().unwrap();
for request in variants.into_iter().skip(1) {
acquirer.acquire(request).await.unwrap();
acquirer.acquire(request).await.into_result().unwrap();
}
assert_eq!(transport.requests.lock().unwrap().len(), 6);

View file

@ -1,3 +1,4 @@
use crate::native::CredentialAttempt;
use litellm_auth_types::{
CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, Sourced,
TokenProviderHandle,
@ -12,25 +13,31 @@ use std::sync::Arc;
use super::native::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest};
use super::types::{AzureAuthInputs, AzureCredentialType, ConfigValue, DEFAULT_AZURE_SCOPE};
const AZURE_AD_TOKEN_ENV: &str = "AZURE_AD_TOKEN";
const AZURE_TENANT_ID_ENV: &str = "AZURE_TENANT_ID";
const AZURE_CLIENT_ID_ENV: &str = "AZURE_CLIENT_ID";
const AZURE_CLIENT_SECRET_ENV: &str = "AZURE_CLIENT_SECRET";
const AZURE_SCOPE_ENV: &str = "AZURE_SCOPE";
const AZURE_AUTHORITY_HOST_ENV: &str = "AZURE_AUTHORITY_HOST";
const AZURE_CREDENTIAL_ENV: &str = "AZURE_CREDENTIAL";
const AZURE_FEDERATED_TOKEN_FILE_ENV: &str = "AZURE_FEDERATED_TOKEN_FILE";
#[derive(strum::IntoStaticStr, strum::VariantArray, strum::VariantNames)]
enum AzureSecretName {
#[strum(serialize = "AZURE_AD_TOKEN")]
AdToken,
#[strum(serialize = "AZURE_TENANT_ID")]
TenantId,
#[strum(serialize = "AZURE_CLIENT_ID")]
ClientId,
#[strum(serialize = "AZURE_CLIENT_SECRET")]
ClientSecret,
#[strum(serialize = "AZURE_SCOPE")]
Scope,
#[strum(serialize = "AZURE_AUTHORITY_HOST")]
AuthorityHost,
#[strum(serialize = "AZURE_CREDENTIAL")]
Credential,
#[strum(serialize = "AZURE_FEDERATED_TOKEN_FILE")]
FederatedTokenFile,
#[strum(serialize = "IDENTITY_ENDPOINT")]
IdentityEndpoint,
#[strum(serialize = "MSI_ENDPOINT")]
MsiEndpoint,
}
pub const SECRET_NAMES: &[&str] = &[
AZURE_AD_TOKEN_ENV,
AZURE_TENANT_ID_ENV,
AZURE_CLIENT_ID_ENV,
AZURE_CLIENT_SECRET_ENV,
AZURE_SCOPE_ENV,
AZURE_AUTHORITY_HOST_ENV,
AZURE_CREDENTIAL_ENV,
AZURE_FEDERATED_TOKEN_FILE_ENV,
];
pub const SECRET_NAMES: &[&str] = <AzureSecretName as strum::VariantNames>::VARIANTS;
#[derive(Clone, Debug)]
pub(crate) enum AzureCredentialPlan {
@ -57,14 +64,14 @@ trait AzureTokenAcquirer: Send + Sync {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>>;
) -> Pin<Box<dyn Future<Output = CredentialAttempt> + Send + '_>>;
}
impl AzureTokenAcquirer for NativeAzureTokenAcquirer {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>> {
) -> Pin<Box<dyn Future<Output = CredentialAttempt> + Send + '_>> {
Box::pin(NativeAzureTokenAcquirer::acquire(self, request))
}
}
@ -83,12 +90,36 @@ impl AzureAuthService {
Self { native }
}
pub async fn identity_token(
&self,
identity: &AzureIdentity,
scope: &str,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Option<ResolvedCredential>, Error> {
let plan = identity_plan(identity, scope, env_lookup)?;
let inputs = AzureAuthInputs::default_credential_for_scope(scope);
Ok(self
.run_plan(plan, &inputs, env_lookup)
.await?
.map(Sourced::into_value))
}
pub async fn get_azure_ad_token(
&self,
inputs: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Option<Sourced<ResolvedCredential>>, Error> {
match select_auth_plan(inputs, env_lookup)? {
self.run_plan(select_auth_plan(inputs, env_lookup)?, inputs, env_lookup)
.await
}
async fn run_plan(
&self,
plan: AzureCredentialPlan,
inputs: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Option<Sourced<ResolvedCredential>>, Error> {
match plan {
AzureCredentialPlan::Supplied(credential) => Ok(Some(credential)),
AzureCredentialPlan::Caller(caller) => {
let credential = caller.acquire().await?;
@ -123,6 +154,7 @@ impl AzureAuthService {
self.native
.acquire(request)
.await
.into_result()
.map(|credential| Sourced::new(credential, source))
.map(Some)
}
@ -131,6 +163,7 @@ impl AzureAuthService {
self.native
.acquire(request)
.await
.into_result()
.map(|credential| Some(Sourced::new(credential, source)))
}
AzureCredentialPlan::Chain(requests) => {
@ -138,8 +171,16 @@ impl AzureAuthService {
for request in requests {
let source = request.credential_source();
match self.native.acquire(request).await {
Ok(credential) => return Ok(Some(Sourced::new(credential, source))),
Err(error) => failures.push(error),
CredentialAttempt::Resolved(credential) => {
return Ok(Some(Sourced::new(credential, source)));
}
CredentialAttempt::Unavailable(error) => failures.push(error),
CredentialAttempt::Failed(error) => {
return Err(Error::CredentialChain {
last: Box::new(error),
earlier: failures,
});
}
}
}
match failures.pop() {
@ -155,39 +196,161 @@ impl AzureAuthService {
}
}
/// The identity azure-identity builds for one explicit credential class, without LiteLLM's provider settings
/// (`AZURE_AD_TOKEN`, `AZURE_CREDENTIAL`, `AZURE_SCOPE`) that `get_azure_ad_token` also reads.
#[derive(Clone, Debug)]
pub enum AzureIdentity {
ClientSecret {
tenant_id: String,
client_id: String,
client_secret: SecretValue,
},
ManagedIdentity {
client_id: String,
},
Default,
}
fn identity_plan(
identity: &AzureIdentity,
scope: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<AzureCredentialPlan, Error> {
let deployment = |value: &str| Sourced::new(value.to_owned(), InputSource::Deployment);
let environment = |name: &str| {
env_lookup(name)
.filter(|value| !value.trim().is_empty())
.map(|value| Sourced::new(value, InputSource::Environment))
};
let scope = deployment(scope);
let authority = environment(AzureSecretName::AuthorityHost.into());
match identity {
AzureIdentity::ClientSecret {
tenant_id,
client_id,
client_secret,
} => Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new(
NativeAzureRequest::ClientSecret {
tenant_id: deployment(tenant_id),
client_id: deployment(client_id),
client_secret: Sourced::new(client_secret.clone(), InputSource::Deployment),
scope,
authority,
},
)?)),
AzureIdentity::ManagedIdentity { client_id } => Ok(AzureCredentialPlan::Native(
ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
client_id: Some(deployment(client_id)),
scope,
selection_source: InputSource::Deployment,
})?,
)),
AzureIdentity::Default => {
let tenant_id = environment(AzureSecretName::TenantId.into());
let client_id = environment(AzureSecretName::ClientId.into());
let from_environment = match (
tenant_id.clone(),
client_id.clone(),
environment(AzureSecretName::ClientSecret.into()),
) {
(Some(tenant_id), Some(client_id), Some(client_secret)) => {
Some(NativeAzureRequest::ClientSecret {
tenant_id,
client_id,
client_secret: client_secret.map(SecretValue::new),
scope: scope.clone(),
authority: authority.clone(),
})
}
_ => None,
};
let probe_imds = env_lookup(AzureSecretName::IdentityEndpoint.into()).is_none()
&& env_lookup(AzureSecretName::MsiEndpoint.into()).is_none();
let AzureCredentialPlan::Chain(chain) = select_native_plan(
Some(Sourced::new(
AzureCredentialType::DefaultAzureCredential,
InputSource::Deployment,
)),
tenant_id,
client_id,
environment(AzureSecretName::FederatedTokenFile.into()),
scope,
authority,
InputSource::Environment,
)?
else {
unreachable!("DefaultAzureCredential always plans a chain")
};
let environment_request = from_environment
.map(ValidatedAzureRequest::new)
.transpose()?;
Ok(AzureCredentialPlan::Chain(
environment_request
.into_iter()
.chain(chain)
.map(|request| request.with_imds_probe(probe_imds))
.collect(),
))
}
}
}
pub(crate) fn select_auth_plan(
inputs: &AzureAuthInputs,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<AzureCredentialPlan, Error> {
let token = configured_secret(&inputs.azure_ad_token, AZURE_AD_TOKEN_ENV, env_lookup);
let tenant_id = configured_string(&inputs.tenant_id, AZURE_TENANT_ID_ENV, env_lookup);
let client_id = configured_string(&inputs.client_id, AZURE_CLIENT_ID_ENV, env_lookup);
let client_secret =
configured_secret(&inputs.client_secret, AZURE_CLIENT_SECRET_ENV, env_lookup);
let scope = configured_string(&inputs.azure_scope, AZURE_SCOPE_ENV, env_lookup)
.unwrap_or_else(|| Sourced::new(DEFAULT_AZURE_SCOPE.to_string(), InputSource::Environment));
let authority = configured_string(
&inputs.azure_authority_host,
AZURE_AUTHORITY_HOST_ENV,
let token = configured_secret(
&inputs.azure_ad_token,
AzureSecretName::AdToken.into(),
env_lookup,
);
let selector = configured_string(&inputs.azure_credential, AZURE_CREDENTIAL_ENV, env_lookup)
.map(|value| {
value
.value()
.parse::<AzureCredentialType>()
.map(|selector| Sourced::new(selector, value.source()))
})
.transpose()
.map_err(|error| {
Error::InvalidConfiguration(ErrorDetail::failed(
"Azure credential selector parsing",
error,
))
})?;
let tenant_id = configured_string(
&inputs.tenant_id,
AzureSecretName::TenantId.into(),
env_lookup,
);
let client_id = configured_string(
&inputs.client_id,
AzureSecretName::ClientId.into(),
env_lookup,
);
let client_secret = configured_secret(
&inputs.client_secret,
AzureSecretName::ClientSecret.into(),
env_lookup,
);
let scope = configured_string(
&inputs.azure_scope,
AzureSecretName::Scope.into(),
env_lookup,
)
.unwrap_or_else(|| Sourced::new(DEFAULT_AZURE_SCOPE.to_string(), InputSource::Environment));
let authority = configured_string(
&inputs.azure_authority_host,
AzureSecretName::AuthorityHost.into(),
env_lookup,
);
let selector = configured_string(
&inputs.azure_credential,
AzureSecretName::Credential.into(),
env_lookup,
)
.map(|value| {
value
.value()
.parse::<AzureCredentialType>()
.map(|selector| Sourced::new(selector, value.source()))
})
.transpose()
.map_err(|error| {
Error::InvalidConfiguration(ErrorDetail::failed(
"Azure credential selector parsing",
error,
))
})?;
let federated_token_file = configured_string(
&inputs.federated_token_file,
AZURE_FEDERATED_TOKEN_FILE_ENV,
AzureSecretName::FederatedTokenFile.into(),
env_lookup,
);
@ -237,7 +400,12 @@ pub(crate) fn select_auth_plan(
return Ok(AzureCredentialPlan::Missing);
}
select_native_plan(
let probe_imds = selector
.as_ref()
.is_none_or(|selector| *selector.value() == AzureCredentialType::DefaultAzureCredential)
&& env_lookup(AzureSecretName::IdentityEndpoint.into()).is_none()
&& env_lookup(AzureSecretName::MsiEndpoint.into()).is_none();
match select_native_plan(
selector,
tenant_id,
client_id,
@ -245,7 +413,15 @@ pub(crate) fn select_auth_plan(
scope,
authority,
inputs.enable_azure_ad_token_refresh.source(),
)
)? {
AzureCredentialPlan::Chain(requests) => Ok(AzureCredentialPlan::Chain(
requests
.into_iter()
.map(|request| request.with_imds_probe(probe_imds))
.collect(),
)),
plan => Ok(plan),
}
}
fn select_native_plan(
@ -509,7 +685,10 @@ mod tests {
#[derive(Debug)]
struct FileResolver;
use crate::native::CredentialAttempt;
struct ChainAcquirer {
terminal_on: Option<&'static str>,
requests: Mutex<Vec<&'static str>>,
succeed_on: Option<&'static str>,
}
@ -518,20 +697,25 @@ mod tests {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> std::pin::Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>>
{
) -> std::pin::Pin<Box<dyn Future<Output = CredentialAttempt> + Send + '_>> {
let kind = request.kind();
self.requests.lock().unwrap().push(kind);
Box::pin(async move {
if self.succeed_on == Some(kind) {
Ok(ResolvedCredential::AccessToken {
CredentialAttempt::Resolved(ResolvedCredential::AccessToken {
token: SecretValue::new("chain-token"),
expires_on: None,
})
} else {
Err(Error::CredentialAcquisition(
litellm_auth_types::ErrorDetail::Unresolved { subject: kind },
))
let error =
Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::Unresolved {
subject: kind,
});
if self.terminal_on == Some(kind) {
CredentialAttempt::Failed(error)
} else {
CredentialAttempt::Unavailable(error)
}
}
})
}
@ -690,9 +874,210 @@ mod tests {
assert_eq!(resolved, Some(SecretValue::new("rotated-assertion")));
}
#[derive(Debug)]
struct UnavailableImdsClient;
impl azure_core::http::HttpClient for UnavailableImdsClient {
fn execute_request<'life0, 'life1, 'async_trait>(
&'life0 self,
_request: &'life1 azure_core::http::Request,
) -> std::pin::Pin<
Box<
dyn Future<Output = azure_core::Result<azure_core::http::AsyncRawResponse>>
+ Send
+ 'async_trait,
>,
>
where
'life0: 'async_trait,
'life1: 'async_trait,
Self: 'async_trait,
{
Box::pin(std::future::pending())
}
}
struct OfflineChainAcquirer(super::NativeAzureTokenAcquirer);
impl AzureTokenAcquirer for OfflineChainAcquirer {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> std::pin::Pin<Box<dyn Future<Output = CredentialAttempt> + Send + '_>> {
Box::pin(async move {
if request.kind() == "developer-tools" {
return CredentialAttempt::Resolved(ResolvedCredential::AccessToken {
token: SecretValue::new("cli-token"),
expires_on: None,
});
}
self.0.acquire(request).await
})
}
}
#[rstest::rstest]
#[case::default_off_azure("DefaultAzureCredential", None, true)]
#[case::explicit_managed_identity("ManagedIdentityCredential", None, false)]
#[case::app_service("DefaultAzureCredential", Some("http://app-service.test"), false)]
#[tokio::test(start_paused = true)]
async fn only_default_imds_discovery_is_bounded(
#[case] selector: &str,
#[case] identity_endpoint: Option<&str>,
#[case] expect_cli: bool,
) {
let native = super::NativeAzureTokenAcquirer::with_transport(
4,
azure_core::http::Transport::new(Arc::new(UnavailableImdsClient)),
);
let service = AzureAuthService::with_acquirer(Arc::new(OfflineChainAcquirer(native)));
let inputs = AzureAuthInputs {
azure_credential: crate::types::ConfigValue::Value(Sourced::new(
selector.into(),
InputSource::Deployment,
)),
enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment),
..Default::default()
};
let environment = |name: &str| {
(name == "IDENTITY_ENDPOINT")
.then_some(identity_endpoint)
.flatten()
.map(str::to_string)
};
let outcome = tokio::time::timeout(
std::time::Duration::from_secs(2),
service.get_azure_ad_token(&inputs, &environment),
)
.await;
assert_eq!(outcome.is_ok(), expect_cli);
if expect_cli {
assert_eq!(
outcome.unwrap().unwrap().unwrap().value().secret().expose(),
"cli-token"
);
}
}
#[rstest::rstest]
#[case::app_service(Some("IDENTITY_ENDPOINT"), false)]
#[case::legacy_app_service(Some("MSI_ENDPOINT"), false)]
#[case::off_azure(None, true)]
#[tokio::test(start_paused = true)]
async fn default_chain_uses_endpoints_from_the_secret_snapshot(
#[case] endpoint: Option<&str>,
#[case] expect_cli: bool,
#[values(false, true)] consumer_identity: bool,
) {
use litellm_secrets::source::SecretSource;
struct EndpointSecrets<'a>(Option<&'a str>);
impl SecretSource for EndpointSecrets<'_> {
fn get_secret_str<'a>(
&'a self,
name: &'a str,
) -> std::pin::Pin<
Box<
dyn Future<Output = Result<Option<SecretValue>, litellm_secrets::Error>>
+ Send
+ 'a,
>,
> {
Box::pin(async move {
Ok((self.0 == Some(name)).then(|| SecretValue::new("http://app-service.test")))
})
}
}
let snapshot = EndpointSecrets(endpoint)
.resolve(SECRET_NAMES)
.await
.unwrap();
let native = super::NativeAzureTokenAcquirer::with_transport(
4,
azure_core::http::Transport::new(Arc::new(UnavailableImdsClient)),
);
let service = AzureAuthService::with_acquirer(Arc::new(OfflineChainAcquirer(native)));
let lookup = |name: &str| snapshot.get(name);
let scope = "https://service.test/.default";
let outcome = tokio::time::timeout(std::time::Duration::from_secs(2), async {
if consumer_identity {
service
.identity_token(&super::AzureIdentity::Default, scope, &lookup)
.await
} else {
service
.get_azure_ad_token(
&AzureAuthInputs::default_credential_for_scope(scope),
&lookup,
)
.await
.map(|token| token.map(Sourced::into_value))
}
})
.await;
assert_eq!(outcome.is_ok(), expect_cli);
if expect_cli {
assert_eq!(
outcome.unwrap().unwrap().unwrap().secret().expose(),
"cli-token"
);
}
}
#[rstest::rstest]
#[case::litellm_provider_settings_are_ignored(
&[("AZURE_AD_TOKEN", "provider-token"), ("AZURE_CREDENTIAL", "ClientSecretCredential")],
&["managed-identity", "developer-tools"]
)]
#[case::environment_client_secret_comes_first(
&[("AZURE_TENANT_ID", "tenant"), ("AZURE_CLIENT_ID", "client"), ("AZURE_CLIENT_SECRET", "secret")],
&["client-secret", "managed-identity", "developer-tools"]
)]
#[case::aks_workload_identity(
&[("AZURE_TENANT_ID", "tenant"), ("AZURE_CLIENT_ID", "client"), ("AZURE_FEDERATED_TOKEN_FILE", "/var/run/token")],
&["workload-identity", "managed-identity", "developer-tools"]
)]
#[tokio::test]
async fn the_default_identity_follows_azure_identity_defaults(
#[case] environment: &[(&str, &str)],
#[case] expected: &[&str],
) {
let acquirer = Arc::new(ChainAcquirer {
terminal_on: None,
requests: Mutex::new(Vec::new()),
succeed_on: Some("developer-tools"),
});
let service = AzureAuthService::with_acquirer(acquirer.clone());
let lookup = |name: &str| {
environment
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| (*value).to_owned())
};
let credential = service
.identity_token(
&super::AzureIdentity::Default,
"https://redis.azure.com/.default",
&lookup,
)
.await
.unwrap()
.unwrap();
assert_eq!(credential.secret().expose(), "chain-token");
assert_eq!(*acquirer.requests.lock().unwrap(), expected);
}
#[rstest::rstest]
#[tokio::test]
async fn default_chain_uses_declared_order_and_stops_after_success() {
let acquirer = Arc::new(ChainAcquirer {
terminal_on: None,
requests: Mutex::new(Vec::new()),
succeed_on: Some("developer-tools"),
});
@ -715,9 +1100,13 @@ mod tests {
);
}
#[rstest::rstest]
#[case::all_unavailable(false)]
#[case::terminal_failure_after_unavailable(true)]
#[tokio::test]
async fn chain_reports_each_acquisition_failure() {
async fn chain_reports_each_acquisition_failure(#[case] terminal: bool) {
let acquirer = Arc::new(ChainAcquirer {
terminal_on: terminal.then_some("developer-tools"),
requests: Mutex::new(Vec::new()),
succeed_on: None,
});
@ -744,6 +1133,27 @@ mod tests {
);
}
#[rstest::rstest]
#[tokio::test]
async fn a_configured_identity_failure_does_not_fall_back_to_a_developer_identity() {
let acquirer = Arc::new(ChainAcquirer {
terminal_on: Some("managed-identity"),
requests: Mutex::new(Vec::new()),
succeed_on: Some("developer-tools"),
});
let service = AzureAuthService::with_acquirer(acquirer.clone());
let inputs = AzureAuthInputs {
enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment),
..Default::default()
};
let error = service
.get_azure_ad_token(&inputs, &|_| None)
.await
.unwrap_err();
assert_eq!(*acquirer.requests.lock().unwrap(), ["managed-identity"]);
assert!(matches!(error, Error::CredentialChain { earlier, .. } if earlier.is_empty()));
}
#[derive(Debug)]
struct CallerToken(&'static str);

View file

@ -6,7 +6,7 @@ license.workspace = true
repository.workspace = true
[features]
google-sdk = ["dep:google-cloud-auth", "dep:http"]
google-sdk = ["dep:http"]
[dependencies]
litellm-secret-value.workspace = true
@ -21,10 +21,14 @@ tokio.workspace = true
veil.workspace = true
gcp_auth = "0.12.7"
google-cloud-auth = { workspace = true, optional = true }
google-cloud-auth.workspace = true
http = { workspace = true, optional = true }
[dev-dependencies]
time = { workspace = true, features = ["formatting"] }
axum = { workspace = true, features = ["json"] }
litellm-error.workspace = true
rstest.workspace = true
rustls.workspace = true
tempfile.workspace = true
tokio = { workspace = true, features = ["test-util"] }

View file

@ -1,20 +1,11 @@
use std::{future::Future, pin::Pin, sync::Arc};
use std::sync::Arc;
use gcp_auth::{CustomServiceAccount, TokenProvider};
use litellm_auth_types::{CredentialPlacement, Error, ErrorDetail, http::apply_credential};
use moka::future::Cache;
use serde_json::Value;
use sha2::{Digest, Sha256};
use litellm_auth_types::{CredentialPlacement, Error, http::apply_credential};
use crate::{
config::{
CredentialSource, VertexConfig, credential_source, get_vertex_ai_project,
is_base64_encoded_json, names_a_credentials_file, non_empty_env,
},
constants::{
CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT, VERTEX_AI_API_KEY_ENV,
VERTEXAI_API_KEY_ENV,
},
config::{VertexConfig, credential_source, get_vertex_ai_project, non_empty_env},
constants::{VERTEX_AI_API_KEY_ENV, VERTEXAI_API_KEY_ENV},
identity::{GcpAuth, GcpProviderLoader, GcpTokenSource},
};
pub struct VertexEnvironment {
@ -22,28 +13,15 @@ pub struct VertexEnvironment {
pub project_id: String,
}
struct VertexAccessToken {
token: String,
project_id: String,
}
#[derive(Clone)]
#[derive(Clone, Default)]
pub struct VertexAuth {
providers: Cache<CredentialCacheKey, Arc<dyn VertexTokenSource>>,
loader: Arc<dyn VertexProviderLoader>,
}
impl Default for VertexAuth {
fn default() -> Self {
Self::new(Arc::new(GcpProviderLoader))
}
identity: GcpAuth,
}
impl VertexAuth {
pub fn new(loader: Arc<dyn VertexProviderLoader>) -> Self {
pub fn new(loader: Arc<dyn GcpProviderLoader>) -> Self {
Self {
providers: Cache::builder().max_capacity(64).build(),
loader,
identity: GcpAuth::new(loader),
}
}
@ -52,7 +30,10 @@ impl VertexAuth {
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<String, Error> {
self.load_provider(config, env_lookup).await?.token().await
self.identity
.access_token(credential_source(config, env_lookup))
.await
.map(|token| token.expose().to_owned())
}
pub async fn validate_environment(
@ -77,10 +58,18 @@ impl VertexAuth {
(true, _) => None,
(false, Some(token)) => Some(token),
(false, None) => {
let access = self.get_access_token(config, env_lookup).await?;
let provider = self.load_provider(config, env_lookup).await?;
let (token, project_id) = match project_id {
Some(project_id) => (provider.token().await?, project_id),
None => tokio::try_join!(provider.token(), provider.project_id())?,
};
return Ok(VertexEnvironment {
headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?,
project_id: project_id.unwrap_or(access.project_id),
headers: apply_credential(
headers,
token.expose(),
CredentialPlacement::Bearer,
)?,
project_id,
});
}
};
@ -104,248 +93,171 @@ impl VertexAuth {
})
}
async fn get_access_token(
&self,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<VertexAccessToken, Error> {
let provider = self.load_provider(config, env_lookup).await?;
let (token, project_id) = tokio::try_join!(provider.token(), provider.project_id())?;
Ok(VertexAccessToken { token, project_id })
}
async fn load_provider(
&self,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Arc<dyn VertexTokenSource>, Error> {
let source = credential_source(config, env_lookup);
let key = source.cache_key();
self.providers
.try_get_with(key, self.loader.load(source))
) -> Result<Arc<dyn GcpTokenSource>, Error> {
self.identity
.load_provider(credential_source(config, env_lookup))
.await
.map_err(|error| (*error).clone())
}
}
pub trait VertexTokenSource: Send + Sync {
fn project_id(&self) -> VertexAuthFuture<'_, String>;
fn token(&self) -> VertexAuthFuture<'_, String>;
}
pub trait VertexProviderLoader: Send + Sync {
fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>>;
}
pub type VertexAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
struct GcpTokenSource(Arc<dyn TokenProvider>);
impl VertexTokenSource for GcpTokenSource {
fn project_id(&self) -> VertexAuthFuture<'_, String> {
Box::pin(async move {
self.0
.project_id()
.await
.map(|project| project.to_string())
.map_err(auth_acquisition_error)
})
}
fn token(&self) -> VertexAuthFuture<'_, String> {
Box::pin(async move {
self.0
.token(&[CLOUD_PLATFORM_SCOPE])
.await
.map(|token| token.as_str().to_string())
.map_err(auth_acquisition_error)
})
}
}
struct GcpProviderLoader;
impl VertexProviderLoader for GcpProviderLoader {
fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
Box::pin(async move {
let provider: Arc<dyn TokenProvider> = match source {
CredentialSource::Inline(configured) => Arc::new(
CustomServiceAccount::from_json(validate_request_credentials(
configured.expose(),
)?)
.map_err(auth_acquisition_error)?,
),
CredentialSource::Trusted(configured) => {
Arc::new(load_trusted_service_account(configured.expose())?)
}
CredentialSource::ApplicationCredentials(path) => {
Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?)
}
CredentialSource::Adc => {
gcp_auth::provider().await.map_err(auth_acquisition_error)?
}
};
Ok(Arc::new(GcpTokenSource(provider)) as Arc<dyn VertexTokenSource>)
})
}
}
fn validate_request_credentials(configured: &str) -> Result<&str, Error> {
let token_uri = serde_json::from_str::<Value>(configured)
.ok()
.and_then(|credentials| {
credentials
.get("token_uri")
.and_then(Value::as_str)
.map(str::to_string)
});
if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) {
return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType {
field: "request-controlled vertex_credentials token_uri".into(),
expected: "the canonical Google OAuth token endpoint",
}));
}
Ok(configured)
}
impl CredentialSource {
fn cache_key(&self) -> CredentialCacheKey {
match self {
Self::Inline(configured) => {
CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into())
}
Self::Trusted(configured) => {
CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into())
}
Self::ApplicationCredentials(path) => {
CredentialCacheKey::ApplicationCredentials(path.clone())
}
Self::Adc => CredentialCacheKey::Adc,
}
}
}
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
enum CredentialCacheKey {
Inline([u8; 32]),
Trusted([u8; 32]),
ApplicationCredentials(String),
Adc,
}
fn load_trusted_service_account(configured: &str) -> Result<CustomServiceAccount, Error> {
if is_base64_encoded_json(configured) {
return Err(Error::InvalidConfiguration(
ErrorDetail::Base64EncodedJson {
field: "vertex_credentials",
},
));
}
if names_a_credentials_file(configured) {
CustomServiceAccount::from_file(configured)
} else {
CustomServiceAccount::from_json(configured)
}
.map_err(auth_acquisition_error)
}
fn auth_acquisition_error(error: gcp_auth::Error) -> Error {
Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed(
"Vertex AI credentials",
error,
))
}
#[cfg(test)]
mod tests {
use std::{
collections::{BTreeMap, BTreeSet},
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
};
use crate::{
CredentialSource, GcpAuthFuture, GcpProviderLoader, GcpTokenSource, VertexAuth,
VertexConfig, get_vertex_ai_location, secret_names,
};
use litellm_secret_value::SecretValue;
use rstest::rstest;
use serde_json::{Value, json};
use super::*;
#[rstest::rstest]
#[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)]
#[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)]
#[case::missing_endpoint("{}", false)]
fn request_credentials_require_canonical_token_endpoint(
#[case] credentials: &str,
#[case] accepted: bool,
) {
assert_eq!(validate_request_credentials(credentials).is_ok(), accepted);
struct FakeProvider {
project: Option<&'static str>,
calls: Arc<AtomicUsize>,
}
#[rstest::rstest]
#[case::missing_file(
"/nonexistent/vertex-key.json",
"failed to open application credentials file"
)]
#[case::malformed_json(
" {\"type\": \"service_account\",",
"failed to deserialize ApplicationCredentials"
)]
#[case::quoted_json(
"'{\"type\": \"service_account\"}'",
"failed to deserialize ApplicationCredentials"
)]
#[tokio::test]
async fn trusted_credentials_are_read_as_a_file_unless_they_are_json(
#[case] configured: &str,
#[case] expected: &str,
) {
let Err(error) = GcpProviderLoader
.load(CredentialSource::Trusted(SecretValue::new(configured)))
.await
else {
panic!("{configured:?} should not load");
};
let report = litellm_error::Report::new(&error).to_string();
assert!(report.contains(expected), "{report}");
}
impl GcpTokenSource for FakeProvider {
fn project_id(&self) -> GcpAuthFuture<'_, String> {
self.calls.fetch_add(1, Ordering::SeqCst);
Box::pin(async move {
self.project.map(str::to_owned).ok_or(
litellm_auth_types::Error::CredentialAcquisition(
litellm_auth_types::ErrorDetail::Unresolved {
subject: "credential project",
},
),
)
})
}
#[rstest::rstest]
#[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")]
#[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")]
#[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")]
#[tokio::test]
async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) {
assert!(
trusted_load_error(configured)
.await
.contains("is base64-encoded JSON, not the JSON itself")
);
}
async fn trusted_load_error(configured: &str) -> String {
match GcpProviderLoader
.load(CredentialSource::Trusted(SecretValue::new(configured)))
.await
{
Ok(_) => panic!("{configured:?} should not load"),
Err(error) => litellm_error::Report::new(&error).to_string(),
fn token(&self) -> GcpAuthFuture<'_, SecretValue> {
self.calls.fetch_add(1, Ordering::SeqCst);
Box::pin(async { Ok(SecretValue::new("adc-token")) })
}
}
#[rstest::rstest]
#[tokio::test]
async fn a_relative_key_file_named_like_json_is_read_as_a_file() {
let key = tempfile::Builder::new()
.prefix("{vertex")
.suffix(".json")
.tempfile_in(".")
.unwrap();
std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap();
let relative_name = key.path().file_name().unwrap().to_str().unwrap();
let absolute_path = std::fs::canonicalize(key.path()).unwrap();
assert_eq!(
trusted_load_error(relative_name).await,
trusted_load_error(absolute_path.to_str().unwrap()).await
);
struct FakeLoader {
loads: Arc<AtomicUsize>,
provider: Arc<dyn GcpTokenSource>,
}
#[rstest::rstest]
fn inline_and_trusted_credentials_never_share_a_cache_entry() {
assert_ne!(
CredentialSource::Inline(SecretValue::new("same-value")).cache_key(),
CredentialSource::Trusted(SecretValue::new("same-value")).cache_key()
);
impl GcpProviderLoader for FakeLoader {
fn load(&self, _source: CredentialSource) -> GcpAuthFuture<'_, Arc<dyn GcpTokenSource>> {
let loads = self.loads.clone();
let provider = self.provider.clone();
Box::pin(async move {
loads.fetch_add(1, Ordering::SeqCst);
Ok(provider)
})
}
}
fn config(value: Value) -> VertexConfig {
VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new())
.unwrap()
}
fn auth(calls: Arc<AtomicUsize>, loads: Arc<AtomicUsize>) -> VertexAuth {
let provider: Arc<dyn GcpTokenSource> = Arc::new(FakeProvider {
calls,
project: Some("adc-project"),
});
VertexAuth::new(Arc::new(FakeLoader { loads, provider }))
}
#[rstest]
#[tokio::test]
async fn secret_names_cover_environment_reads() {
let seen = Arc::new(std::sync::Mutex::new(BTreeSet::<String>::new()));
let recorded = seen.clone();
let env = |name: &str| {
recorded.lock().unwrap().insert(name.to_string());
None
};
let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0)));
auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env)
.await
.unwrap();
get_vertex_ai_location(&VertexConfig::default(), &env);
let allowed = secret_names()
.iter()
.map(|name| (*name).to_owned())
.collect::<BTreeSet<_>>();
assert!(seen.lock().unwrap().is_subset(&allowed));
}
#[rstest]
#[tokio::test]
async fn explicit_token_and_header_do_not_acquire_adc() {
let loads = Arc::new(AtomicUsize::new(0));
let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone());
let configured = config(json!({"vertex_project":"project-1"}));
let explicit = auth
.validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None)
.await
.unwrap();
assert_eq!(explicit.headers[0].1, "Bearer access-token");
let existing = auth
.validate_environment(
vec![("authorization".into(), "Bearer existing".into())],
None,
&configured,
&|_| None,
)
.await
.unwrap();
assert_eq!(existing.headers[0].1, "Bearer existing");
assert_eq!(loads.load(Ordering::SeqCst), 0);
}
#[rstest]
#[tokio::test]
async fn provider_is_reused_across_authentication_calls() {
let calls = Arc::new(AtomicUsize::new(0));
let loads = Arc::new(AtomicUsize::new(0));
let auth = auth(calls.clone(), loads.clone());
let config = VertexConfig::default();
let environment = |_: &str| None;
let (first, second) = tokio::try_join!(
auth.validate_environment(Vec::new(), None, &config, &environment),
auth.validate_environment(Vec::new(), None, &config, &environment),
)
.unwrap();
assert_eq!(first.project_id, "adc-project");
assert_eq!(first.headers[0].1, "Bearer adc-token");
assert_eq!(first.project_id, second.project_id);
assert_eq!(first.headers, second.headers);
assert_eq!(loads.load(Ordering::SeqCst), 1);
assert_eq!(calls.load(Ordering::SeqCst), 4);
}
#[rstest]
#[tokio::test]
async fn an_explicit_project_does_not_require_project_metadata_in_the_credentials() {
let auth = VertexAuth::new(Arc::new(FakeLoader {
loads: Arc::new(AtomicUsize::new(0)),
provider: Arc::new(FakeProvider {
calls: Arc::new(AtomicUsize::new(0)),
project: None,
}),
}));
let configured = VertexConfig::new(None, Some("configured-project".into()), None);
let environment = auth
.validate_environment(Vec::new(), None, &configured, &|_| None)
.await
.unwrap();
assert_eq!(environment.project_id, "configured-project");
assert_eq!(environment.headers[0].1, "Bearer adc-token");
}
}

View file

@ -146,6 +146,20 @@ pub enum CredentialSource {
Adc,
}
impl CredentialSource {
pub fn from_vertex_environment(env_lookup: &dyn Fn(&str) -> Option<String>) -> Self {
VertexParams::CREDENTIALS
.resolve(&|_| None, env_lookup)
.map(|configured| Self::Trusted(SecretValue::new(configured)))
.unwrap_or_else(|| Self::application_default(env_lookup))
}
pub fn application_default(env_lookup: &dyn Fn(&str) -> Option<String>) -> Self {
non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV)
.map_or(Self::Adc, Self::ApplicationCredentials)
}
}
pub(crate) fn credential_source(
config: &VertexConfig,
env_lookup: &dyn Fn(&str) -> Option<String>,
@ -158,12 +172,7 @@ pub(crate) fn credential_source(
}
};
}
if let Some(configured) = VertexParams::CREDENTIALS.resolve(&|_| None, env_lookup) {
return CredentialSource::Trusted(SecretValue::new(configured));
}
non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV)
.map(CredentialSource::ApplicationCredentials)
.unwrap_or(CredentialSource::Adc)
CredentialSource::from_vertex_environment(env_lookup)
}
fn optional_credentials(
@ -319,6 +328,24 @@ mod tests {
));
}
#[rstest::rstest]
fn application_default_ignores_vertex_credentials() {
let env = |name: &str| match name {
VERTEXAI_CREDENTIALS_ENV => Some("environment-json".to_string()),
GOOGLE_APPLICATION_CREDENTIALS_ENV => Some("adc.json".to_string()),
_ => None,
};
assert!(
matches!(CredentialSource::application_default(&env), CredentialSource::ApplicationCredentials(path) if path == "adc.json")
);
assert!(matches!(
CredentialSource::application_default(
&|name| (name == VERTEXAI_CREDENTIALS_ENV).then(|| "environment-json".into())
),
CredentialSource::Adc
));
}
#[rstest::rstest]
fn params_become_a_deployment_sourced_config() {
let params = VertexParams {

View file

@ -0,0 +1,445 @@
use crate::{
config::{CredentialSource, is_base64_encoded_json, names_a_credentials_file},
constants::{CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT},
};
use gcp_auth::{CustomServiceAccount, TokenProvider};
use litellm_auth_types::{Error, ErrorDetail};
use litellm_secret_value::SecretValue;
use moka::future::Cache;
use serde_json::Value;
use sha2::{Digest, Sha256};
use std::{future::Future, pin::Pin, sync::Arc};
#[derive(Clone)]
pub struct GcpAuth {
providers: Cache<CredentialCacheKey, Arc<dyn GcpTokenSource>>,
loader: Arc<dyn GcpProviderLoader>,
}
impl Default for GcpAuth {
fn default() -> Self {
Self::new(Arc::new(NativeProviderLoader))
}
}
impl GcpAuth {
pub fn new(loader: Arc<dyn GcpProviderLoader>) -> Self {
Self {
providers: Cache::builder().max_capacity(64).build(),
loader,
}
}
pub async fn access_token(&self, source: CredentialSource) -> Result<SecretValue, Error> {
self.load_provider(source).await?.token().await
}
pub(crate) async fn load_provider(
&self,
source: CredentialSource,
) -> Result<Arc<dyn GcpTokenSource>, Error> {
self.providers
.try_get_with(source.cache_key(), self.loader.load(source))
.await
.map_err(|error| (*error).clone())
}
}
pub trait GcpTokenSource: Send + Sync {
fn project_id(&self) -> GcpAuthFuture<'_, String>;
fn token(&self) -> GcpAuthFuture<'_, SecretValue>;
}
pub trait GcpProviderLoader: Send + Sync {
fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc<dyn GcpTokenSource>>;
}
pub type GcpAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
struct ServiceAccountTokenSource(Arc<dyn TokenProvider>);
impl GcpTokenSource for ServiceAccountTokenSource {
fn project_id(&self) -> GcpAuthFuture<'_, String> {
Box::pin(async move {
self.0
.project_id()
.await
.map(|project| project.to_string())
.map_err(auth_acquisition_error)
})
}
fn token(&self) -> GcpAuthFuture<'_, SecretValue> {
Box::pin(async move {
self.0
.token(&[CLOUD_PLATFORM_SCOPE])
.await
.map(|token| SecretValue::new(token.as_str()))
.map_err(auth_acquisition_error)
})
}
}
struct SdkTokenSource {
credentials: google_cloud_auth::credentials::AccessTokenCredentials,
project_id: Option<String>,
}
impl GcpTokenSource for SdkTokenSource {
fn project_id(&self) -> GcpAuthFuture<'_, String> {
Box::pin(async move {
self.project_id
.clone()
.ok_or(Error::CredentialAcquisition(ErrorDetail::Unresolved {
subject: "GCP resource project",
}))
})
}
fn token(&self) -> GcpAuthFuture<'_, SecretValue> {
Box::pin(async move {
self.credentials
.access_token()
.await
.map(|token| SecretValue::new(token.token))
.map_err(auth_acquisition_error)
})
}
}
async fn load_application_credentials(path: &str) -> Result<Arc<dyn GcpTokenSource>, Error> {
let contents = tokio::fs::read_to_string(path)
.await
.map_err(auth_acquisition_error)?;
load_credentials_json(&contents)
}
fn load_credentials_json(contents: &str) -> Result<Arc<dyn GcpTokenSource>, Error> {
let configured: Value = serde_json::from_str(contents).map_err(auth_acquisition_error)?;
let kind = configured.get("type").and_then(Value::as_str);
let project_id = configured
.get("project_id")
.and_then(Value::as_str)
.map(str::to_string);
let credentials = match kind {
Some("authorized_user") => {
google_cloud_auth::credentials::user_account::Builder::new(configured)
.with_scopes([CLOUD_PLATFORM_SCOPE])
.build_access_token_credentials()
}
Some("impersonated_service_account") => {
google_cloud_auth::credentials::impersonated::Builder::new(configured)
.with_scopes([CLOUD_PLATFORM_SCOPE])
.build_access_token_credentials()
}
Some("external_account") => {
google_cloud_auth::credentials::external_account::Builder::new(configured)
.with_scopes([CLOUD_PLATFORM_SCOPE])
.build_access_token_credentials()
}
Some("service_account") | None => {
let provider =
CustomServiceAccount::from_json(contents).map_err(auth_acquisition_error)?;
return Ok(Arc::new(ServiceAccountTokenSource(Arc::new(provider))));
}
Some(_) => {
return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType {
field: "Google application credentials type".into(),
expected: "a supported Google credential type",
}));
}
}
.map_err(auth_acquisition_error)?;
Ok(Arc::new(SdkTokenSource {
credentials,
project_id,
}))
}
struct NativeProviderLoader;
impl GcpProviderLoader for NativeProviderLoader {
fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc<dyn GcpTokenSource>> {
Box::pin(async move {
let provider: Arc<dyn TokenProvider> = match source {
CredentialSource::Inline(configured) => Arc::new(
CustomServiceAccount::from_json(validate_request_credentials(
configured.expose(),
)?)
.map_err(auth_acquisition_error)?,
),
CredentialSource::Trusted(configured) => {
return load_trusted_credentials(configured.expose()).await;
}
CredentialSource::ApplicationCredentials(path) => {
return load_application_credentials(&path).await;
}
CredentialSource::Adc => {
gcp_auth::provider().await.map_err(auth_acquisition_error)?
}
};
Ok(Arc::new(ServiceAccountTokenSource(provider)) as Arc<dyn GcpTokenSource>)
})
}
}
fn validate_request_credentials(configured: &str) -> Result<&str, Error> {
let token_uri = serde_json::from_str::<Value>(configured)
.ok()
.and_then(|credentials| {
credentials
.get("token_uri")
.and_then(Value::as_str)
.map(str::to_string)
});
if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) {
return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType {
field: "request-controlled vertex_credentials token_uri".into(),
expected: "the canonical Google OAuth token endpoint",
}));
}
Ok(configured)
}
impl CredentialSource {
fn cache_key(&self) -> CredentialCacheKey {
match self {
Self::Inline(configured) => {
CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into())
}
Self::Trusted(configured) => {
CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into())
}
Self::ApplicationCredentials(path) => {
CredentialCacheKey::ApplicationCredentials(path.clone())
}
Self::Adc => CredentialCacheKey::Adc,
}
}
}
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
enum CredentialCacheKey {
Inline([u8; 32]),
Trusted([u8; 32]),
ApplicationCredentials(String),
Adc,
}
async fn load_trusted_credentials(configured: &str) -> Result<Arc<dyn GcpTokenSource>, Error> {
if is_base64_encoded_json(configured) {
return Err(Error::InvalidConfiguration(
ErrorDetail::Base64EncodedJson {
field: "vertex_credentials",
},
));
}
if names_a_credentials_file(configured) {
load_application_credentials(configured).await
} else {
load_credentials_json(configured)
}
}
fn auth_acquisition_error(error: impl std::error::Error + Send + Sync + 'static) -> Error {
Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed(
"Google credentials",
error,
))
}
#[cfg(test)]
mod tests {
use litellm_secret_value::SecretValue;
use super::*;
#[rstest::rstest]
#[case::authorized_user("authorized_user", 2)]
#[case::external_account("external_account", 2)]
#[case::impersonated_service_account("impersonated_service_account", 4)]
#[tokio::test]
async fn explicit_application_credentials_refresh_expired_tokens(
#[case] kind: &str,
#[case] expected_requests: usize,
#[values(false, true)] inline: bool,
) {
use axum::{Json, Router, routing::post};
use std::sync::atomic::{AtomicUsize, Ordering};
let _ = rustls::crypto::ring::default_provider().install_default();
let requests = Arc::new(AtomicUsize::new(0));
let recorded = requests.clone();
let impersonations = requests.clone();
let app = Router::new().route("/token", post(move || {
let recorded = recorded.clone();
async move {
Json(serde_json::json!({
"access_token": format!("token-{}", recorded.fetch_add(1, Ordering::SeqCst)),
"expires_in": 3600,
"token_type": "Bearer",
"issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
}))
}
}));
let app = app.route("/v1/projects/-/serviceAccounts/test-principal:generateAccessToken", post(move || {
let recorded = impersonations.clone();
async move {
let expiration = (time::OffsetDateTime::now_utc() + time::Duration::hours(1)).format(&time::format_description::well_known::Rfc3339).unwrap();
Json(serde_json::json!({
"accessToken": format!("impersonated-{}", recorded.fetch_add(1, Ordering::SeqCst)),
"expireTime": expiration,
}))
}
}));
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
let file = tempfile::NamedTempFile::new().unwrap();
let subject_token = tempfile::NamedTempFile::new().unwrap();
std::fs::write(subject_token.path(), "subject-token").unwrap();
let configured = if kind == "external_account" {
serde_json::json!({
"type": "external_account", "audience": "//iam.googleapis.com/projects/test/locations/global/workloadIdentityPools/test/providers/test",
"subject_token_type": "urn:ietf:params:oauth:token-type:jwt", "token_url": format!("http://{address}/token"),
"credential_source": { "file": subject_token.path().to_str().unwrap() }, "quota_project_id": "quota-project"
})
} else {
let source = serde_json::json!({
"type": "authorized_user", "client_id": "user-client", "client_secret": "user-secret",
"refresh_token": "user-refresh", "quota_project_id": "quota-project", "token_uri": format!("http://{address}/token")
});
if kind == "impersonated_service_account" {
serde_json::json!({
"type": kind, "source_credentials": source, "quota_project_id": "quota-project",
"service_account_impersonation_url": format!("http://{address}/v1/projects/-/serviceAccounts/test-principal:generateAccessToken")
})
} else {
source
}
};
std::fs::write(file.path(), configured.to_string()).unwrap();
let source = if inline {
CredentialSource::Trusted(SecretValue::new(configured.to_string()))
} else {
CredentialSource::ApplicationCredentials(file.path().to_str().unwrap().into())
};
let provider = NativeProviderLoader.load(source).await.unwrap();
assert!(matches!(
provider.project_id().await,
Err(Error::CredentialAcquisition(ErrorDetail::Unresolved {
subject: "GCP resource project"
}))
));
let original = provider.token().await.unwrap();
tokio::time::pause();
tokio::time::advance(std::time::Duration::from_secs(3601)).await;
tokio::time::resume();
let refreshed = provider.token().await.unwrap();
assert_ne!(original, refreshed);
assert_eq!(requests.load(Ordering::SeqCst), expected_requests);
server.abort();
}
#[rstest::rstest]
#[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)]
#[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)]
#[case::missing_endpoint("{}", false)]
fn request_credentials_require_canonical_token_endpoint(
#[case] credentials: &str,
#[case] accepted: bool,
) {
assert_eq!(validate_request_credentials(credentials).is_ok(), accepted);
}
#[rstest::rstest]
#[case::missing_file("/nonexistent/vertex-key.json", "Google credentials")]
#[case::malformed_json(" {\"type\": \"service_account\",", "Google credentials")]
#[case::quoted_json("'{\"type\": \"service_account\"}'", "Google credentials")]
#[tokio::test]
async fn trusted_credentials_are_read_as_a_file_unless_they_are_json(
#[case] configured: &str,
#[case] expected: &str,
) {
let Err(error) = NativeProviderLoader
.load(CredentialSource::Trusted(SecretValue::new(configured)))
.await
else {
panic!("{configured:?} should not load");
};
let report = litellm_error::Report::new(&error).to_string();
assert!(report.contains(expected), "{report}");
}
#[rstest::rstest]
#[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")]
#[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")]
#[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")]
#[tokio::test]
async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) {
assert!(
trusted_load_error(configured)
.await
.contains("is base64-encoded JSON, not the JSON itself")
);
}
async fn trusted_load_error(configured: &str) -> String {
match NativeProviderLoader
.load(CredentialSource::Trusted(SecretValue::new(configured)))
.await
{
Ok(_) => panic!("{configured:?} should not load"),
Err(error) => litellm_error::Report::new(&error).to_string(),
}
}
#[rstest::rstest]
#[tokio::test]
async fn a_relative_key_file_named_like_json_is_read_as_a_file() {
let key = tempfile::Builder::new()
.prefix("{vertex")
.suffix(".json")
.tempfile_in(".")
.unwrap();
std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap();
let relative_name = key.path().file_name().unwrap().to_str().unwrap();
let absolute_path = std::fs::canonicalize(key.path()).unwrap();
assert_eq!(
trusted_load_error(relative_name).await,
trusted_load_error(absolute_path.to_str().unwrap()).await
);
}
#[rstest::rstest]
#[tokio::test]
async fn unknown_application_credential_types_are_rejected_without_service_account_fallback() {
let file = tempfile::NamedTempFile::new().unwrap();
std::fs::write(
file.path(),
r#"{"type":"unsupported","private_key":"credential-value"}"#,
)
.unwrap();
let result = GcpAuth::default()
.access_token(CredentialSource::ApplicationCredentials(
file.path().to_str().unwrap().into(),
))
.await;
let error = result.unwrap_err();
assert!(matches!(error, Error::InvalidConfiguration(_)));
assert!(
!litellm_error::Report::new(&error)
.to_string()
.contains("credential-value")
);
}
#[rstest::rstest]
fn inline_and_trusted_credentials_never_share_a_cache_entry() {
assert_ne!(
CredentialSource::Inline(SecretValue::new("same-value")).cache_key(),
CredentialSource::Trusted(SecretValue::new("same-value")).cache_key()
);
}
}

View file

@ -1,16 +1,16 @@
mod auth;
mod config;
pub mod constants;
mod identity;
#[cfg(feature = "google-sdk")]
mod sdk;
pub use auth::{
VertexAuth, VertexAuthFuture, VertexEnvironment, VertexProviderLoader, VertexTokenSource,
};
pub use auth::{VertexAuth, VertexEnvironment};
pub use config::{
CredentialSource, VertexConfig, get_vertex_ai_location, get_vertex_ai_project,
get_vertex_ai_project_from_credentials,
};
pub use constants::secret_names;
pub use identity::{GcpAuth, GcpAuthFuture, GcpProviderLoader, GcpTokenSource};
#[cfg(feature = "google-sdk")]
pub use sdk::GoogleCredentials;

View file

@ -1,120 +0,0 @@
use std::{
collections::{BTreeMap, BTreeSet},
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
};
use litellm_auth_gcp::{
CredentialSource, VertexAuth, VertexAuthFuture, VertexConfig, VertexProviderLoader,
VertexTokenSource, get_vertex_ai_location, secret_names,
};
use rstest::rstest;
use serde_json::{Value, json};
struct FakeProvider {
calls: Arc<AtomicUsize>,
}
impl VertexTokenSource for FakeProvider {
fn project_id(&self) -> VertexAuthFuture<'_, String> {
self.calls.fetch_add(1, Ordering::SeqCst);
Box::pin(async { Ok("adc-project".into()) })
}
fn token(&self) -> VertexAuthFuture<'_, String> {
self.calls.fetch_add(1, Ordering::SeqCst);
Box::pin(async { Ok("adc-token".into()) })
}
}
struct FakeLoader {
loads: Arc<AtomicUsize>,
provider: Arc<dyn VertexTokenSource>,
}
impl VertexProviderLoader for FakeLoader {
fn load(&self, _source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
let loads = self.loads.clone();
let provider = self.provider.clone();
Box::pin(async move {
loads.fetch_add(1, Ordering::SeqCst);
Ok(provider)
})
}
}
fn config(value: Value) -> VertexConfig {
VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new())
.unwrap()
}
fn auth(calls: Arc<AtomicUsize>, loads: Arc<AtomicUsize>) -> VertexAuth {
let provider: Arc<dyn VertexTokenSource> = Arc::new(FakeProvider { calls });
VertexAuth::new(Arc::new(FakeLoader { loads, provider }))
}
#[rstest]
#[tokio::test]
async fn secret_names_cover_environment_reads() {
let seen = Arc::new(std::sync::Mutex::new(BTreeSet::<String>::new()));
let recorded = seen.clone();
let env = |name: &str| {
recorded.lock().unwrap().insert(name.to_string());
None
};
let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0)));
auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env)
.await
.unwrap();
get_vertex_ai_location(&VertexConfig::default(), &env);
assert!(
seen.lock()
.unwrap()
.iter()
.all(|name| secret_names().contains(&name.as_str()))
);
}
#[rstest]
#[tokio::test]
async fn explicit_token_and_header_do_not_acquire_adc() {
let loads = Arc::new(AtomicUsize::new(0));
let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone());
let configured = config(json!({"vertex_project":"project-1"}));
let explicit = auth
.validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None)
.await
.unwrap();
assert_eq!(explicit.headers[0].1, "Bearer access-token");
let existing = auth
.validate_environment(
vec![("authorization".into(), "Bearer existing".into())],
None,
&configured,
&|_| None,
)
.await
.unwrap();
assert_eq!(existing.headers[0].1, "Bearer existing");
assert_eq!(loads.load(Ordering::SeqCst), 0);
}
#[rstest]
#[tokio::test]
async fn provider_is_reused_across_authentication_calls() {
let calls = Arc::new(AtomicUsize::new(0));
let loads = Arc::new(AtomicUsize::new(0));
let auth = auth(calls.clone(), loads.clone());
for _ in 0..2 {
let environment = auth
.validate_environment(Vec::new(), None, &VertexConfig::default(), &|_| None)
.await
.unwrap();
assert_eq!(environment.project_id, "adc-project");
assert_eq!(environment.headers[0].1, "Bearer adc-token");
}
assert_eq!(loads.load(Ordering::SeqCst), 1);
assert_eq!(calls.load(Ordering::SeqCst), 4);
}

View file

@ -92,6 +92,11 @@ impl Error {
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum ErrorDetail {
#[error("{operation} returned HTTP {status}")]
UnexpectedStatus {
operation: &'static str,
status: u16,
},
#[error("{field} must be {expected}")]
InvalidType {
field: String,

View file

@ -53,6 +53,9 @@ pub struct AwsParams {
}
impl AwsParams {
pub const ROLE_ARN_ENV: &str = "AWS_ROLE_ARN";
pub const WEB_IDENTITY_TOKEN_FILE_ENV: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
pub const ACCESS_KEY_ID: ParamSpec = spec(
"access key id",
ParamKind::Credential,
@ -149,6 +152,10 @@ impl AwsParams {
AwsParams::SPECS
.iter()
.flat_map(|spec| spec.env.iter().copied())
.chain([
AwsParams::ROLE_ARN_ENV,
AwsParams::WEB_IDENTITY_TOKEN_FILE_ENV,
])
.fold(Vec::new(), |mut names, name| {
if !names.contains(&name) {
names.push(name);

View file

@ -3,21 +3,13 @@ use rstest::rstest;
use serde_json::{Map, Value};
#[rstest]
fn fields_name_every_param_once_and_in_declaration_order() {
fn every_named_param_roundtrips_without_loss() {
let filled: Map<String, Value> = AwsParams::fields()
.map(|name| (name.to_string(), Value::from(format!("value-of-{name}"))))
.collect();
let typed = AwsParams::from_optional_params(&filled);
let serialized = serde_json::to_value(&typed).unwrap();
assert_eq!(
serialized
.as_object()
.unwrap()
.keys()
.map(String::as_str)
.collect::<Vec<_>>(),
AwsParams::fields().collect::<Vec<_>>()
);
assert_eq!(filled.len(), AwsParams::fields().count());
assert_eq!(serialized, Value::Object(filled.clone()));
assert_eq!(
serde_json::from_value::<AwsParams>(Value::Object(filled)).unwrap(),

View file

@ -8,7 +8,7 @@ use litellm_cache::{
use litellm_http::Client;
use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode};
use crate::{GcpTokenSource, TokenSource};
use crate::{GcsTokenSource, TokenSource};
pub const DEFAULT_ENDPOINT: &str = "https://storage.googleapis.com";
@ -62,7 +62,7 @@ pub struct GcsCache<S: CacheCodec> {
impl<S: CacheCodec> GcsCache<S> {
pub fn new(config: GcsConfig, client: Client, codec: S) -> Self {
let token = Arc::new(GcpTokenSource::new(config.path_service_account.clone()));
let token = Arc::new(GcsTokenSource::new(config.path_service_account.clone()));
Self::with_token_source(config, client, codec, token)
}

View file

@ -2,4 +2,4 @@ mod cache;
mod token;
pub use cache::{DEFAULT_ENDPOINT, GcsCache, GcsConfig, key_prefix};
pub use token::{GcpTokenSource, StaticTokenSource, TokenSource};
pub use token::{GcsTokenSource, StaticTokenSource, TokenSource};

View file

@ -1,7 +1,6 @@
use std::{future::Future, pin::Pin};
use litellm_auth_gcp::{VertexAuth, VertexConfig};
use litellm_auth_types::{InputSource, Sourced};
use litellm_auth_gcp::{CredentialSource, GcpAuth};
use litellm_secret_value::SecretValue;
use litellm_cache::Error;
@ -10,28 +9,30 @@ pub trait TokenSource: Send + Sync + 'static {
fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>>;
}
pub struct GcpTokenSource {
auth: VertexAuth,
config: VertexConfig,
pub struct GcsTokenSource {
auth: GcpAuth,
credentials: Option<String>,
}
impl GcpTokenSource {
impl GcsTokenSource {
pub fn new(path_service_account: Option<String>) -> Self {
let credentials = path_service_account
.map(|path| Sourced::new(SecretValue::new(path), InputSource::Deployment));
Self {
auth: VertexAuth::default(),
config: VertexConfig::new(credentials, None, None),
auth: GcpAuth::default(),
credentials: path_service_account.filter(|path| !path.trim().is_empty()),
}
}
}
impl TokenSource for GcpTokenSource {
impl TokenSource for GcsTokenSource {
fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>> {
Box::pin(async move {
self.auth
.access_token(&self.config, &|name| std::env::var(name).ok())
.access_token(match &self.credentials {
Some(path) => CredentialSource::Trusted(SecretValue::new(path.clone())),
None => CredentialSource::application_default(&|name| std::env::var(name).ok()),
})
.await
.map(|token| token.expose().to_owned())
.map_err(Error::unavailable)
})
}

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)();

View file

@ -36,9 +36,12 @@ impl ProvideCredentials for S3Credentials {
self.config.secret_access_key.clone(),
) {
return Ok(AwsCredentials::new(
access_key_id,
secret_access_key,
self.config.session_token.clone(),
access_key_id.expose(),
secret_access_key.expose(),
self.config
.session_token
.as_ref()
.map(|value| value.expose().to_owned()),
None,
"litellm-s3-cache",
));

View file

@ -5,9 +5,9 @@ use rstest::rstest;
fn explicit(session_token: Option<&str>) -> AwsAuthConfig {
AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
session_token: session_token.map(str::to_string),
access_key_id: Some("key".into()),
secret_access_key: Some("secret".into()),
session_token: session_token.map(Into::into),
region_name: Some("us-east-1".to_string()),
..Default::default()
}

View file

@ -23,8 +23,8 @@ pub fn config(endpoint: &str) -> S3CacheConfig {
url: endpoint.to_string(),
}),
auth: AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
access_key_id: Some("key".into()),
secret_access_key: Some("secret".into()),
region_name: Some("us-east-1".to_string()),
..Default::default()
},

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`

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"] }

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,
},
}

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,

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(),
}
}

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(_))
));
}

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"));
}

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

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

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

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;
}
}

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()
}
}

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()
}
}

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()
}
}

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(())
}

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"));
}

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;
}

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),
}
}

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 => {}
}
}

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);
}

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]);
}

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")
);
}

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");
}

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());
}

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\"}"
);
}

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();
}

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}"
);
}

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");
}

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