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

merge: refresh detection on bounded analytics fixes

Browse source
This commit is contained in:
moe-berri 2026-11-06 08:52:49 +08:00
commit db66b0a3a4
501 changed files with 29020 additions and 47239 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

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

2
AGENTS.md
View file

@ -45,7 +45,7 @@ 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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

59
litellm-rust/Cargo.lock generated
View file

@ -547,7 +547,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 +736,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",
@ -1000,7 +1000,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",
@ -2255,7 +2255,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 +2391,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 +2833,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 +2860,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 +2889,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 +2904,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 +2923,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 +2942,7 @@ checksum = "986c5ce3b994526b3cd75578e62554abd09f0899d6206de48b3e96ab34ccc8c7"
dependencies = [
"hex",
"http-body-util",
"hyper 1.10.1",
"hyper 1.12.0",
"hyper-util",
"pin-project-lite",
"tokio",
@ -3769,6 +3769,7 @@ dependencies = [
"litellm-config",
"litellm-error",
"litellm-gateway-auth",
"litellm-gateway-forward",
"litellm-gateway-inference",
"litellm-gateway-mcp",
"litellm-gateway-ui",
@ -3815,6 +3816,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"
@ -6031,7 +6055,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",
@ -7512,6 +7536,7 @@ dependencies = [
"bytes",
"futures-core",
"futures-sink",
"futures-util",
"pin-project-lite",
"tokio",
]
@ -7591,7 +7616,7 @@ dependencies = [
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"hyper 1.10.1",
"hyper 1.12.0",
"hyper-timeout",
"hyper-util",
"percent-encoding",
@ -8167,9 +8192,9 @@ dependencies = [
[[package]]
name = "want"
version = "0.3.1"
version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e"
checksum = "ec4cdd0dd910afe868b7ef477227d8d538b46b3075031afee8a9f2acb0a2ed0b"
dependencies = [
"try-lock",
]
@ -8551,7 +8576,7 @@ dependencies = [
"futures",
"http 1.4.2",
"http-body-util",
"hyper 1.10.1",
"hyper 1.12.0",
"hyper-util",
"log",
"once_cell",

1
litellm-rust/Cargo.toml
View file

@ -26,6 +26,7 @@ litellm-inference-ocr = { path = "crates/inference-ocr" }
litellm-inference-testing = { path = "crates/inference-testing" }
litellm-gateway-mcp = { path = "crates/gateway-mcp" }
litellm-gateway = { path = "crates/gateway" }
litellm-gateway-forward = { path = "crates/gateway-forward" }
litellm-gateway-inference = { path = "crates/gateway-inference" }
litellm-gateway-auth = { path = "crates/gateway-auth" }
litellm-gateway-management = { path = "crates/gateway-management" }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@ -12,6 +12,7 @@ envy = "0.4.2"
http-body-util = "0.1"
litellm-inference.workspace = true
litellm-gateway-inference.workspace = true
litellm-gateway-forward.workspace = true
litellm-gateway-mcp.workspace = true
tokio-util = "0.7"
thiserror.workspace = true

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

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

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

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

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

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

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

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

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

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

210
litellm-rust/crates/router-core/src/call.rs
View file

@ -4,13 +4,7 @@
//! `should_retry_this_error` and `_time_to_sleep_before_retry` in `litellm/router.py`, and
//! `run_async_fallback` in `router_utils/fallback_event_handlers.py`.
use std::{
borrow::Cow,
future::Future,
pin::Pin,
sync::Arc,
time::{Duration, Instant},
};
use std::{borrow::Cow, future::Future, pin::Pin, sync::Arc, time::Duration};
use serde_json::{Map, Value};
@ -29,7 +23,7 @@ use crate::{
host_index::python_str,
lookup::Lookup,
mock::MockFailure,
outcome::OutcomeWrite,
outcome::{OutcomeCallbacks, OutcomeWrite},
prerouting::{PreRoute, PreRouted},
registry::{ResolveKind, ResolveScope},
retry::{RetryContext, resolve_retry_policy, should_retry},
@ -63,6 +57,12 @@ pub struct RouterResponse {
struct Failed {
error: Box<Error>,
deployment_id: Option<String>,
/// The group a pre-routing hook rewrote the level's name to for the failed attempt (an auto
/// router's tier), which its order fallbacks and fallback lookup key off
/// (`get_pre_routing_selection`).
routed: Option<String>,
/// The group the host stamped that attempt as served for ([`Candidates::served_as`]).
served_as: Option<String>,
}
impl From<Error> for Failed {
@ -70,6 +70,8 @@ impl From<Error> for Failed {
Self {
error: Box::new(error),
deployment_id: None,
routed: None,
served_as: None,
}
}
}
@ -103,6 +105,8 @@ struct Target {
/// `_failover_excluded_ids` in the level's metadata: the deployments weighted failover
/// already excluded on the way here.
failover_excluded: Arc<[String]>,
/// The group a pre-routing hook selected for the level's failed attempt, if it rewrote it.
routed: Option<String>,
}
/// A step of a level's fallback chain: a configured target, or an order level of the group.
@ -166,10 +170,18 @@ fn unattempted_chain(
target: &Target,
at: FallbackHop,
attempted: &[Tried],
served: (Option<&str>, Option<&str>),
) -> bool {
let mut lookup = vec![target.group.as_str()];
if request.model != target.group {
lookup.push(request.model.as_str());
let (routed, served_as) = served;
let mut lookup: Vec<&str> = Vec::new();
for group in routed
.into_iter()
.chain(served_as)
.chain([target.group.as_str(), request.model.as_str()])
{
if !lookup.contains(&group) {
lookup.push(group);
}
}
list.lookup_groups(&lookup, lists.keys)
.is_some_and(|targets| {
@ -285,6 +297,7 @@ impl Core {
params: None,
excluded: Arc::from([]),
failover_excluded: Arc::from([]),
routed: None,
};
let at = FallbackHop::default();
// A fallbacks mock stands in for the requested group's retry loop: no pick, no attempt.
@ -343,6 +356,8 @@ impl Core {
let Failed {
error,
deployment_id: failed_id,
routed,
served_as,
} = failed;
let error = *error;
self.level_failed(at, &target, &error, false);
@ -360,7 +375,14 @@ impl Core {
self.level_failed(at, &target, &raised, true);
return Err(raised);
}
let (targets, order_path) = self.fallback_targets(snapshot, request, &target, &error);
let (targets, order_path) = self.fallback_targets(
snapshot,
request,
&target,
routed.as_deref(),
served_as.as_deref(),
&error,
);
if !order_path {
self.host.observe(&RouterEvent::FallingBack {
level: at.level,
@ -368,6 +390,7 @@ impl Core {
});
}
let mut target = target;
target.routed = routed;
// `_maybe_run_weighted_failover`: its own `run_async_fallback`, before the regular one.
let failover = if order_path {
None
@ -463,7 +486,12 @@ impl Core {
if at.depth >= max_fallbacks {
return Err(None);
}
let failed = Tried::Target(FallbackTarget::Name(target.group.clone()));
let failed = Tried::Target(FallbackTarget::Name(
target
.routed
.clone()
.unwrap_or_else(|| target.group.clone()),
));
if !chain.attempted.contains(&failed) {
chain.attempted.push(failed);
}
@ -488,17 +516,34 @@ impl Core {
let mut unroutable: Option<Value> = None;
for next in targets {
let (tried, group, order) = match &next {
Next::Order(group, order) => (
Tried::Order(group.clone(), *order),
group.clone(),
Some(*order),
),
Next::Order(group, order) => {
// An order level of the group a pre-routing hook selected passes the
// fallback gates as any group the level did not serve does.
if *group != target.group {
match self.host.allow_fallback(request, group).await {
Ok(true) => {}
Ok(false) => continue,
Err(raised) => return Err(Some(Error::Attempt(raised))),
}
}
(
Tried::Order(group.clone(), *order),
group.clone(),
Some(*order),
)
}
Next::Failover(group, excluded) => {
let mut key = excluded.to_vec();
key.sort();
(Tried::Failover(group.clone(), key), group.clone(), None)
}
Next::Fallback(FallbackTarget::Name(name)) if *name == target.group => continue,
// The group that just failed: the pre-routing selection when there was one, so
// a subagent bound to another tier can fall back to the model it asked for.
Next::Fallback(FallbackTarget::Name(name))
if *name == target.routed.as_deref().unwrap_or(&target.group) =>
{
continue;
}
// A provider-scoped request never leaves its group; a dict target with no
// `model` names no group, so it is skipped too.
Next::Fallback(fallback)
@ -573,6 +618,7 @@ impl Core {
params: params.clone(),
excluded,
failover_excluded: target.failover_excluded.clone(),
routed: None,
};
let parent = at;
let at = FallbackHop {
@ -704,6 +750,8 @@ impl Core {
snapshot: &RouterStateSnapshot,
request: &RouterRequest,
target: &Target,
routed: Option<&str>,
served_as: Option<&str>,
error: &Error,
) -> (Option<Vec<Next>>, bool) {
// A pick's error decides by its exception's classes as an attempt's does: a pre-call
@ -711,11 +759,19 @@ impl Core {
let failure = retried_failure(error);
let is_a = |name: &str| failure.as_ref().is_some_and(|f| f.is_a(name));
let lists = Lists::of(snapshot, request);
// `fallback_lookup_groups`: the group being served, then the group the request was for.
let mut lookup = vec![target.group.as_str()];
if request.model != target.group {
lookup.push(request.model.as_str());
// `fallback_lookup_groups`: the group a pre-routing hook selected, the group being served,
// then the group the request was for.
let mut lookup: Vec<&str> = routed.into_iter().collect();
for group in served_as
.into_iter()
.chain([target.group.as_str(), request.model.as_str()])
{
if !lookup.contains(&group) {
lookup.push(group);
}
}
// `order_model_group`: the pre-routing selection's order levels, else the level's own.
let order_group = routed.unwrap_or(target.group.as_str());
let chain = |list: &Fallbacks| {
if list.is_client_side(lists.keys) {
Some(list.as_chain())
@ -734,7 +790,7 @@ impl Core {
if !skip_order {
// A routing group's members, with a priority group's priorities as their order.
let listing_team = request.listing_team_id.as_deref();
let mut orders: Vec<i64> = group_members(snapshot, &target.group, listing_team)
let mut orders: Vec<i64> = group_members(snapshot, order_group, listing_team)
.iter()
.filter(|d| d.listed_for_team(listing_team))
.filter_map(|d| d.order())
@ -742,11 +798,14 @@ impl Core {
orders.sort_unstable();
orders.dedup();
if orders.len() > 1 {
let current = target.order.unwrap_or(orders[0]);
let current = target
.order
.filter(|_| routed.is_none())
.unwrap_or(orders[0]);
let mut targets: Vec<Next> = orders
.into_iter()
.filter(|order| *order > current)
.map(|order| Next::Order(target.group.clone(), order))
.map(|order| Next::Order(order_group.to_owned(), order))
.collect();
if let Some(list) = lists.fallbacks {
targets.extend(fallbacks_of(chain(list).unwrap_or_default()));
@ -879,6 +938,8 @@ impl Core {
return Err(Failed {
error: Box::new(Error::Attempt(AttemptError::host_raised(&["Exception"]))),
deployment_id: None,
routed: None,
served_as: None,
});
}
let Some(first_failure) = retried_failure(&failed.error) else {
@ -965,6 +1026,8 @@ impl Core {
return Err(Failed {
error: Box::new(Error::Attempt(AttemptError::host_raised(&["Exception"]))),
deployment_id: None,
routed: None,
served_as: None,
});
}
if num_retries == 0 {
@ -1162,7 +1225,7 @@ impl Core {
.unrouted_attempt(snapshot, op, request, target, (retry, max_retries), at)
.await;
}
let pick_started = Instant::now();
let pick_started = self.health.now_time();
let (deployment, candidates) = match self
.picked(snapshot, op, request, target, inputs, retry, at)
.await
@ -1179,7 +1242,11 @@ impl Core {
self.host.observe(&RouterEvent::Selected {
model_group: candidates.model_group.clone(),
deployment_id: deployment.id.clone(),
pick: pick_started.elapsed(),
pick: self
.health
.now_time()
.duration_since(pick_started)
.unwrap_or_default(),
});
if !self.host.supports(op, deployment.provider()) {
return Err(Error::Config(ConfigError::Unsupported(Unsupported::Op {
@ -1205,7 +1272,7 @@ impl Core {
},
None,
snapshot.settings.dynamic.default_ref(),
jiff::Timestamp::now(),
self.health.timestamp(),
),
Ok(None)
);
@ -1219,6 +1286,7 @@ impl Core {
} else {
candidates.model_group.clone()
},
routed: routed_name(target, &candidates),
target_order: target.order,
candidate_ids: candidates.all.iter().map(|d| d.id.clone()).collect(),
retry,
@ -1236,23 +1304,43 @@ impl Core {
// the stream's end.
match self.host.invoke(&attempt).await {
Ok(output) => {
// `_should_raise_content_policy_error`: a blocked response is raised only when
// a content-policy fallback would serve it, and counts as no failure.
let routed = routed_name(target, &candidates);
let served = (routed.as_deref(), candidates.served_as.as_deref());
let raised = match output.blocked {
Some(Blocked::ContentFilter) => self.has_content_policy_fallback(
snapshot, request, target, at, attempted, served,
),
Some(Blocked::Refusal) => self.refusal_fallback_available(
snapshot, request, target, at, attempted, served,
),
None => false,
};
let tracks = output
.tracks_usage
.unwrap_or_else(|| deployment.tracks_usage());
// A shared logging object logs the success of the response that serves the
// request alone, so a raised blocked response counts its request and no success.
let callbacks = match raised && request.outcome_callbacks.shared_logging {
true => OutcomeCallbacks {
success: false,
..request.outcome_callbacks
},
false => request.outcome_callbacks,
};
let recorded = match (output.deferred, tracks) {
(true, true) => {
let tokens = output.total_tokens as f64;
self.health.add_usage(&deployment, None, tokens, 1.0).await
self.health
.add_usage(&deployment, None, tokens, 1.0)
.await
.map(|_| ())
}
(true, false) => Ok(()),
(false, _) => {
self.health
.record_success(
request.outcome_callbacks,
&deployment,
output.total_tokens,
tracks,
)
.record_success(callbacks, &deployment, output.total_tokens, tracks)
.await
}
};
@ -1262,16 +1350,6 @@ impl Core {
self.host.observe(&RouterEvent::Succeeded {
deployment_id: deployment.id.clone(),
});
// `_should_raise_content_policy_error`: a blocked response is raised only when
// a content-policy fallback would serve it, and counts as no failure.
let raised =
match output.blocked {
Some(Blocked::ContentFilter) => self
.has_content_policy_fallback(snapshot, request, target, at, attempted),
Some(Blocked::Refusal) => self
.refusal_fallback_available(snapshot, request, target, at, attempted),
None => false,
};
if let Some(blocked) = output.blocked.filter(|_| raised) {
self.host.observe(&RouterEvent::Blocked {
level: at.level,
@ -1284,6 +1362,8 @@ impl Core {
return Err(Failed {
error: Box::new(Error::Attempt(failure)),
deployment_id: Some(deployment.id.clone()),
routed: routed_name(target, &candidates),
served_as: candidates.served_as.clone(),
});
}
Ok(RouterResponse {
@ -1303,6 +1383,8 @@ impl Core {
Err(Failed {
error: Box::new(Error::Attempt(failure)),
deployment_id: Some(deployment.id.clone()),
routed: routed_name(target, &candidates),
served_as: candidates.served_as.clone(),
})
}
}
@ -1325,6 +1407,9 @@ impl Core {
&snapshot.settings.routing_strategy,
request,
)?;
self.host.observe(&RouterEvent::PickStarted {
model: target.group.clone(),
});
let routed = self
.pre_routed(PreRoute {
op,
@ -1334,22 +1419,25 @@ impl Core {
retry,
})
.await?;
let (routed, served_as) = match routed {
PreRouted::Bound { model, group } => (PreRouted::Model(model), Some(group)),
other => (other, None),
};
match routed {
PreRouted::Model(name) => {
snapshot.missing_routing_group(request, &name)?;
let candidates = self
let mut candidates = self
.candidates(snapshot, op, request, &name, inputs)
.await?;
candidates.served_as = served_as;
if candidates.healthy.is_empty() {
return Err(Error::NoHealthyDeployments {
model: candidates.failed_model(&snapshot.registry),
}
.into());
return Err(candidates.none_healthy(&snapshot.registry).into());
}
let deployment = self.pick(snapshot, op, request, &candidates).await?;
Ok((deployment, candidates))
}
// A custom strategy's pick is served as is, for the requested name.
PreRouted::Bound { .. } => unreachable!("mapped to a model above"),
PreRouted::Deployment(deployment) => {
let candidates = Candidates {
name: target.group.clone(),
@ -1358,6 +1446,7 @@ impl Core {
all: vec![deployment.clone()],
healthy: vec![deployment.clone()],
wildcard_request: None,
served_as: None,
};
Ok((deployment, candidates))
}
@ -1383,6 +1472,7 @@ impl Core {
deployment,
name: target.group.clone(),
model_group: target.group.clone(),
routed: None,
target_order: target.order,
candidate_ids: Vec::new(),
retry,
@ -1417,10 +1507,11 @@ impl Core {
target: &Target,
at: FallbackHop,
attempted: &[Tried],
served: (Option<&str>, Option<&str>),
) -> bool {
let lists = Lists::of(snapshot, request);
match lists.content_policy {
Some(list) => unattempted_chain(list, &lists, request, target, at, attempted),
Some(list) => unattempted_chain(list, &lists, request, target, at, attempted, served),
None => snapshot.settings.dynamic.fallbacks.has_star(),
}
}
@ -1434,20 +1525,22 @@ impl Core {
target: &Target,
at: FallbackHop,
attempted: &[Tried],
served: (Option<&str>, Option<&str>),
) -> bool {
if request.disable_fallbacks {
return false;
}
let lists = Lists::of(snapshot, request);
if lists.content_policy.is_some() {
return self.has_content_policy_fallback(snapshot, request, target, at, attempted);
return self
.has_content_policy_fallback(snapshot, request, target, at, attempted, served);
}
if snapshot.settings.dynamic.fallbacks.has_star() {
return true;
}
lists
.fallbacks
.is_some_and(|list| unattempted_chain(list, &lists, request, target, at, attempted))
lists.fallbacks.is_some_and(|list| {
unattempted_chain(list, &lists, request, target, at, attempted, served)
})
}
/// `async_deployment_callback_on_failure`, when it runs for this failure.
@ -1553,3 +1646,8 @@ impl Core {
Ok(())
}
}
/// The name a pre-routing hook rewrote the level's group to for this pick, if it did.
fn routed_name(target: &Target, candidates: &Candidates) -> Option<String> {
(candidates.name != target.group).then(|| candidates.name.clone())
}

22
litellm-rust/crates/router-core/src/call/stream.rs
View file

@ -37,6 +37,9 @@ pub struct StreamOrigin {
pub attempted_retries: u32,
/// The budget the first stream retry committed to (`committed_retry_budget_for_request`).
pub committed_budget: Option<u32>,
/// The group a pre-routing hook selected for `group`, which the fallback lookup keys on first.
#[serde(default)]
pub routed: Option<String>,
}
/// What may still take over a stream that fails before content.
@ -61,14 +64,22 @@ impl StreamOrigin {
params: None,
excluded: Arc::from([]),
failover_excluded: Arc::from([]),
routed: self.routed.clone(),
}
}
/// `fallback_lookup_groups`: the served group, then the group the request was for.
fn lookup<'a>(&'a self, request: &'a RouterRequest) -> Vec<&'a str> {
let mut lookup = vec![self.group.as_str()];
if request.model != self.group {
lookup.push(request.model.as_str());
let mut lookup: Vec<&str> = Vec::new();
for group in self
.routed
.as_deref()
.into_iter()
.chain([self.group.as_str(), request.model.as_str()])
{
if !lookup.contains(&group) {
lookup.push(group);
}
}
lookup
}
@ -186,7 +197,10 @@ impl Router {
scope: 0,
},
&mut chain,
Failed::from(Error::Attempt(failure)),
Failed {
routed: origin.routed.clone(),
..Failed::from(Error::Attempt(failure))
},
)
.await
}

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

@ -40,7 +40,13 @@ pub enum Error {
#[error("no deployment serves model `{model}`")]
UnknownModel { model: String },
#[error("no deployment of `{model}` is available")]
NoHealthyDeployments { model: String },
NoHealthyDeployments {
model: String,
/// The ids of the deployments the request could route to before the cooldown filter
/// (`request_deployments`), when the candidate list came up empty; `None` when the
/// strategy refused every candidate.
request_ids: Option<Vec<String>>,
},
#[error("the attempt failed")]
Attempt(#[from] AttemptError),
#[error("the routing cache failed")]
@ -201,8 +207,9 @@ impl Error {
Self::UnknownModel { model } => Self::UnknownModel {
model: model.clone(),
},
Self::NoHealthyDeployments { model } => Self::NoHealthyDeployments {
Self::NoHealthyDeployments { model, request_ids } => Self::NoHealthyDeployments {
model: model.clone(),
request_ids: request_ids.clone(),
},
Self::Paused { model } => Self::Paused {
model: model.clone(),

28
litellm-rust/crates/router-core/src/health.rs
View file

@ -247,6 +247,10 @@ impl Health {
self.clock.now()
}
pub(crate) fn timestamp(&self) -> jiff::Timestamp {
jiff::Timestamp::try_from(self.clock.now()).unwrap_or(jiff::Timestamp::UNIX_EPOCH)
}
fn now(&self) -> f64 {
unix_seconds(self.clock.now())
}
@ -554,35 +558,35 @@ impl Health {
return Ok(());
}
callbacks
.record(
OutcomeWrite::AttemptUsage,
self.add_usage(deployment, None, total_tokens as f64, 1.0),
)
.record(OutcomeWrite::AttemptUsage, async {
self.add_usage(deployment, None, total_tokens as f64, 1.0)
.await
.map(|_| ())
})
.await
}
/// `_increment_deployment_usage`, for a deployment that tracks usage: tokens and requests
/// under `model` (metadata's `deployment`, which a wildcard route fills; else the
/// deployment's own), unless there is nothing to add.
/// deployment's own), unless there is nothing to add. Returns the tokens counter's key when
/// it added to it.
pub async fn add_usage(
&self,
deployment: &Deployment,
model: Option<&str>,
total_tokens: f64,
requests: f64,
) -> Result<(), Error> {
) -> Result<Option<String>, Error> {
if total_tokens <= 0.0 && requests <= 0.0 {
return Ok(());
return Ok(None);
}
let minute = current_minute(self.clock.now());
let model = model.unwrap_or(&deployment.params.model);
let tpm_key = usage_key(&deployment.id, model, UsageMetric::Tpm, &minute);
self.cache
.async_increment_usage(
&[
(
usage_key(&deployment.id, model, UsageMetric::Tpm, &minute),
total_tokens,
),
(tpm_key.clone(), total_tokens),
(
usage_key(&deployment.id, model, UsageMetric::Rpm, &minute),
requests,
@ -591,7 +595,7 @@ impl Health {
Some(USAGE_TTL),
)
.await?;
Ok(())
Ok(Some(tpm_key))
}
/// `async_deployment_callback_on_failure`: a failed request counts against the deployment's

21
litellm-rust/crates/router-core/src/host.rs
View file

@ -83,6 +83,13 @@ pub struct RouterRequest {
/// A Decisions call carrying audio or videos: the deployments that refuse them are skipped
/// unless that leaves nothing (`_without_decisions_deployments_refusing_media`).
pub decisions_media: bool,
/// A batch create (`batch_create_routing`): when the caller may use a deployment declared
/// `model_info.mode: batch`, the Azure deployments that are not are skipped
/// (`without_deployments_that_cannot_batch`).
pub batch_create: bool,
/// An embedding pick: the host keeps the candidates whose `max_embedding_batch_size` admits
/// the input (`deployments_admitting_embedding_batch`).
pub embedding_batch: bool,
/// `f"{provider}/{model}"` for a lookup that cannot ask the host
/// (`_core_get_routable_deployments`); a pick asks [`Host::provider_request`].
pub wildcard_request: Option<String>,
@ -291,6 +298,9 @@ pub struct Attempt {
pub name: String,
/// The model group being served, after fallbacks.
pub model_group: String,
/// The group a pre-routing hook selected for `name`, when it differs (the tier stamp the
/// fallback lookup keys on first).
pub routed: Option<String>,
/// The `order` level an order fallback targets in the same group.
pub target_order: Option<i64>,
/// Every deployment `name` resolved to for this caller, the picked one included.
@ -383,11 +393,22 @@ pub enum SelectStage {
/// `_decisions_deployment_refuses_media`: return the candidates that refuse the audio or
/// videos a Decisions call carries.
MediaRefusing,
/// `deployments_admitting_embedding_batch`: return the candidates that admit the embedding
/// input, or what the pick keeps when none does (`deployments_when_none_admits_embedding_batch`).
EmbeddingBatch,
/// `_async_override_selector_pre_call_check`: a deployment named by id under a per-request
/// strategy override still passes that strategy's pre-call check (rpm accounting); return the
/// candidates, or raise its refusal.
OverrideCheck(RoutingStrategy),
}
/// Routing events, delivered through [`Host::observe`].
#[derive(Clone, Debug, PartialEq)]
pub enum RouterEvent {
/// A pick for `model` (the level's group) begins: its route phase opens.
PickStarted {
model: String,
},
Selected {
model_group: String,
deployment_id: String,

36
litellm-rust/crates/router-core/src/lib.rs
View file

@ -366,7 +366,7 @@ impl Router {
})
.await?;
let name = match routed {
PreRouted::Model(name) => name,
PreRouted::Model(name) | PreRouted::Bound { model: name, .. } => name,
PreRouted::Deployment(deployment) => return Ok(deployment),
};
snapshot.missing_routing_group(request, &name)?;
@ -374,6 +374,10 @@ impl Router {
.core
.candidates(&snapshot, op, request, &name, &SelectInputs::default())
.await?;
// An empty candidate list names what the request could route to before the cooldown filter.
if candidates.healthy.is_empty() && !candidates.kind.is_single() {
return Err(candidates.none_healthy(&snapshot.registry));
}
self.core.pick(&snapshot, op, request, &candidates).await
}
@ -398,9 +402,7 @@ impl Router {
.await?;
// `_core_async_get_healthy_deployments` raises when nothing is left.
if candidates.healthy.is_empty() {
return Err(Error::NoHealthyDeployments {
model: candidates.failed_model(&snapshot.registry),
});
return Err(candidates.none_healthy(&snapshot.registry));
}
Ok(candidates)
}
@ -434,7 +436,7 @@ impl Router {
/// A success a logging callback saw (`deployment_callback_on_success`): counted for a
/// configured deployment, with its tokens and requests when it tracks usage. Returns the
/// success counter's key, or `None` for an unknown deployment.
/// tokens counter's key when it added usage (`_increment_deployment_usage`), else `None`.
pub async fn report_success(
&self,
id: &str,
@ -462,22 +464,22 @@ impl Router {
let Some(deployment) = snapshot.registry.deployment(id) else {
return Ok(None);
};
let key = self.core.health.count_success(id).await?;
if report
self.core.health.count_success(id).await?;
if !report
.tracks_usage
.unwrap_or_else(|| deployment.tracks_usage())
{
self.core
.health
.add_usage(
&deployment,
report.usage_model.as_deref(),
total_tokens,
requests,
)
.await?;
return Ok(None);
}
Ok(Some(key))
self.core
.health
.add_usage(
&deployment,
report.usage_model.as_deref(),
total_tokens,
requests,
)
.await
}
/// `sync_deployment_callback_on_success`: count a success against `id`, configured or not,

4
litellm-rust/crates/router-core/src/lookup.rs
View file

@ -433,6 +433,10 @@ impl<'a> Lookup<'a> {
.collect()
}
pub fn has_team_wildcards(&self, team_id: &str) -> bool {
self.host.team_patterns.iter().any(|(t, _)| t == team_id)
}
/// The deployments `team_id`'s own wildcard pattern routes `name` to
/// (`_core_team_pattern_routers`, keyed by `team_public_model_name`).
pub fn team_wildcard(&self, name: &str, team_id: &str) -> Vec<Arc<Deployment>> {

9
litellm-rust/crates/router-core/src/outcome.rs
View file

@ -28,6 +28,10 @@ pub struct OutcomeCallbacks {
pub failure: bool,
/// The attempt runs through `make_call`, which counts its usage whatever the callbacks do.
pub make_call: bool,
/// The request's attempts log on one logging object the caller passed (`litellm_logging_obj`),
/// which logs one success a request, the response that serves it: an attempt whose blocked
/// response is raised for a fallback counts no success.
pub shared_logging: bool,
}
impl Default for OutcomeCallbacks {
@ -36,6 +40,7 @@ impl Default for OutcomeCallbacks {
success: true,
failure: true,
make_call: true,
shared_logging: false,
}
}
}
@ -78,21 +83,25 @@ mod tests {
success: true,
failure: true,
make_call: true,
shared_logging: false,
};
const UNBOUND: OutcomeCallbacks = OutcomeCallbacks {
success: false,
failure: false,
make_call: true,
shared_logging: false,
};
const UNBOUND_WITHOUT_MAKE_CALL: OutcomeCallbacks = OutcomeCallbacks {
success: false,
failure: false,
make_call: false,
shared_logging: false,
};
const BOUND_SYNC: OutcomeCallbacks = OutcomeCallbacks {
success: false,
failure: true,
make_call: true,
shared_logging: false,
};
#[rstest]

3
litellm-rust/crates/router-core/src/prerouting.rs
View file

@ -33,6 +33,9 @@ pub struct PreRoute<'a> {
pub enum PreRouted {
/// Resolve this model group instead of the requested name.
Model(String),
/// Resolve `model`, served as `group`: the group a session binding stamped on the request,
/// which the level's fallbacks also look up.
Bound { model: String, group: String },
/// A custom strategy's pick: serve this deployment, without resolving or filtering.
Deployment(Arc<Deployment>),
}

17
litellm-rust/crates/router-core/src/registry.rs
View file

@ -234,6 +234,23 @@ impl Deployment {
/// `_deployment_supports_web_search`: `model_info.supports_web_search` when it is set, read
/// for truthiness, else true.
/// `is_batch_capable`: declared `model_info.mode: batch`.
pub(crate) fn is_batch_capable(&self) -> bool {
self.model_info.raw.get("mode").and_then(Value::as_str) == Some("batch")
}
/// `_cannot_run_batches`: an Azure (or Azure AI) deployment not declared batch-capable, by
/// its `custom_llm_provider`, else its model's provider prefix.
pub(crate) fn cannot_run_batches(&self) -> bool {
let provider = self.params.custom_llm_provider.as_deref().or_else(|| {
self.params
.model
.split_once('/')
.map(|(provider, _)| provider)
});
matches!(provider, Some("azure" | "azure_ai")) && !self.is_batch_capable()
}
pub(crate) fn supports_web_search(&self) -> bool {
self.model_info
.raw

80
litellm-rust/crates/router-core/src/resolve.rs
View file

@ -223,14 +223,20 @@ pub(crate) fn wildcard_name<'a>(
registry: &'a Registry,
name: &'a str,
specific_deployment: bool,
team_id: Option<&str>,
) -> Option<&'a str> {
let model = registry.alias_of(name).unwrap_or(name);
(registry.has_wildcards()
let lookup = crate::lookup::Lookup::new(registry);
let team_patterns = team_id.is_some_and(|team| lookup.has_team_wildcards(team));
let team_misses = || team_id.is_none_or(|team| lookup.team_wildcard(model, team).is_empty());
let global_misses = || registry.has_wildcards() && registry.wildcard_for(model).is_none();
((team_patterns || registry.has_wildcards())
&& !specific_deployment
&& !registry.has_name(model)
&& (registry.has_name(name) || registry.deployment(name).is_none())
&& registry.routing_group_named(model).is_none()
&& registry.wildcard_for(model).is_none())
&& team_misses()
&& (team_patterns || global_misses()))
.then_some(model)
}
@ -339,15 +345,21 @@ pub fn resolve_request(
let kept = reserved(model, deployments, caller, now)?;
return resolved(ResolveKind::Group, model, kept);
}
// `_core_wildcard_route_deployments`: the team's own pattern, else the proxy-wide one.
let team_owned = team_id
.map(|team| crate::lookup::Lookup::new(registry).team_wildcard(model, team))
.filter(|deployments| !deployments.is_empty());
if let Some(deployments) = team_owned {
// `_core_wildcard_route_deployments`: the team's own pattern, else the proxy-wide one,
// each tried as spelled and then under the name's provider (`get_pattern`).
let team_owned = team_id.and_then(|team| {
let lookup = crate::lookup::Lookup::new(registry);
[Some(model), caller.wildcard_request]
.into_iter()
.flatten()
.map(|request| (request, lookup.team_wildcard(request, team)))
.find(|(_, deployments)| !deployments.is_empty())
});
if let Some((request, deployments)) = team_owned {
let routed = deployments.iter().flat_map(|deployment| {
let public = deployment.model_info.team_public_model_name.as_deref();
let pattern = public.unwrap_or(model);
routed_wildcard(pattern, model, std::slice::from_ref(deployment))
let pattern = public.unwrap_or(request);
routed_wildcard(pattern, request, std::slice::from_ref(deployment))
});
let kept = reserved(model, routed.collect(), caller, now)?;
return resolved(ResolveKind::Wildcard, model, kept);
@ -428,6 +440,7 @@ pub fn resolve_request(
mod tests {
use std::collections::BTreeMap;
use rstest::rstest;
use serde_json::{Value, json};
use super::*;
@ -544,8 +557,11 @@ mod tests {
#[test]
fn a_name_no_pattern_matches_is_tried_under_its_provider() {
let registry = registry(wildcard_group());
assert_eq!(wildcard_name(&registry, "gpt-4o", false), Some("gpt-4o"));
assert_eq!(wildcard_name(&registry, "openai/gpt-4o", false), None);
assert_eq!(
wildcard_name(&registry, "gpt-4o", false, None),
Some("gpt-4o")
);
assert_eq!(wildcard_name(&registry, "openai/gpt-4o", false, None), None);
let caller = Caller {
wildcard_request: Some("openai/gpt-4o"),
..Caller::default()
@ -564,6 +580,48 @@ mod tests {
assert_eq!(resolution.deployments[0].params.model, "openai/gpt-4o");
}
#[rstest]
fn a_bare_name_is_tried_under_its_provider_against_the_teams_own_wildcard() {
let registry = registry(json!([
{"model_name": "model_name_tw_1", "litellm_params": {"model": "openai/*"},
"model_info": {"id": "tw", "team_id": "TW", "team_public_model_name": "openai/*"}},
{"model_name": "g", "litellm_params": {"model": "openai/g"}, "model_info": {"id": "g"}},
]));
// No proxy-wide wildcard: only the team's patterns call for the provider spelling.
assert_eq!(
wildcard_name(&registry, "gpt-4o-mini", false, Some("TW")),
Some("gpt-4o-mini")
);
assert_eq!(wildcard_name(&registry, "gpt-4o-mini", false, None), None);
assert_eq!(
wildcard_name(&registry, "openai/gpt-4o-mini", false, Some("TW")),
None
);
let resolve = |team_id| {
let caller = Caller {
team_id,
wildcard_request: Some("openai/gpt-4o-mini"),
..Caller::default()
};
resolve_request(
&registry,
"gpt-4o-mini",
false,
caller,
None,
None,
Timestamp::now(),
)
.unwrap()
};
let resolution = resolve(Some("TW")).unwrap();
assert_eq!(resolution.kind, ResolveKind::Wildcard);
assert_eq!(resolution.group, "gpt-4o-mini");
assert_eq!(ids(Some(resolution.clone())), ["tw"]);
assert_eq!(resolution.deployments[0].params.model, "openai/gpt-4o-mini");
assert!(resolve(Some("OTHER")).is_none());
}
#[test]
fn a_deployment_named_by_id_refuses_the_endpoint_like_any_other() {
let registry = registry(json!([

90
litellm-rust/crates/router-core/src/select.rs
View file

@ -16,7 +16,7 @@ use crate::{
host::{RouterRequest, RoutingContext, SelectStage},
registry::{Deployment, Registry, ResolveKind},
resolve::{Caller, resolve_request},
strategy::{StrategyScope, routing_context},
strategy::{StrategyScope, request_override, routing_context},
};
/// Randomness for the shuffle and retry jitter, injectable so tests are deterministic.
@ -77,11 +77,23 @@ pub struct Candidates {
/// The provider-prefixed spelling the name resolved under (`get_pattern`), when the host
/// was asked for one.
pub wildcard_request: Option<String>,
/// The group the host stamped the pick as served for when a pre-routing hook rewrote the name
/// (a Claude Code subagent's bound router): a fallback lookup key.
pub served_as: Option<String>,
}
impl Candidates {
/// The `model` a no-deployments error names: what `_common_checks_available_deployment`
/// returned, so the alias-resolved name a wildcard serves and not its pattern.
/// The error an empty candidate list raises: it names the deployments the request could
/// route to before the cooldown filter.
pub(crate) fn none_healthy(&self, registry: &Registry) -> Error {
Error::NoHealthyDeployments {
model: self.failed_model(registry),
request_ids: Some(self.all.iter().map(|d| d.id.clone()).collect()),
}
}
pub(crate) fn failed_model(&self, registry: &Registry) -> String {
let model = match self.kind {
ResolveKind::Wildcard => registry.alias_of(&self.name).unwrap_or(&self.name),
@ -269,8 +281,12 @@ impl Core {
// No deployment for the name: route to the first default fallback.
let default_target = crate::resolve::default_target(&snapshot.settings.dynamic);
// `get_pattern`: a name no pattern matches is tried again under its inferred provider.
let wildcard_name =
crate::resolve::wildcard_name(registry, name, request.specific_deployment);
let wildcard_name = crate::resolve::wildcard_name(
registry,
name,
request.specific_deployment,
request.team_id.as_deref(),
);
let wildcard_request = match wildcard_name {
Some(model) => {
let scope = StrategyScope::Default;
@ -296,7 +312,7 @@ impl Core {
caller,
default_target,
snapshot.settings.dynamic.default_ref(),
jiff::Timestamp::now(),
self.health.timestamp(),
)?
.ok_or_else(|| Error::UnknownModel {
// The name `_common_checks_available_deployment` looked for, past the alias.
@ -466,6 +482,7 @@ impl Core {
all,
healthy,
wildcard_request,
served_as: None,
});
}
@ -503,6 +520,7 @@ impl Core {
all,
healthy,
wildcard_request,
served_as: None,
});
}
@ -602,6 +620,34 @@ impl Core {
}
}
// `deployments_admitting_embedding_batch`: an embedding pick keeps the deployments whose
// input limit admits the batch; when none does, the host keeps what the pick waits for or
// the 400 names.
if request.embedding_batch && !resolution.kind.is_single() && !healthy.is_empty() {
let mut ctx = self.context(
snapshot,
op,
request,
&resolution.group,
name,
StrategyScope::Default,
);
ctx.pool = &all;
healthy = self
.host
.select(&SelectStage::EmbeddingBatch, &ctx, healthy)
.await?;
}
// `without_deployments_that_cannot_batch`: a batch create on a group the caller may use a
// batch deployment of skips the Azure deployments that cannot run batches.
if request.batch_create
&& !resolution.kind.is_single()
&& healthy.iter().any(|d| d.is_batch_capable())
{
healthy.retain(|d| !d.cannot_run_batches());
}
if let Some(pinned) = request.pinned_deployment_id.as_deref()
&& healthy.iter().any(|d| d.id == pinned)
{
@ -653,6 +699,7 @@ impl Core {
all,
healthy,
wildcard_request,
served_as: None,
})
}
@ -691,14 +738,43 @@ impl Core {
) -> Result<Arc<Deployment>, Error> {
let no_deployment = || Error::NoHealthyDeployments {
model: candidates.failed_model(&snapshot.registry),
request_ids: None,
};
// A deployment named by its id is returned as is, before any strategy runs.
// A deployment named by its id is returned as is, before any strategy runs, past a
// per-request strategy override's pre-call check.
if candidates.kind.is_single() {
return candidates
let named = candidates
.healthy
.first()
.cloned()
.ok_or_else(no_deployment);
.ok_or_else(no_deployment)?;
if let Some(strategy) = request_override(request.routing_strategy.as_deref())
&& !matches!(
strategy,
RoutingStrategy::SimpleShuffle | RoutingStrategy::Unset
)
{
let ctx = self.context(
snapshot,
op,
request,
&candidates.model_group,
&candidates.name,
StrategyScope::Override,
);
let kept = self
.host
.select(
&SelectStage::OverrideCheck(strategy),
&ctx,
vec![named.clone()],
)
.await?;
if kept.is_empty() {
return Err(no_deployment());
}
}
return Ok(named);
}
let (strategy, scope) = routing_context(
&snapshot.registry,

54
litellm-rust/crates/router-core/tests/call_loop.rs
View file

@ -1755,3 +1755,57 @@ async fn a_listed_pass_through_pick_names_the_group_an_alias_resolves_to(config:
Error::NoPassThroughDeployments { model } if model == "primary"
));
}
struct StoppedClock(u64);
impl litellm_router_core::health::Clock for StoppedClock {
fn now(&self) -> std::time::SystemTime {
std::time::UNIX_EPOCH + Duration::from_secs(self.0)
}
}
fn utc_hour(hour: u64) -> u64 {
1_767_225_600 + hour * 3_600
}
#[rstest]
#[case::inside_the_window(utc_hour(10), vec!["open"])]
#[case::outside_the_window(utc_hour(20), vec!["open", "reserved"])]
#[tokio::test]
async fn access_windows_follow_the_routers_clock(#[case] now: u64, #[case] expected: Vec<&str>) {
let config: RouterConfig = serde_json::from_value(json!({
"model_list": [
{"model_name": "chat", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "open"}},
{"model_name": "chat", "litellm_params": {"model": "openai/b"}, "model_info": {
"id": "reserved",
"access_windows": [{"start": "09:00", "end": "17:00", "timezone": "UTC", "team_ids": ["batch"]}],
}},
]
}))
.unwrap();
let router = Router::builder(
config,
RouterDynamicConfig::default(),
Arc::new(ScriptedHost::default()),
)
.randomness(Arc::new(SeededRandomness::new(7)))
.clock(Arc::new(StoppedClock(now)))
.build()
.unwrap();
let mut picked = std::collections::BTreeSet::new();
for _ in 0..40 {
let deployment = router
.select(&Op::from("acompletion"), &request_for("chat"))
.await
.unwrap();
picked.insert(deployment.id.clone());
}
assert_eq!(picked.into_iter().collect::<Vec<_>>(), expected);
}
fn request_for(model: &str) -> RouterRequest {
RouterRequest {
model: model.to_owned(),
..RouterRequest::default()
}
}

2
litellm/accounting/application.py
View file

@ -7,6 +7,7 @@ from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_b
from litellm.proxy.db.db_span import db_span
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendTransaction
from litellm.proxy.db.model_usage_rollup import queue_model_usage_upsert
from litellm.proxy.db.rollup_lock_timeout import apply_rollup_lock_timeout
from litellm.repositories.accounting_delivery_repository import DeliveryProjection, LedgerTransaction
if TYPE_CHECKING:
@ -25,6 +26,7 @@ async def apply_deliveries(
from litellm.proxy.guardrails.usage_tracking import UsageUnitIncrement, UsageUnitKey, upsert_usage_unit_row
contributions: Final = project_deliveries(deliveries)
await apply_rollup_lock_timeout(transaction)
await writer.apply_accounting_entities(
prisma_client, proxy_logging_obj, transaction, entity_transactions(contributions)
)

73
litellm/accounting/metrics.py
View file

@ -1,30 +1,55 @@
from typing import Final
from prometheus_client import Gauge, Histogram
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final
from litellm.repositories.accounting_delivery_repository import DeliveryMetrics
_PENDING: Final = Gauge("litellm_accounting_pending", "Pending SpendLogs deliveries", multiprocess_mode="livemax")
_PARKED: Final = Gauge("litellm_accounting_parked", "Parked SpendLogs deliveries", multiprocess_mode="livemax")
_OLDEST: Final = Gauge(
"litellm_accounting_oldest_pending_seconds",
"Age of the oldest pending SpendLogs delivery",
multiprocess_mode="livemax",
)
_BATCH_TIME: Final = Histogram("litellm_accounting_batch_seconds", "SpendLogs delivery worker batch time")
_BATCH_COUNT: Final = Gauge(
"litellm_accounting_batch_deliveries", "Retained batch deliveries", multiprocess_mode="livemax"
)
_TABLE_BYTES: Final = Gauge(
"litellm_accounting_delivery_table_bytes", "Delivery table size in bytes", multiprocess_mode="livemax"
)
if TYPE_CHECKING:
from prometheus_client import Gauge, Histogram
@dataclass(frozen=True, slots=True)
class _Instruments:
pending: "Gauge"
parked: "Gauge"
oldest: "Gauge"
batch_time: "Histogram"
batch_count: "Gauge"
table_bytes: "Gauge"
def _instruments() -> _Instruments | None:
try:
# prometheus-client ships only in the proxy-runtime extra, so a plain litellm[proxy] install lands here
from prometheus_client import Gauge, Histogram
except ImportError:
return None
return _Instruments(
pending=Gauge("litellm_accounting_pending", "Pending SpendLogs deliveries", multiprocess_mode="livemax"),
parked=Gauge("litellm_accounting_parked", "Parked SpendLogs deliveries", multiprocess_mode="livemax"),
oldest=Gauge(
"litellm_accounting_oldest_pending_seconds",
"Age of the oldest pending SpendLogs delivery",
multiprocess_mode="livemax",
),
batch_time=Histogram("litellm_accounting_batch_seconds", "SpendLogs delivery worker batch time"),
batch_count=Gauge(
"litellm_accounting_batch_deliveries", "Retained batch deliveries", multiprocess_mode="livemax"
),
table_bytes=Gauge(
"litellm_accounting_delivery_table_bytes", "Delivery table size in bytes", multiprocess_mode="livemax"
),
)
_INSTRUMENTS: Final = _instruments()
def record_deliveries(metrics: DeliveryMetrics, seconds: float) -> None:
_PENDING.set(metrics.pending_count)
_PARKED.set(metrics.parked_count)
_OLDEST.set(max(0.0, metrics.oldest_pending_seconds))
_BATCH_TIME.observe(seconds)
_BATCH_COUNT.set(metrics.batch_count)
_TABLE_BYTES.set(metrics.table_bytes)
if _INSTRUMENTS is None:
return
_INSTRUMENTS.pending.set(metrics.pending_count)
_INSTRUMENTS.parked.set(metrics.parked_count)
_INSTRUMENTS.oldest.set(max(0.0, metrics.oldest_pending_seconds))
_INSTRUMENTS.batch_time.observe(seconds)
_INSTRUMENTS.batch_count.set(metrics.batch_count)
_INSTRUMENTS.table_bytes.set(metrics.table_bytes)

2
litellm/accounting/snapshot.py
View file

@ -14,6 +14,7 @@ if TYPE_CHECKING:
from litellm.router import Router
_FIELDS: Final = TypeAdapter(dict[str, object])
_JSON_SAFE_INTEGER: Final = 2**53 - 1
def _mapping(value: object) -> Mapping[str, object] | None:
@ -86,6 +87,7 @@ def build_accounting_payload(
"metadata": accounting_metadata,
},
non_finite_as_null=True,
integer_limit=_JSON_SAFE_INTEGER,
)
)

35
litellm/accounting/worker.py
View file

@ -14,6 +14,7 @@ from litellm.accounting.metrics import record_deliveries
from litellm.accounting.projection import Contribution
from litellm.repositories.accounting_delivery_repository import (
AccountingDeliveryRepository,
DeliveryBatchResult,
DeliveryProjection,
LedgerDatabase,
LedgerTransaction,
@ -34,6 +35,14 @@ class _AccountingAlerts(Protocol):
_JOB: Final = "accounting_delivery"
def _filled(result: DeliveryBatchResult, batch_size: int) -> bool:
return result.applied + len(result.parked) >= batch_size
def _progressed(result: DeliveryBatchResult) -> bool:
return bool(result.applied or result.failed)
async def flush_accounting_deliveries(
writer: "DBSpendUpdateWriter", prisma_client: "PrismaClient", proxy_logging_obj: "ProxyLogging"
) -> None:
@ -54,9 +63,9 @@ async def flush_accounting_deliveries(
)
repository: Final = AccountingDeliveryRepository(database)
started: Final = time.monotonic()
batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "10000"))
batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "500"))
async def run_batch(*, retry_only: bool) -> bool:
async def run_batch(*, retry_only: bool) -> DeliveryBatchResult:
result: Final = await repository.apply_pending(apply, batch_size=batch_size, retry_only=retry_only)
if result.applied:
await writer.invalidate_accounting_entities(applied, proxy_logging_obj)
@ -68,17 +77,25 @@ async def flush_accounting_deliveries(
proxy_logging_obj,
message=f"Accounting delivery parked {len(result.parked)} requests. Included spend totals require repair.",
)
return bool(result.applied or result.failed)
return result
await run_batch(retry_only=False)
fresh_deadline: Final = time.monotonic() + 5
for _ in range(batch_size):
if not _filled(await run_batch(retry_only=False), batch_size) or time.monotonic() >= fresh_deadline:
break
retry_deadline: Final = time.monotonic() + 5
for _ in range(batch_size):
if time.monotonic() >= retry_deadline or not await run_batch(retry_only=True):
if time.monotonic() >= retry_deadline or not _progressed(await run_batch(retry_only=True)):
break
retention_days: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1"))
cleanup_deadline: Final = time.monotonic() + 5
for _ in range(batch_size):
if (
await repository.cleanup_applied(retention_days=retention_days, batch_size=batch_size * 2)
< batch_size * 2
or time.monotonic() >= cleanup_deadline
):
break
await repository.cleanup_applied(
retention_days=int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1")),
batch_size=batch_size * 2,
)
metrics: Final = await repository.metrics()
record_deliveries(metrics, time.monotonic() - started)
if metrics.batch_count >= int(

2
litellm/constants.py
View file

@ -28,7 +28,7 @@ GITHUB_COPILOT_PER_USER_AUTH_TYPE: Final = "per_user_oauth"
GITHUB_COPILOT_AUTH_TYPE_KEY: Final = "github_copilot_auth_type"
GITHUB_COPILOT_HOST_KEY: Final = "github_copilot_host"
GITHUB_COPILOT_DEFAULT_HOST: Final = "github.com"
GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS: Final = 60
GITHUB_COPILOT_USER_SESSION_TTL_SECONDS: Final = 1800
GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS: Final = 60
GITHUB_COPILOT_DEVICE_FLOW_CACHE_PREFIX: Final = "github_copilot_device_flow"
USER_PROVIDER_CREDENTIAL_CACHE_PREFIX: Final = "user_provider_credential"

5
litellm/integrations/custom_logger.py
View file

@ -898,6 +898,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
def redacts_messages_itself(self) -> bool:
return False
def exempt_from_message_redaction(self) -> bool:
"""Router pre-call checks that key routing off the request and never log it return True, so
`turn_off_message_logging` leaves the success event they receive unredacted."""
return False
def redact_standard_logging_payload_from_model_call_details(self, model_call_details: dict) -> dict:
"""
Redacts or excludes fields from StandardLoggingPayload before callbacks receive it.

3
litellm/integrations/otel/model/spans.py
View file

@ -341,6 +341,9 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
"LiteLLM_SpendLogDelivery",
"LiteLLM_ClaudeCodePluginTable",
"LiteLLM_ClaudeCodePluginUpload",
"LiteLLM_ClaudeCodePluginVersion",
"LiteLLM_SkillSubmission",
"LiteLLM_ClaudeCodePluginDownload",
"LiteLLM_MemoryTable",
"LiteLLM_AdaptiveRouterState",
"LiteLLM_AdaptiveRouterSession",

4
litellm/interactions/background_cost_polling.py
View file

@ -71,12 +71,12 @@ if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
_TERMINAL_STATUSES: Final = frozenset(
{"completed", "failed", "cancelled", "incomplete", "budget_exceeded", "requires_action"}
{"completed", "failed", "cancelled", "incomplete", "budget_exceeded", "requires_action", "requires_continuation"}
)
_POLLABLE_STATUSES: Final = frozenset({"in_progress", "queued"})
_STATUSES_THAT_PRODUCED_OUTPUT: Final = frozenset({"completed", "requires_action"})
_STATUSES_THAT_PRODUCED_OUTPUT: Final = frozenset({"completed", "requires_action", "requires_continuation"})
SettlementOutcome: TypeAlias = Literal["billed", "released", "unsettled"]

5
litellm/interactions/main.py
View file

@ -95,6 +95,7 @@ async def acreate(
response_mime_type: str | None = None,
# Continuation
previous_interaction_id: str | None = None,
continuation_token: str | None = None,
# Extra params
extra_headers: dict[str, str] | None = None,
extra_body: Mapping[str, object] | None = None,
@ -129,6 +130,7 @@ async def acreate(
response_format: JSON schema for response format
response_mime_type: MIME type of the response
previous_interaction_id: ID of previous interaction for continuation
continuation_token: Token from a `requires_continuation` interaction to resume its long decode
extra_headers: Additional headers
extra_body: Additional body parameters
timeout: Request timeout
@ -164,6 +166,7 @@ async def acreate(
response_format=response_format,
response_mime_type=response_mime_type,
previous_interaction_id=previous_interaction_id,
continuation_token=continuation_token,
extra_headers=extra_headers,
extra_body=extra_body,
timeout=timeout,
@ -224,6 +227,7 @@ def create(
response_mime_type: str | None = None,
# Continuation
previous_interaction_id: str | None = None,
continuation_token: str | None = None,
# Extra params
extra_headers: dict[str, str] | None = None,
extra_body: Mapping[str, object] | None = None,
@ -262,6 +266,7 @@ def create(
response_format: JSON schema for response format
response_mime_type: MIME type of the response
previous_interaction_id: ID of previous interaction for continuation
continuation_token: Token from a `requires_continuation` interaction to resume its long decode
extra_headers: Additional headers
extra_body: Additional body parameters
timeout: Request timeout

1
litellm/interactions/utils.py
View file

@ -20,6 +20,7 @@ INTERACTIONS_API_OPTIONAL_PARAMS: Final = {
"response_format",
"response_mime_type",
"previous_interaction_id",
"continuation_token",
"agent_config",
}

28
litellm/litellm_core_utils/litellm_logging.py
View file

@ -4074,12 +4074,11 @@ class Logging(LiteLLMLoggingBaseClass):
start_time=start_time,
end_time=end_time,
)
if (
callable(callback) is True and is_sync_request and customLogger is not None
): # custom logger functions
if callable(callback) is True and is_sync_request: # custom logger functions
function_logger = customLogger if customLogger is not None else CustomLogger()
print_verbose(f"success callbacks: Running Custom Callback Function - {callback}")
customLogger.log_event(
function_logger.log_event(
kwargs=call_details,
response_obj=result,
start_time=start_time,
@ -4406,7 +4405,11 @@ class Logging(LiteLLMLoggingBaseClass):
if isinstance(callback, CustomLogger): # custom logger class
from litellm.integrations.shadow_eval_logger import ShadowEvalLogger
model_call_details: dict = self.model_call_details
model_call_details: dict = (
self._model_call_details_with_request_messages()
if callback.exempt_from_message_redaction()
else self.model_call_details
)
##################################
# call redaction hook for custom logger
model_call_details = callback.redact_standard_logging_payload_from_model_call_details(
@ -4970,6 +4973,21 @@ class Logging(LiteLLMLoggingBaseClass):
) -> None:
executor.submit(self.success_handler, result, start_time, end_time, cache_hit)
def _model_call_details_with_request_messages(self) -> dict[str, object]: # mutable-ok: loggers take a dict
"""The success-event payload with the request messages put back where `turn_off_message_logging`
redacted them in place, for a callback that routes on the request and never logs it."""
request_messages: Final = cast(object, self.messages) # cast-ok: the request messages are only put back as-is
standard_logging_object: Final = self.model_call_details.get("standard_logging_object")
return {
**self.model_call_details,
"messages": request_messages,
**(
{"standard_logging_object": {**standard_logging_object, "messages": request_messages}}
if isinstance(standard_logging_object, dict)
else {}
),
}
def _get_completed_success_callbacks(
self, *, asynchronous: bool
) -> tuple[str | Callable[..., object] | CustomLogger, ...]:

2
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -5879,7 +5879,7 @@ def _bedrock_tools_pt(
description = name
defs = parameters.pop("$defs", {})
defs_copy = copy.deepcopy(defs) if isinstance(defs, dict) else {}
defs_copy = defs if isinstance(defs, dict) else {}
# Expand $ref references in parameters using the definitions
# Note: We don't pre-flatten defs as that causes exponential memory growth
# with circular references (see issue #19098). unpack_defs handles nested

11
litellm/litellm_core_utils/safe_json_dumps.py
View file

@ -19,6 +19,12 @@ def _finite_or_none(value: float | None) -> float | None:
return None if isinstance(value, float) and not math.isfinite(value) else value
def _bounded(value: float | None, integer_limit: int | None) -> float | None:
if integer_limit is None or isinstance(value, bool) or not isinstance(value, int):
return value
return max(-integer_limit, min(integer_limit, value))
def safe_json_structure(
data: object,
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH,
@ -26,12 +32,14 @@ def safe_json_structure(
key: str | None = None,
*,
non_finite_as_null: bool = False,
integer_limit: int | None = None,
) -> object:
"""
Rebuild data out of JSON-native pieces while detecting circular references.
If a circular reference is detected then a marker string is returned.
NUL bytes are stripped from strings to prevent PostgreSQL 22P05 errors. With
non_finite_as_null, NaN and infinite floats become null, since jsonb rejects them.
With integer_limit, integers are clamped to that magnitude.
value_transform, when given, is applied to every string leaf (and to the
str() fallback for non-serializable objects) with the mapping key the leaf
@ -51,7 +59,8 @@ def safe_json_structure(
cleaned = obj.replace("\x00", "") if "\x00" in obj else obj
return _transform(key, cleaned)
if isinstance(obj, (int, float, bool, type(None))):
return _finite_or_none(obj) if non_finite_as_null else obj
bounded: Final = _bounded(obj, integer_limit)
return _finite_or_none(bounded) if non_finite_as_null else bounded
# Check for circular reference.
if id(obj) in seen:
return "CircularReference Detected"

4
litellm/llms/anthropic/chat/transformation.py
View file

@ -827,11 +827,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
return cls._raise_invalid_reasoning_effort(model, value, llm_provider)
def get_supported_openai_params(self, model: str):
sampling_params: Final = ("temperature", "top_p") if self._supports_sampling_params(model) else ()
params: Final = [
"stream",
"stop",
"temperature",
"top_p",
*sampling_params,
"max_tokens",
"max_completion_tokens",
"tools",

3
litellm/llms/anthropic/messages/adapters/handler.py
View file

@ -123,10 +123,9 @@ def _extract_proxy_litellm_metadata(
def _requested_router(kwargs: Mapping[str, object]) -> "Router | None":
from litellm.router import Router
from litellm.router_python_core import PythonCore
requested: Final = kwargs.get("litellm_router")
return requested if isinstance(requested, (Router, PythonCore)) else None # pyright: ignore[reportReturnType] # the router passes its core, which carries Router's members
return requested if isinstance(requested, Router) else None
def request_safeguards_evaluator(

4
litellm/llms/gemini/interactions/transformation.py
View file

@ -126,6 +126,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
"response_format",
"response_mime_type",
"previous_interaction_id",
"continuation_token",
]
def validate_environment(
@ -206,6 +207,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
"environment",
"response_modalities",
"previous_interaction_id",
"continuation_token",
]
dropped_keys: Final = (
frozenset({"agent_config"})
@ -301,7 +303,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
) -> InteractionsAPIStreamingResponse:
"""Parse streaming chunk."""
verbose_logger.debug("Google AI Interactions streaming chunk: %s", parsed_chunk)
return InteractionsAPIStreamingResponse(**parsed_chunk)
return InteractionsAPIStreamingResponse.model_validate(parsed_chunk)
# GET / DELETE / CANCEL - just build URLs, responses match spec directly

225
litellm/llms/github_copilot/authenticator.py
View file

@ -1,62 +1,30 @@
import hashlib
import json
import os
import time
from dataclasses import dataclass
from datetime import datetime
from typing import Any, Final, TypeVar
from collections.abc import Mapping
from typing import Final, Protocol, cast # noqa: TID251 # pins untyped clients and caches to local Protocols
import httpx
from pydantic import BaseModel, ConfigDict, TypeAdapter, with_config
from pydantic import ConfigDict, TypeAdapter, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import GITHUB_COPILOT_USER_SESSION_TTL_SECONDS
from litellm.litellm_core_utils.asyncify import can_block_current_thread
from litellm.llms.custom_httpx.http_handler import get_httpx_client
from .common_utils import (
GetAccessTokenError,
GetAPIKeyError,
GetDeviceCodeError,
RefreshAPIKeyError,
)
# Constants (default values — overridable via environment variables at call time)
DEFAULT_GITHUB_CLIENT_ID: Final = "Iv1.b507a08c87ecfe98"
DEFAULT_GITHUB_DEVICE_CODE_URL: Final = "https://github.com/login/device/code"
DEFAULT_GITHUB_ACCESS_TOKEN_URL: Final = "https://github.com/login/oauth/access_token"
DEFAULT_GITHUB_API_KEY_URL: Final = "https://api.github.com/copilot_internal/v2/token"
class CopilotEndpoints(BaseModel):
model_config = ConfigDict(frozen=True, strict=True, extra="ignore", hide_input_in_errors=True)
api: str | None = None
class CopilotCachedKey(BaseModel):
model_config = ConfigDict(frozen=True, strict=True, extra="ignore", hide_input_in_errors=True)
token: str
expires_at: float
endpoints: CopilotEndpoints | None = None
class CopilotCachedEndpoints(BaseModel):
model_config = ConfigDict(frozen=True, strict=True, extra="ignore", hide_input_in_errors=True)
endpoints: CopilotEndpoints | None = None
ApiKeyFileView = TypeVar("ApiKeyFileView", CopilotCachedKey, CopilotCachedEndpoints)
@dataclass(frozen=True, slots=True)
class _UnsavedApiKey:
access_token: str
key: CopilotCachedKey
_UNSAVED_API_KEYS: Final[dict[str, _UnsavedApiKey]] = {} # mutable-ok: refreshed keys api-key.json could not hold
DEFAULT_GITHUB_COPILOT_USER_URL: Final = "https://api.github.com/copilot_internal/user"
def _read_token_file(path: str) -> str:
@ -78,11 +46,66 @@ class _AccessTokenPoll(TypedDict):
error: ReadOnly[NotRequired[object]]
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True))
_JSON_OBJECT: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(
Mapping[str, object], config=ConfigDict(strict=True)
)
_DEVICE_CODE: Final = TypeAdapter(_DeviceCode)
_ACCESS_TOKEN_POLL: Final = TypeAdapter(_AccessTokenPoll)
class _CopilotAPIBaseCache(Protocol):
def get_cache(self, key: str) -> object: ...
def set_cache(self, key: str, value: str, ttl: float | None = None) -> object: ...
class _SyncGetClient(Protocol):
def get(self, url: str, *, headers: Mapping[str, str]) -> httpx.Response: ...
_COPILOT_API_BASE_CACHE: Final[_CopilotAPIBaseCache] = cast( # cast-ok: [LIT006] InMemoryCache methods lack type stubs
_CopilotAPIBaseCache,
InMemoryCache(
max_size_in_memory=1000,
default_ttl=GITHUB_COPILOT_USER_SESSION_TTL_SECONDS,
),
)
def _copilot_api_base_for_token(access_token: str) -> str | None:
cache_key: Final = hashlib.sha256(access_token.encode()).hexdigest()
cached_api_base: Final[object] = _COPILOT_API_BASE_CACHE.get_cache(cache_key)
if isinstance(cached_api_base, str):
return cached_api_base
try:
response: Final = cast( # cast-ok: [LIT006] HTTPHandler.get has no return annotation
_SyncGetClient,
get_httpx_client(),
).get(
DEFAULT_GITHUB_COPILOT_USER_URL,
headers=github_api_headers(access_token),
)
if not response.is_success:
return None
payload: Final[Mapping[str, object]] = _JSON_OBJECT.validate_python(response.json())
endpoints: Final[Mapping[str, object]] = _JSON_OBJECT.validate_python(payload.get("endpoints"))
api_base: Final[object] = endpoints.get("api")
if not isinstance(api_base, str):
return None
from .github_host import github_host_endpoints
from .per_user_auth import validated_copilot_api_base
validated: Final = validated_copilot_api_base(api_base, github_host_endpoints("github.com"))
if validated != api_base.rstrip("/"):
return None
_COPILOT_API_BASE_CACHE.set_cache(cache_key, validated)
except Exception:
return None
else:
return validated
def github_api_headers(
access_token: str | None = None,
) -> dict[str, str]: # mutable-ok: returned straight into httpx handlers whose headers params require dict
@ -111,7 +134,6 @@ class Authenticator:
self.token_dir,
os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_FILE", "access-token"),
)
self.api_key_file = os.path.join(self.token_dir, os.getenv("GITHUB_COPILOT_API_KEY_FILE", "api-key.json"))
self._ensure_token_dir()
def get_access_token(self) -> str:
@ -155,7 +177,7 @@ class Authenticator:
except OSError:
verbose_logger.error("Error saving access token to file")
return access_token
except (GetDeviceCodeError, GetAccessTokenError, RefreshAPIKeyError) as e:
except (GetDeviceCodeError, GetAccessTokenError) as e:
verbose_logger.warning("Failed attempt %s: %s", attempt + 1, e)
continue
@ -165,69 +187,15 @@ class Authenticator:
)
def get_api_key(self) -> str:
"""
Get the API key, refreshing if necessary.
Returns:
str: The GitHub Copilot API key.
Raises:
GetAPIKeyError: If unable to obtain an API key.
"""
unsaved: Final = self._unsaved_api_key()
if unsaved is not None:
return unsaved.token
cached: Final = self._read_api_key_file(CopilotCachedKey)
if cached is not None and cached.expires_at > datetime.now().timestamp():
return cached.token
if cached is not None:
verbose_logger.warning("API key expired, refreshing")
try:
api_key_info: Final = self._refresh_api_key()
except RefreshAPIKeyError as e:
raise GetAPIKeyError(
message=f"Failed to refresh API key: {e}",
status_code=401,
) from e
serialized: Final = json.dumps(api_key_info)
try:
with open(self.api_key_file, "w") as f:
f.write(serialized)
except OSError as e:
verbose_logger.warning(
"Keeping the refreshed API key in memory, it could not be saved to %s: %s", self.api_key_file, e
)
self._keep_unsaved_api_key(serialized)
token: Final = api_key_info.get("token")
if not token:
raise GetAPIKeyError(
message="API key response missing token",
status_code=401,
)
return token
"""Return the stored GitHub access token, loading it with device flow when needed."""
return self.get_access_token()
def get_api_base(self) -> str | None:
"""
Get the API endpoint of the current API key.
Returns:
Optional[str]: The GitHub Copilot API endpoint, or None if not found.
"""
cached: Final = self._unsaved_api_key() or self._read_api_key_file(CopilotCachedEndpoints)
if cached is None or cached.endpoints is None:
return None
return cached.endpoints.api
def _read_api_key_file(self, view: type[ApiKeyFileView]) -> ApiKeyFileView | None:
try:
return view.model_validate_json(_read_token_file(self.api_key_file))
except OSError:
verbose_logger.warning("No API key file found or error opening file")
return None
except ValueError as e:
verbose_logger.warning("Ignoring %s, not a Copilot API key file: %s", self.api_key_file, e)
"""Return the Copilot API endpoint for a stored token without starting device flow."""
access_token: Final = self._stored_access_token()
if access_token is None:
return None
return _copilot_api_base_for_token(access_token)
def _stored_access_token(self) -> str | None:
try:
@ -236,59 +204,6 @@ class Authenticator:
return None
return access_token or None
def _unsaved_api_key(self) -> CopilotCachedKey | None:
unsaved: Final = _UNSAVED_API_KEYS.get(self.api_key_file)
if unsaved is None or unsaved.key.expires_at <= time.time():
return None
return unsaved.key if unsaved.access_token == self._stored_access_token() else None
def _keep_unsaved_api_key(self, serialized: str) -> None:
access_token: Final = self._stored_access_token()
if access_token is None:
return
try:
key: Final = CopilotCachedKey.model_validate_json(serialized)
except ValueError:
return
_UNSAVED_API_KEYS[self.api_key_file] = _UnsavedApiKey(access_token=access_token, key=key)
def _refresh_api_key(self) -> dict[str, Any]:
"""
Refresh the API key using the access token.
Returns:
Dict[str, Any]: The API key information including token and expiration.
Raises:
RefreshAPIKeyError: If unable to refresh the API key.
"""
access_token: Final = self.get_access_token()
headers: Final = self._get_github_headers(access_token)
api_key_url: Final = os.getenv("GITHUB_COPILOT_API_KEY_URL", DEFAULT_GITHUB_API_KEY_URL)
max_retries: Final = 3
for attempt in range(max_retries):
try:
sync_client = get_httpx_client()
response = sync_client.get(api_key_url, headers=headers)
response.raise_for_status()
response_json = response.json()
if "token" in response_json:
return response_json
else:
verbose_logger.warning("API key response missing token: %s", response_json)
except httpx.HTTPStatusError as e:
verbose_logger.error("HTTP error refreshing API key (attempt %s/%s): %s", attempt + 1, max_retries, e)
except Exception as e:
verbose_logger.error("Unexpected error refreshing API key: %s", e)
raise RefreshAPIKeyError(
message="Failed to refresh API key after maximum retries",
status_code=401,
)
def _ensure_token_dir(self) -> None:
"""Ensure the token directory exists."""
if not os.path.exists(self.token_dir):

6
litellm/llms/github_copilot/chat/transformation.py
View file

@ -16,7 +16,7 @@ from litellm.types.utils import ModelResponse
from ..authenticator import Authenticator
from ..common_utils import (
DEFAULT_GITHUB_COPILOT_API_BASE,
GetAPIKeyError,
GetAccessTokenError,
get_copilot_default_headers,
)
from ..per_user_auth import require_github_copilot_user_session
@ -56,7 +56,7 @@ class GithubCopilotConfig(OpenAIConfig):
dynamic_api_base: Final = self.api_base_without_login(api_base)
try:
dynamic_api_key: Final = self.authenticator.get_api_key()
except GetAPIKeyError as e:
except GetAccessTokenError as e:
raise AuthenticationError(
model=model,
llm_provider=custom_llm_provider,
@ -104,7 +104,7 @@ class GithubCopilotConfig(OpenAIConfig):
return get_copilot_default_headers(session_token)
try:
return get_copilot_default_headers(self.authenticator.get_api_key())
except GetAPIKeyError:
except GetAccessTokenError:
return {}
def validate_environment(

12
litellm/llms/github_copilot/common_utils.py
View file

@ -46,18 +46,6 @@ class GetAccessTokenError(GithubCopilotError):
pass
class APIKeyExpiredError(GithubCopilotError):
pass
class RefreshAPIKeyError(GithubCopilotError):
pass
class GetAPIKeyError(GithubCopilotError):
pass
def get_copilot_default_headers(
api_key: str,
) -> dict[str, str]: # mutable-ok: callers merge caller headers into the returned dict

4
litellm/llms/github_copilot/embedding/transformation.py
View file

@ -30,7 +30,7 @@ from litellm.utils import convert_to_model_response_object
from ..authenticator import Authenticator
from ..common_utils import (
DEFAULT_GITHUB_COPILOT_API_BASE,
GetAPIKeyError,
GetAccessTokenError,
get_copilot_default_headers,
pin_session_authorization,
)
@ -100,7 +100,7 @@ class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig):
return merged_headers
except GetAPIKeyError as e:
except GetAccessTokenError as e:
raise AuthenticationError(
model=model,
llm_provider="github_copilot",

11
litellm/llms/github_copilot/github_host.py
View file

@ -7,8 +7,8 @@ from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST
from .authenticator import (
DEFAULT_GITHUB_ACCESS_TOKEN_URL,
DEFAULT_GITHUB_API_KEY_URL,
DEFAULT_GITHUB_CLIENT_ID,
DEFAULT_GITHUB_COPILOT_USER_URL,
DEFAULT_GITHUB_DEVICE_CODE_URL,
)
@ -20,11 +20,9 @@ class GithubHostEndpoints:
host: str
device_code_url: str
access_token_url: str
api_key_url: str
copilot_user_url: str
user_url: str
client_id: str
exchanges_copilot_token: bool
copilot_api_hosts: frozenset[str]
copilot_api_host_suffix: str | None
default_copilot_api_base: str
@ -51,17 +49,14 @@ def github_host_endpoints(host: str) -> GithubHostEndpoints:
if normalized is None:
raise ValueError("host must be github.com or a single-subdomain GHE.com hostname")
client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
exchanges_copilot_token: Final = client_id == DEFAULT_GITHUB_CLIENT_ID
if normalized == GITHUB_COPILOT_DEFAULT_HOST:
return GithubHostEndpoints(
host=normalized,
device_code_url=os.getenv("GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL),
access_token_url=os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL),
api_key_url=os.getenv("GITHUB_COPILOT_API_KEY_URL", DEFAULT_GITHUB_API_KEY_URL),
copilot_user_url="https://api.github.com/copilot_internal/user",
copilot_user_url=DEFAULT_GITHUB_COPILOT_USER_URL,
user_url="https://api.github.com/user",
client_id=client_id,
exchanges_copilot_token=exchanges_copilot_token,
copilot_api_hosts=frozenset({"githubcopilot.com"}),
copilot_api_host_suffix=".githubcopilot.com",
default_copilot_api_base="https://api.githubcopilot.com",
@ -74,11 +69,9 @@ def github_host_endpoints(host: str) -> GithubHostEndpoints:
host=normalized,
device_code_url=f"https://{normalized}/login/device/code",
access_token_url=f"https://{normalized}/login/oauth/access_token",
api_key_url=f"https://{api_host}/copilot_internal/v2/token",
copilot_user_url=f"https://{api_host}/copilot_internal/user",
user_url=f"https://{api_host}/user",
client_id=client_id,
exchanges_copilot_token=exchanges_copilot_token,
copilot_api_hosts=frozenset({copilot_api_host}),
copilot_api_host_suffix=None,
default_copilot_api_base=f"https://{copilot_api_host}",

4
litellm/llms/github_copilot/messages/transformation.py
View file

@ -10,7 +10,7 @@ from litellm.llms.anthropic.messages.native.transformation import (
from ..authenticator import Authenticator
from ..common_utils import (
DEFAULT_GITHUB_COPILOT_API_BASE,
GetAPIKeyError,
GetAccessTokenError,
get_copilot_default_headers,
pin_session_authorization,
)
@ -84,7 +84,7 @@ class GithubCopilotAnthropicMessagesConfig(AnthropicMessagesConfig):
dynamic_api_key: Final = (
user_session.token if user_session is not None else self.authenticator.get_api_key()
)
except GetAPIKeyError as e:
except GetAccessTokenError as e:
raise AuthenticationError(
model=model,
llm_provider="github_copilot",

146
litellm/llms/github_copilot/per_user_auth.py
View file

@ -1,7 +1,6 @@
"""Per-user GitHub Copilot OAuth: exchange a caller's stored GitHub token for a
short-lived Copilot token, plus the device-flow helpers the proxy endpoints use
to establish that connection. The shared device login (Authenticator) is never
touched on the per-user path.
"""Per-user GitHub Copilot OAuth sessions and the device-flow helpers the proxy
endpoints use to establish those connections. The shared device login
(Authenticator) is never touched on the per-user path.
"""
import asyncio
@ -28,7 +27,7 @@ from litellm.constants import (
GITHUB_COPILOT_AUTH_TYPE_KEY,
GITHUB_COPILOT_HOST_KEY,
GITHUB_COPILOT_PER_USER_AUTH_TYPE,
GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS,
GITHUB_COPILOT_USER_SESSION_TTL_SECONDS,
)
from litellm.exceptions import (
APIConnectionError,
@ -39,6 +38,7 @@ from litellm.exceptions import (
)
from .authenticator import github_api_headers
from .common_utils import get_copilot_default_headers
from .github_host import GithubHostEndpoints, github_host_endpoints, normalized_github_host
GITHUB_COPILOT_USER_SESSION_KWARG_KEY: Final = "github_copilot_user_session"
@ -46,7 +46,6 @@ _LLM_PROVIDER: Final = "github_copilot"
_DEVICE_FLOW_SCOPE: Final = "read:user"
_SESSION_CACHE_MAX_SIZE: Final = 1000
_SYNC_LOCK_COUNT: Final = 64
_ACCOUNT_LOOKUP_TTL_SECONDS: Final = 1800
class _SessionCache(Protocol):
@ -88,19 +87,19 @@ _SESSION_CACHE: Final = cast( # cast-ok: pins the untyped module cache to _Sess
_SessionCache,
InMemoryCache(
max_size_in_memory=_SESSION_CACHE_MAX_SIZE,
default_ttl=GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS,
default_ttl=GITHUB_COPILOT_USER_SESSION_TTL_SECONDS,
),
)
_EXCHANGE_LOCKS: Final = tuple(threading.Lock() for _ in range(_SYNC_LOCK_COUNT))
_RESOLVE_LOCKS: Final = tuple(threading.Lock() for _ in range(_SYNC_LOCK_COUNT))
_IN_FLIGHT: Final[ # mutable-ok: single-flight dedup registry mutated under its own locking
dict[str, tuple[str, "asyncio.Future[GithubCopilotUserSession]"]]
] = {}
@with_config(ConfigDict(extra="allow", strict=True))
class _CopilotTokenPayload(TypedDict):
token: ReadOnly[str]
expires_at: ReadOnly[NotRequired[object]]
class _CopilotUserPayload(TypedDict):
access_type_sku: ReadOnly[NotRequired[object]]
chat_enabled: ReadOnly[NotRequired[object]]
endpoints: ReadOnly[NotRequired[object]]
@ -126,20 +125,12 @@ class _CopilotEndpointsPayload(TypedDict):
api: ReadOnly[NotRequired[object]]
@with_config(ConfigDict(extra="allow", strict=True))
class _CopilotAccountPayload(TypedDict):
access_type_sku: ReadOnly[NotRequired[object]]
chat_enabled: ReadOnly[NotRequired[object]]
endpoints: ReadOnly[NotRequired[object]]
@with_config(ConfigDict(extra="allow", strict=True))
class _GithubUserPayload(TypedDict):
login: ReadOnly[str]
_COPILOT_TOKEN: Final = TypeAdapter(_CopilotTokenPayload)
_COPILOT_ACCOUNT: Final = TypeAdapter(_CopilotAccountPayload)
_COPILOT_USER: Final = TypeAdapter(_CopilotUserPayload)
_DEVICE_CODE: Final = TypeAdapter(_DeviceCodePayload)
_ACCESS_TOKEN_POLL: Final = TypeAdapter(_AccessTokenPollPayload)
_GITHUB_USER: Final = TypeAdapter(_GithubUserPayload)
@ -237,13 +228,13 @@ def _cached_session(cache_key: str, github_token: str) -> GithubCopilotUserSessi
return None
if not hmac.compare_digest(cached_github_token, github_token):
return None
if expires_at - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS <= time.time():
if expires_at <= time.time():
return None
return GithubCopilotUserSession(token=token, api_base=api_base)
def _cache_session(cache_key: str, github_token: str, session: GithubCopilotUserSession, expires_at: float) -> None:
ttl: Final = expires_at - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS - time.time()
ttl: Final = expires_at - time.time()
if ttl <= 0:
return
_SESSION_CACHE.set_cache(cache_key, (github_token, session.token, session.api_base, expires_at), ttl=ttl)
@ -272,57 +263,32 @@ def _session_from_response(
raise _reject_connection(credential_name)
if response.status_code == 429:
raise CallerCredentialRateLimitError(
message="GitHub rate limited the GitHub Copilot token exchange",
message="GitHub rate limited the GitHub Copilot user lookup",
llm_provider=_LLM_PROVIDER,
model="",
)
if not response.is_success:
raise ServiceUnavailableError(
message=f"GitHub Copilot token exchange failed with status {response.status_code}",
message=f"GitHub Copilot user lookup failed with status {response.status_code}",
llm_provider=_LLM_PROVIDER,
model="",
)
if not endpoints.exchanges_copilot_token:
return _session_from_account(response, github_token, credential_name, cache_key, endpoints)
try:
payload: Final = _COPILOT_TOKEN.validate_python(response.json())
payload: Final = _COPILOT_USER.validate_python(response.json())
except Exception as e:
raise ServiceUnavailableError(
message="GitHub Copilot token exchange returned an invalid response",
message="GitHub Copilot user lookup returned an invalid response",
llm_provider=_LLM_PROVIDER,
model="",
) from e
api_base: Final = validated_copilot_api_base(_endpoints_api(payload.get("endpoints")), endpoints)
session: Final = GithubCopilotUserSession(token=payload["token"], api_base=api_base)
expires_at: Final = _expires_at_seconds(payload.get("expires_at"))
if expires_at > 0:
_cache_session(cache_key, github_token, session, expires_at)
return session
def _session_from_account(
response: httpx.Response,
github_token: str,
credential_name: str,
cache_key: str,
endpoints: GithubHostEndpoints,
) -> GithubCopilotUserSession:
try:
payload: Final = _COPILOT_ACCOUNT.validate_python(response.json())
except Exception as e:
raise ServiceUnavailableError(
message="GitHub Copilot account lookup returned an invalid response",
llm_provider=_LLM_PROVIDER,
model="",
) from e
if payload.get("chat_enabled") is False or payload.get("access_type_sku") == "no_access":
_SESSION_CACHE.delete_cache(cache_key)
raise _reject_connection(credential_name)
api_base: Final = validated_copilot_api_base(_endpoints_api(payload.get("endpoints")), endpoints)
session: Final = GithubCopilotUserSession(token=github_token, api_base=api_base)
_cache_session(cache_key, github_token, session, time.time() + _ACCOUNT_LOOKUP_TTL_SECONDS)
_cache_session(cache_key, github_token, session, time.time() + GITHUB_COPILOT_USER_SESSION_TTL_SECONDS)
return session
@ -332,10 +298,6 @@ def _endpoints_api(payload_endpoints: object) -> object:
return _COPILOT_ENDPOINTS.validate_python(payload_endpoints).get("api")
def _session_lookup_url(endpoints: GithubHostEndpoints) -> str:
return endpoints.api_key_url if endpoints.exchanges_copilot_token else endpoints.copilot_user_url
def _to_int(raw: object) -> int:
if isinstance(raw, bool):
return 0
@ -351,20 +313,7 @@ def _to_int(raw: object) -> int:
return 0
def _expires_at_seconds(raw: object) -> float:
if isinstance(raw, bool):
return 0.0
if isinstance(raw, (int, float)):
return float(raw)
if isinstance(raw, str):
try:
return float(raw)
except ValueError:
return 0.0
return 0.0
def _do_exchange(
def _do_resolve(
github_token: str,
credential_name: str,
cache_key: str,
@ -377,12 +326,12 @@ def _do_exchange(
_SyncGetClient, litellm.module_level_client
)
response: Final = client.get(
_session_lookup_url(endpoints),
endpoints.copilot_user_url,
headers=github_api_headers(github_token),
)
except Exception as e:
raise APIConnectionError(
message="GitHub Copilot token exchange request failed",
message="GitHub Copilot user lookup request failed",
llm_provider=_LLM_PROVIDER,
model="",
) from e
@ -398,7 +347,7 @@ def _new_async_github_client() -> _AsyncGitHubClient:
)(_LLM_PROVIDER)
async def _do_aexchange(
async def _do_aresolve(
github_token: str,
credential_name: str,
cache_key: str,
@ -407,19 +356,19 @@ async def _do_aexchange(
client: Final = _new_async_github_client()
try:
response: Final = await client.get(
_session_lookup_url(endpoints),
endpoints.copilot_user_url,
headers=github_api_headers(github_token),
)
except Exception as e:
raise APIConnectionError(
message="GitHub Copilot token exchange request failed",
message="GitHub Copilot user lookup request failed",
llm_provider=_LLM_PROVIDER,
model="",
) from e
return _session_from_response(response, github_token, credential_name, cache_key, endpoints)
def exchange_github_token(
def resolve_copilot_user_session(
user_id: str,
github_token: str,
credential_name: str,
@ -429,15 +378,15 @@ def exchange_github_token(
cached: Final = _cached_session(cache_key, github_token)
if cached is not None:
return cached
lock: Final = _EXCHANGE_LOCKS[hash(cache_key) % _SYNC_LOCK_COUNT]
lock: Final = _RESOLVE_LOCKS[hash(cache_key) % _SYNC_LOCK_COUNT]
with lock:
rechecked: Final = _cached_session(cache_key, github_token)
if rechecked is not None:
return rechecked
return _do_exchange(github_token, credential_name, cache_key, endpoints)
return _do_resolve(github_token, credential_name, cache_key, endpoints)
async def aexchange_github_token(
async def aresolve_copilot_user_session(
user_id: str,
github_token: str,
credential_name: str,
@ -452,14 +401,14 @@ async def aexchange_github_token(
if in_flight_entry is not None:
in_flight_token, in_flight = in_flight_entry
if not hmac.compare_digest(in_flight_token, github_token):
return await _do_aexchange(github_token, credential_name, cache_key, endpoints)
return await _do_aresolve(github_token, credential_name, cache_key, endpoints)
if not in_flight.done() and in_flight.get_loop() is loop:
return await asyncio.shield(in_flight)
future: Final = loop.create_future()
entry: Final = (github_token, future)
_IN_FLIGHT[cache_key] = entry
try:
session: Final = await _do_aexchange(github_token, credential_name, cache_key, endpoints)
session: Final = await _do_aresolve(github_token, credential_name, cache_key, endpoints)
except BaseException as e:
if not future.done():
future.set_exception(e)
@ -610,12 +559,41 @@ async def acheck_copilot_seat(
credential_name: str,
endpoints: GithubHostEndpoints,
) -> GithubCopilotUserSession:
return await aexchange_github_token(
session: Final = await aresolve_copilot_user_session(
user_id=user_id,
github_token=github_token,
credential_name=credential_name,
endpoints=endpoints,
)
client: Final = _new_async_github_client()
models_url: Final = f"{session.api_base}/models"
try:
response: Final = await client.get(
models_url,
headers=get_copilot_default_headers(session.token),
)
except Exception as e:
raise APIConnectionError(
message="GitHub Copilot models check request failed",
llm_provider=_LLM_PROVIDER,
model="",
) from e
if response.status_code in (401, 403, 404):
_SESSION_CACHE.delete_cache(_session_cache_key(user_id, credential_name, endpoints.host))
raise _reject_connection(credential_name)
if response.status_code == 429:
raise CallerCredentialRateLimitError(
message="GitHub rate limited the GitHub Copilot models check",
llm_provider=_LLM_PROVIDER,
model="",
)
if not response.is_success:
raise ServiceUnavailableError(
message=f"GitHub Copilot models check failed with status {response.status_code}",
llm_provider=_LLM_PROVIDER,
model="",
)
return session
def github_copilot_auth_mode(litellm_credential_name: object, auth_type_value: object) -> bool:
@ -804,7 +782,7 @@ def attach_github_copilot_user_session(
credential_name: Final = credential_name_obj if isinstance(credential_name_obj, str) else ""
endpoints: Final = github_copilot_host_endpoints(credential_name)
user_id, github_token = _caller_connection(kwargs, credential_name, endpoints)
session: Final = exchange_github_token(
session: Final = resolve_copilot_user_session(
user_id=user_id,
github_token=github_token,
credential_name=credential_name,
@ -826,7 +804,7 @@ async def aattach_github_copilot_user_session(
credential_name: Final = credential_name_obj if isinstance(credential_name_obj, str) else ""
endpoints: Final = github_copilot_host_endpoints(credential_name)
user_id, github_token = _caller_connection(kwargs, credential_name, endpoints)
session: Final = await aexchange_github_token(
session: Final = await aresolve_copilot_user_session(
user_id=user_id,
github_token=github_token,
credential_name=credential_name,

4
litellm/llms/github_copilot/responses/transformation.py
View file

@ -38,7 +38,7 @@ from litellm.utils import cached_get_model_info_helper
from ..authenticator import Authenticator
from ..common_utils import (
DEFAULT_GITHUB_COPILOT_API_BASE,
GetAPIKeyError,
GetAccessTokenError,
get_copilot_default_headers,
pin_session_authorization,
)
@ -256,7 +256,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig):
return merged_headers
except GetAPIKeyError as e:
except GetAccessTokenError as e:
raise AuthenticationError(
model=model,
llm_provider="github_copilot",

7
litellm/llms/openai/openai.py
View file

@ -1,3 +1,4 @@
import inspect
import time
import types
from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping, Sequence
@ -379,6 +380,11 @@ _EXTRA_HEADERS_ADAPTER: Final = TypeAdapter(dict[str, str] | None)
_EXTRA_QUERY_ADAPTER: Final = TypeAdapter(dict[str, object] | None)
_NO_EXTRA_HEADERS: Final[Mapping[str, str]] = types.MappingProxyType({})
_SDK_OPTION_KEYS: Final = frozenset(("extra_headers", "extra_query", "extra_body"))
_CLIENT_API_KEY_ONLY_AUTH: Final[Mapping[str, Mapping[str, bool]]] = (
types.MappingProxyType({"security": types.MappingProxyType({"bearer_auth": True})})
if "security" in inspect.signature(make_request_options).parameters
else types.MappingProxyType({})
)
def _raw_chat_response(
@ -398,6 +404,7 @@ def _request_without_sdk_transform(
extra_query=_EXTRA_QUERY_ADAPTER.validate_python(data.get("extra_query")),
extra_body=data.get("extra_body"),
timeout=timeout,
**_CLIENT_API_KEY_ONLY_AUTH, # pyright: ignore[reportArgumentType] # openai<3 make_request_options has no security kwarg
)
return body, options

392
litellm/model_prices_and_context_window_backup.json
View file

@ -4046,6 +4046,8 @@
"cache_creation_input_token_cost": 1.25e-05,
"cache_creation_input_token_cost_above_1hr": 2e-05,
"cache_read_input_token_cost": 2.5e-07,
"deprecation_date": "2027-09-01",
"deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement",
"litellm_provider": "azure_ai",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
@ -4068,7 +4070,7 @@
"supports_xhigh_reasoning_effort": true,
"supports_max_reasoning_effort": false,
"prompt_cache_min_tokens": 512,
"source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/claude-models"
"source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement"
},
"azure_ai/claude-opus-5": {
"deprecation_date": "2027-07-08",
@ -13458,13 +13460,16 @@
"supports_reasoning": true
},
"azure_ai/mistral-document-ai-2512": {
"deprecation_date": "2027-02-15",
"deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement",
"litellm_provider": "azure_ai",
"ocr_cost_per_page": 0.003,
"mode": "ocr",
"supported_endpoints": [
"/v1/ocr"
],
"source": "https://ai.azure.com/catalog/models/mistral-document-ai-2512"
"source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement",
"successor_model": "azure_ai/mistral-ocr-4-0"
},
"azure_ai/mistral-ocr-4-0": {
"litellm_provider": "azure_ai",
@ -32320,7 +32325,9 @@
"deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations",
"successor_model": "gemini/gemini-3.8-flash",
"input_cost_per_token_above_200k_tokens_batches": 1.25e-06,
"output_cost_per_token_above_200k_tokens_batches": 7.5e-06
"input_cost_per_token_above_200k_tokens_flex": 1.25e-06,
"output_cost_per_token_above_200k_tokens_batches": 7.5e-06,
"output_cost_per_token_above_200k_tokens_flex": 7.5e-06
},
"gemini/gemini-2.5-computer-use-preview-10-2025": {
"deprecation_date": "2026-07-28",
@ -32811,7 +32818,7 @@
"output_cost_per_token": 9e-06,
"output_cost_per_video_token": 1.75e-05,
"rpm": 2000,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"source": "https://ai.google.dev/gemini-api/docs/deprecations",
"supported_endpoints": [
"/v1beta/interactions"
],
@ -32833,7 +32840,7 @@
"tpm": 800000,
"deprecation_date": "2026-10-22",
"deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations",
"successor_model": "gemini/gemini-omni-1.1-flash"
"successor_model": "gemini/gemini-omni-1.2-flash"
},
"gemini/gemini-3.1-pro-preview": {
"prompt_cache_min_tokens": 4096,
@ -32899,7 +32906,9 @@
"output_cost_per_token_flex": 6e-06,
"cache_read_input_token_cost_above_200k_tokens_batches": 4e-07,
"input_cost_per_token_above_200k_tokens_batches": 2e-06,
"output_cost_per_token_above_200k_tokens_batches": 9e-06
"input_cost_per_token_above_200k_tokens_flex": 2e-06,
"output_cost_per_token_above_200k_tokens_batches": 9e-06,
"output_cost_per_token_above_200k_tokens_flex": 9e-06
},
"gemini/gemini-3.1-pro-preview-customtools": {
"prompt_cache_min_tokens": 4096,
@ -32965,7 +32974,9 @@
"output_cost_per_token_flex": 6e-06,
"cache_read_input_token_cost_above_200k_tokens_batches": 4e-07,
"input_cost_per_token_above_200k_tokens_batches": 2e-06,
"output_cost_per_token_above_200k_tokens_batches": 9e-06
"input_cost_per_token_above_200k_tokens_flex": 2e-06,
"output_cost_per_token_above_200k_tokens_batches": 9e-06,
"output_cost_per_token_above_200k_tokens_flex": 9e-06
},
"gemini-3-flash-preview": {
"cache_read_input_audio_token_cost": 1e-07,
@ -33613,14 +33624,14 @@
"gemini/veo-3.1-fast-generate-preview": {
"deprecation_date": "2026-10-22",
"deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations",
"successor_model": "gemini/gemini-omni-1.1-flash",
"successor_model": "gemini/gemini-omni-1.2-flash",
"litellm_provider": "gemini",
"max_input_tokens": 1024,
"mode": "video_generation",
"output_cost_per_second": 0.1,
"output_cost_per_second_1080p": 0.12,
"output_cost_per_second_4k": 0.3,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"source": "https://ai.google.dev/gemini-api/docs/deprecations",
"supported_modalities": [
"text"
],
@ -33631,13 +33642,13 @@
"gemini/veo-3.1-generate-preview": {
"deprecation_date": "2026-10-22",
"deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations",
"successor_model": "gemini/gemini-omni-1.1-flash",
"successor_model": "gemini/gemini-omni-1.2-flash",
"litellm_provider": "gemini",
"max_input_tokens": 1024,
"mode": "video_generation",
"output_cost_per_second": 0.4,
"output_cost_per_second_4k": 0.6,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"source": "https://ai.google.dev/gemini-api/docs/deprecations",
"supported_modalities": [
"text"
],
@ -33648,13 +33659,13 @@
"gemini/veo-3.1-lite-generate-preview": {
"deprecation_date": "2026-10-22",
"deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations",
"successor_model": "gemini/gemini-omni-1.1-flash",
"successor_model": "gemini/gemini-omni-1.2-flash",
"litellm_provider": "gemini",
"max_input_tokens": 1024,
"mode": "video_generation",
"output_cost_per_second": 0.05,
"output_cost_per_second_1080p": 0.08,
"source": "https://ai.google.dev/gemini-api/docs/video",
"source": "https://ai.google.dev/gemini-api/docs/deprecations",
"supported_modalities": [
"text"
],
@ -40887,6 +40898,22 @@
"output_cost_per_token": 4e-07,
"supports_tool_choice": true
},
"jina-reranker-v3.5": {
"input_cost_per_token": 5e-08,
"litellm_provider": "jina_ai",
"max_input_tokens": 131072,
"mode": "rerank",
"output_cost_per_token": 0.0,
"source": "https://api.jina.ai/v1/models"
},
"jina-reranker-m0": {
"input_cost_per_token": 5e-08,
"litellm_provider": "jina_ai",
"max_input_tokens": 10240,
"mode": "rerank",
"output_cost_per_token": 0.0,
"source": "https://api.jina.ai/v1/models"
},
"jina-reranker-v2-base-multilingual": {
"input_cost_per_token": 5e-08,
"litellm_provider": "jina_ai",
@ -43765,7 +43792,8 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_vision": true
},
"mistral/zai-glm-5-3": {
"cache_read_input_token_cost": 1.4e-07,
@ -45264,6 +45292,18 @@
"source": "https://tokenfactory.nebius.com/api/public/models_info",
"cache_read_input_token_cost": 1.5e-07
},
"nebius/moonshotai/Kimi-K3.1-Pro": {
"input_cost_per_token": 2e-06,
"output_cost_per_token": 1e-05,
"cache_read_input_token_cost": 2e-07,
"litellm_provider": "nebius",
"mode": "chat",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_vision": true,
"supports_prompt_caching": true,
"source": "https://tokenfactory.nebius.com/api/public/models_info"
},
"nebius/moonshotai/Kimi-K3": {
"max_input_tokens": 1048576,
"max_output_tokens": 1024000,
@ -60260,6 +60300,7 @@
"rpm": 10,
"gemini_audio_only_live": true,
"input_cost_per_second": 8.33333333333e-05,
"search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014},
"supports_response_schema": false,
"supports_reasoning": true
},
@ -67653,17 +67694,20 @@
"source": "https://deepinfra.com/pricing"
},
"gemini/gemini-omni-1.1-flash": {
"deprecation_date": "2026-12-03",
"deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations",
"input_cost_per_audio_token": 1.5e-06,
"input_cost_per_token": 1.5e-06,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_input_tokens": 131072,
"max_output_tokens": 65536,
"mode": "chat",
"output_cost_per_reasoning_token": 9e-06,
"output_cost_per_token": 9e-06,
"output_cost_per_video_token": 1.75e-05,
"rpm": 2000,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"source": "https://ai.google.dev/gemini-api/docs/models/gemini-omni-flash",
"successor_model": "gemini/gemini-omni-1.2-flash",
"supported_endpoints": [
"/v1beta/interactions"
],
@ -71179,6 +71223,8 @@
},
"gpt-5.5-cyber-preview": {
"cache_read_input_token_cost": 1.25e-06,
"deprecation_date": "2027-02-01",
"deprecation_source": "https://developers.openai.com/api/docs/deprecations",
"input_cost_per_token": 1.25e-05,
"litellm_provider": "openai",
"mode": "chat",
@ -72138,8 +72184,10 @@
"output_cost_per_token": 4.5e-06,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supports_audio_input": true,
"supports_audio_output": true,
"tpm": 250000,
"rpm": 10,
"search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014},
"supports_function_calling": true,
"supports_response_schema": false,
"supports_vision": true,
@ -72159,8 +72207,10 @@
"output_cost_per_token": 4.5e-06,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supports_audio_input": true,
"supports_audio_output": true,
"tpm": 250000,
"rpm": 10,
"search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014},
"supports_function_calling": true,
"supports_response_schema": false,
"supports_vision": true,
@ -76457,6 +76507,33 @@
"supports_tool_choice": true,
"supports_vision": true
},
"us-gov.xai.grok-4.7": {
"supports_regex_lookaround": false,
"supported_endpoints": ["/v1/chat/completions"],
"supports_bedrock_runtime_chat_completions_tools_with_reasoning": true,
"supports_bedrock_runtime_chat_completions_response_format": true,
"cache_read_input_token_cost": 6.6e-07,
"cache_read_input_token_cost_flex": 3.3e-07,
"cache_read_input_token_cost_priority": 1.155e-06,
"input_cost_per_token": 2.64e-06,
"input_cost_per_token_batches": 1.32e-06,
"input_cost_per_token_flex": 1.32e-06,
"input_cost_per_token_priority": 4.62e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 524288,
"max_output_tokens": 524288,
"mode": "chat",
"output_cost_per_token": 7.92e-06,
"output_cost_per_token_batches": 3.96e-06,
"output_cost_per_token_flex": 3.96e-06,
"output_cost_per_token_priority": 1.386e-05,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json",
"supports_function_calling": true,
"supports_prompt_caching": false,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"baseten/deepseek-ai/DeepSeek-V4.1-Flash-Fast": {
"cache_read_input_token_cost": 1.4e-08,
"input_cost_per_token": 6e-07,
@ -78425,7 +78502,7 @@
"cache_creation_input_token_cost": 1e-05,
"cache_creation_input_token_cost_above_1hr": 1.6e-05,
"cache_read_input_token_cost": 4e-07,
"deprecation_date": "2027-10-01",
"deprecation_date": "2027-01-15",
"deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement",
"input_cost_per_token": 8e-06,
"litellm_provider": "azure_ai",
@ -78435,7 +78512,7 @@
"output_cost_per_token": 4e-05,
"prompt_cache_min_tokens": 512.0,
"search_context_cost_per_query": {"search_context_size_high":0.02,"search_context_size_low":0.02,"search_context_size_medium":0.02},
"source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/claude-models",
"source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement",
"supports_adaptive_thinking": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -82770,6 +82847,39 @@
"supports_vision": true,
"supports_xhigh_reasoning_effort": true
},
"gpt-6-cyber": {
"cache_creation_input_token_cost": 2.5e-05,
"cache_creation_input_token_cost_above_12hr": 5e-05,
"cache_read_input_token_cost": 2e-06,
"input_cost_per_token": 2e-05,
"litellm_provider": "openai",
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 0.0001,
"source": "https://developers.openai.com/api/docs/models/gpt-6-cyber",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_computer_use": true,
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_search": true,
"supports_vision": true,
"supports_web_search": true
},
"gpt-6-astra-law": {
"cache_read_input_token_cost": 1e-06,
"input_cost_per_token": 1e-05,
@ -83130,6 +83240,11 @@
"supports_vision": true,
"supports_reasoning": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.033,
"search_context_size_medium": 0.033,
"search_context_size_high": 0.033
},
"supports_prompt_caching": true,
"supports_function_calling": true,
"supports_response_schema": true,
@ -83146,6 +83261,11 @@
"supports_system_messages": true,
"supports_reasoning": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.033,
"search_context_size_medium": 0.033,
"search_context_size_high": 0.033
},
"supports_prompt_caching": true,
"supports_function_calling": true,
"supports_response_schema": true,
@ -83155,6 +83275,7 @@
"litellm_provider": "empiriolabs",
"mode": "chat",
"max_input_tokens": 202745,
"max_output_tokens": 128000,
"input_cost_per_token": 8.25e-07,
"output_cost_per_token": 3.301e-06,
"input_cost_per_token_above_32k_tokens": 1.1e-06,
@ -83178,6 +83299,11 @@
"supports_system_messages": true,
"supports_reasoning": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.033,
"search_context_size_medium": 0.033,
"search_context_size_high": 0.033
},
"supports_function_calling": true,
"supports_response_schema": true,
"source": "https://empiriolabs.ai/models/glm-5-2"
@ -83230,6 +83356,11 @@
"supports_system_messages": true,
"supports_reasoning": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.033,
"search_context_size_medium": 0.033,
"search_context_size_high": 0.033
},
"supports_function_calling": true,
"supports_response_schema": true,
"source": "https://empiriolabs.ai/models/glm-5-3"
@ -83245,6 +83376,11 @@
"supports_vision": true,
"supports_reasoning": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.033,
"search_context_size_medium": 0.033,
"search_context_size_high": 0.033
},
"supports_function_calling": true,
"supports_response_schema": true,
"source": "https://empiriolabs.ai/models/glm-5-3-flash"
@ -83342,6 +83478,28 @@
"supports_response_schema": true,
"source": "https://empiriolabs.ai/models/kimi-k3"
},
"empiriolabs/kimi-k3-1-pro": {
"litellm_provider": "empiriolabs",
"mode": "chat",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"input_cost_per_token": 2e-06,
"output_cost_per_token": 1e-05,
"cache_read_input_token_cost": 1.8e-06,
"supports_system_messages": true,
"supports_vision": true,
"supports_reasoning": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.015,
"search_context_size_medium": 0.015,
"search_context_size_high": 0.015
},
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_function_calling": true,
"source": "https://empiriolabs.ai/models/kimi-k3-1-pro"
},
"empiriolabs/mimo-v2-6-flash": {
"litellm_provider": "empiriolabs",
"mode": "chat",
@ -83721,6 +83879,11 @@
"supports_vision": true,
"supports_reasoning": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01,
"search_context_size_high": 0.01
},
"supports_response_schema": true,
"source": "https://empiriolabs.ai/models/qwen3-5-122b-a10b"
},
@ -83737,6 +83900,11 @@
"supports_vision": true,
"supports_reasoning": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01,
"search_context_size_high": 0.01
},
"supports_response_schema": true,
"source": "https://empiriolabs.ai/models/qwen3-5-27b"
},
@ -83753,6 +83921,11 @@
"supports_vision": true,
"supports_reasoning": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01,
"search_context_size_high": 0.01
},
"supports_response_schema": true,
"source": "https://empiriolabs.ai/models/qwen3-5-35b-a3b"
},
@ -83769,6 +83942,11 @@
"supports_vision": true,
"supports_reasoning": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01,
"search_context_size_high": 0.01
},
"supports_response_schema": true,
"source": "https://empiriolabs.ai/models/qwen3-5-397b-a17b"
},
@ -83830,6 +84008,10 @@
"max_output_tokens": 65536,
"input_cost_per_token": 2.9e-08,
"output_cost_per_token": 2.87e-07,
"input_cost_per_token_above_128k_tokens": 1.15e-07,
"output_cost_per_token_above_128k_tokens": 1.147e-06,
"input_cost_per_token_above_256k_tokens": 1.72e-07,
"output_cost_per_token_above_256k_tokens": 1.72e-06,
"supports_system_messages": true,
"supports_vision": true,
"supports_reasoning": true,
@ -83852,6 +84034,11 @@
"supports_system_messages": true,
"supports_vision": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.015,
"search_context_size_medium": 0.015,
"search_context_size_high": 0.015
},
"supports_audio_input": true,
"source": "https://empiriolabs.ai/models/qwen3-5-omni-flash"
},
@ -83865,6 +84052,11 @@
"supports_system_messages": true,
"supports_vision": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.015,
"search_context_size_medium": 0.015,
"search_context_size_high": 0.015
},
"supports_response_schema": true,
"supports_audio_input": true,
"source": "https://empiriolabs.ai/models/qwen3-5-omni-plus"
@ -83924,6 +84116,11 @@
"supports_vision": true,
"supports_reasoning": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01,
"search_context_size_high": 0.01
},
"supports_response_schema": true,
"source": "https://empiriolabs.ai/models/qwen3-6-27b"
},
@ -84421,7 +84618,7 @@
"max_output_tokens": 131072,
"input_cost_per_token": 1e-06,
"output_cost_per_token": 2.7e-06,
"cache_read_input_token_cost": 5e-08,
"cache_read_input_token_cost": 1e-07,
"supports_system_messages": true,
"supports_vision": true,
"supports_reasoning": true,
@ -85310,6 +85507,24 @@
"output_cost_per_token": 6.09e-07,
"source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts"
},
"gmi/XiaomiMiMo/MiMo-V2.6-Flash": {
"cache_read_input_token_cost": 2.24e-09,
"input_cost_per_token": 1.12e-07,
"litellm_provider": "gmi",
"max_input_tokens": 1050000,
"mode": "chat",
"output_cost_per_token": 2.24e-07,
"source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts"
},
"gmi/XiaomiMiMo/MiMo-V2.6-Pro": {
"cache_read_input_token_cost": 2.88e-09,
"input_cost_per_token": 3.48e-07,
"litellm_provider": "gmi",
"max_input_tokens": 1050000,
"mode": "chat",
"output_cost_per_token": 6.96e-07,
"source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts"
},
"gmi/deepseek-ai/DeepSeek-Prover-V2-671B": {
"input_cost_per_token": 5e-07,
"litellm_provider": "gmi",
@ -85434,6 +85649,15 @@
"output_cost_per_token": 3.168e-06,
"source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts"
},
"gmi/deepseek-ai/DeepSeek-V4.1-Flash": {
"cache_read_input_token_cost": 3.6e-09,
"input_cost_per_token": 1.8e-07,
"litellm_provider": "gmi",
"max_input_tokens": 1048575,
"mode": "chat",
"output_cost_per_token": 7.2e-07,
"source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts"
},
"gmi/deepseek-ai/deepseek-v4-flash-vision-exp": {
"cache_read_input_token_cost": 1.4e-08,
"input_cost_per_token": 4.4e-07,
@ -85443,6 +85667,24 @@
"output_cost_per_token": 1.32e-06,
"source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts"
},
"gmi/deepseek-ai/deepseek-v4.1-flash-1017": {
"cache_read_input_token_cost": 3.6e-09,
"input_cost_per_token": 1.8e-07,
"litellm_provider": "gmi",
"max_input_tokens": 1048575,
"mode": "chat",
"output_cost_per_token": 7.2e-07,
"source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts"
},
"gmi/gmicloud/Ripple-65B-Preview": {
"cache_read_input_token_cost": 5e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "gmi",
"max_input_tokens": 1000000,
"mode": "chat",
"output_cost_per_token": 6e-07,
"source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts"
},
"gmi/google/gemini-3.1-flash-lite-preview": {
"cache_read_input_token_cost": 2.5e-08,
"input_cost_per_token": 2.5e-07,
@ -85759,6 +86001,14 @@
"output_cost_per_token": 2.5e-07,
"source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts"
},
"gmi/openai/gpt-oss-20b": {
"input_cost_per_token": 4e-08,
"litellm_provider": "gmi",
"max_input_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.5e-07,
"source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts"
},
"gmi/stepfun-ai/Step-3.7-Flash": {
"cache_read_input_token_cost": 4e-08,
"input_cost_per_token": 2e-07,
@ -85803,6 +86053,15 @@
"output_cost_per_token": 6e-07,
"source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts"
},
"gmi/tencent/hy4": {
"cache_read_input_token_cost": 2.1e-08,
"input_cost_per_token": 4.17e-07,
"litellm_provider": "gmi",
"max_input_tokens": 1000000,
"mode": "chat",
"output_cost_per_token": 1.2505e-06,
"source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts"
},
"gmi/tencent/hy4-preview": {
"cache_read_input_token_cost": 4.2e-08,
"input_cost_per_token": 8.34e-07,
@ -86176,5 +86435,100 @@
"output_cost_per_token": 6e-07,
"source": "https://api.together.ai/v1/models",
"supports_prompt_caching": true
},
"vertex_ai/gemini-omni-1.2-flash": {
"input_cost_per_token": 1.5e-06,
"litellm_provider": "vertex_ai",
"max_input_tokens": 131072,
"max_output_tokens": 260640,
"mode": "chat",
"output_cost_per_reasoning_token": 9e-06,
"output_cost_per_token": 9e-06,
"output_cost_per_video_token": 1.2e-05,
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
"supported_endpoints": [
"/v1beta/interactions"
],
"supported_modalities": [
"text",
"image",
"video"
],
"supported_output_modalities": [
"text",
"video"
],
"supports_audio_input": false,
"supports_audio_output": false,
"supports_computer_use": false,
"supports_function_calling": false,
"supports_prompt_caching": false,
"supports_reasoning": true,
"supports_response_schema": false,
"supports_system_messages": false,
"supports_tool_choice": false,
"supports_url_context": false,
"supports_video_input": true,
"supports_vision": true,
"supports_web_search": false
},
"gemini/gemini-omni-1.2-flash": {
"input_cost_per_image_token": 1.5e-06,
"input_cost_per_token": 1.5e-06,
"input_cost_per_video_token": 1.5e-06,
"litellm_provider": "gemini",
"max_input_tokens": 131072,
"max_output_tokens": 65536,
"mode": "chat",
"output_cost_per_reasoning_token": 9e-06,
"output_cost_per_token": 9e-06,
"output_cost_per_video_token": 1.2e-05,
"source": "https://ai.google.dev/gemini-api/docs/pricing",
"supported_endpoints": [
"/v1beta/interactions"
],
"supported_modalities": [
"text",
"image",
"video"
],
"supported_output_modalities": [
"text",
"video",
"audio"
],
"supports_audio_input": false,
"supports_audio_output": true,
"supports_reasoning": true,
"supports_system_messages": false,
"supports_video_input": true,
"supports_vision": true
},
"databricks/databricks-kimi-k3-1-pro": {
"cache_creation_input_token_cost": 1.99997e-06,
"cache_read_input_token_cost": 1.9999e-07,
"input_cost_per_token": 1.99997e-06,
"input_dbu_cost_per_token": 2.8571e-05,
"litellm_provider": "databricks",
"max_input_tokens": 1048576,
"max_output_tokens": 1048576,
"metadata": {"notes":"Input/output cost per token is dbu cost * $0.070. Billing reads the per-token dollar fields; the '*_dbu_cost_per_token' fields are the published Databricks rates, kept for reference."},
"mode": "chat",
"output_cost_per_token": 9.99999e-06,
"output_dbu_cost_per_token": 0.000142857,
"source": "https://www.databricks.com/product/pricing/foundation-model-serving",
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true,
"token_count_framing": 100.0
}
}

128
litellm/models/plugin_upload.py
View file

@ -1,12 +1,25 @@
from collections.abc import AsyncIterator, Mapping
from dataclasses import dataclass
from typing import Protocol
from datetime import datetime
from enum import Enum
from types import MappingProxyType
from typing import Literal, Protocol
def has_agent_skill_metadata(manifest: Mapping[str, object]) -> bool:
return isinstance(manifest.get("agent_skill"), dict)
def uploaded_plugin_source(base_url: str, name: str, digest: str) -> Mapping[str, str]:
return MappingProxyType(
{
"source": "archive",
"url": f"{base_url.rstrip('/')}/claude-code/plugins/{name}/archive/{digest}.zip",
"sha256": digest,
}
)
@dataclass(frozen=True, slots=True)
class PluginUpload:
name: str
@ -27,13 +40,126 @@ class PluginArchiveMetadata:
version: str | None = None
description: str | None = None
has_agent_skill_metadata: bool = False
updated_at: datetime | None = None
@dataclass(frozen=True, slots=True)
class PluginPublication:
name: str
version: str | None
description: str | None
manifest: Mapping[str, object]
created_by: str | None
content: bytes | None = None
@dataclass(frozen=True, slots=True)
class PluginPublicationBase:
id: str
updated_at: datetime | None
@dataclass(frozen=True, slots=True)
class PublishedPlugin:
id: str
name: str
version: str
enabled: bool
updated_at: datetime
@dataclass(frozen=True, slots=True)
class PluginPublicationError:
kind: Literal["invalid_version", "conflict", "not_found"]
message: str
@dataclass(frozen=True, slots=True)
class PluginVersion:
id: str
version: str | None
description: str | None
manifest: Mapping[str, object]
published_at: datetime | None
published_by: str | None
archive_available: bool
@dataclass(frozen=True, slots=True)
class PluginSnapshot:
id: str
name: str
version: str | None
description: str | None
manifest: Mapping[str, object]
enabled: bool
created_at: datetime | None
updated_at: datetime | None
created_by: str | None
@dataclass(frozen=True, slots=True)
class PluginDownloader:
authenticated: bool
user_id: str | None
user_email: str | None
key_alias: str | None
team_id: str | None
team_alias: str | None
class PluginDownloadResult(Enum):
RECORDED = "recorded"
SKILL_NOT_FOUND = "skill_not_found"
@dataclass(frozen=True, slots=True)
class PluginDownloadEvent:
id: str
downloaded_at: datetime
version: str | None
digest: str
authenticated: bool
user_id: str | None
user_email: str | None
key_alias: str | None
team_id: str | None
team_alias: str | None
@dataclass(frozen=True, slots=True)
class PluginDownloadsResponse:
total_downloads: int
anonymous_downloads: int
downloads: tuple[PluginDownloadEvent, ...]
next_cursor: str | None
class PluginUploadStore(Protocol):
async def create(self, upload: PluginUpload) -> str | None: ...
async def get_plugin(self, name: str) -> PluginSnapshot | None: ...
async def publish(
self, publication: PluginPublication, expected: PluginPublicationBase | None = None
) -> PublishedPlugin | PluginPublicationError: ...
async def versions(self, metadata: PluginArchiveMetadata, limit: int, offset: int) -> tuple[PluginVersion, ...]: ...
async def version(self, metadata: PluginArchiveMetadata, version_id: str) -> PluginVersion | None: ...
async def version_content(self, metadata: PluginArchiveMetadata, version_id: str) -> bytes | None: ...
async def get_metadata(self, name: str) -> PluginArchiveMetadata | None: ...
async def get_content(self, metadata: PluginArchiveMetadata) -> bytes | None: ...
def public_metadata(self) -> AsyncIterator[PluginArchiveMetadata]: ...
async def record_download(
self, metadata: PluginArchiveMetadata, digest: str, downloader: PluginDownloader
) -> PluginDownloadResult: ...
async def downloads(self, name: str, limit: int, cursor: int | None) -> PluginDownloadsResponse | None: ...
async def download_counts(self, plugin_ids: tuple[str, ...]) -> Mapping[str, int]: ...

110
litellm/models/skill_submission.py Normal file
View file

@ -0,0 +1,110 @@
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime
from typing import Literal, Protocol, TypeAlias
from litellm.models.plugin_upload import (
PluginPublication,
PluginPublicationBase,
PluginPublicationError,
PublishedPlugin,
)
from litellm.types.proxy.skill_scanning import SkillScan
SubmissionStatus: TypeAlias = Literal["pending_review", "approved", "rejected"]
SubmissionOperation: TypeAlias = Literal["create", "update"]
@dataclass(frozen=True, slots=True)
class SkillProposal:
name: str
version: str | None
description: str | None
manifest: Mapping[str, object]
@dataclass(frozen=True, slots=True)
class SkillSubmission:
id: str
plugin_name: str
operation: SubmissionOperation
status: SubmissionStatus
proposal: SkillProposal
has_archive: bool
submitted_by: str
created_at: datetime
base_plugin_id: str | None
base_updated_at: datetime | None
reviewed_by: str | None
reviewed_at: datetime | None
rejection_reason: str | None
@dataclass(frozen=True, slots=True)
class SubmissionPage:
submissions: tuple[SkillSubmission, ...]
has_more: bool
@dataclass(frozen=True, slots=True)
class SubmissionLimitExceeded:
message: str
@dataclass(frozen=True, slots=True)
class SubmissionMissing:
pass
@dataclass(frozen=True, slots=True)
class SubmissionConflict:
message: str
@dataclass(frozen=True, slots=True)
class SubmissionReviewed:
submission: SkillSubmission
plugin: PublishedPlugin | None
published_source: Mapping[str, str] | None = None
SubmissionReviewResult: TypeAlias = SubmissionReviewed | SubmissionMissing | SubmissionConflict | PluginPublicationError
@dataclass(frozen=True, slots=True)
class SkillScanInput:
scan: SkillScan
source: Mapping[str, str]
content: bytes | None
class SkillSubmissionStore(Protocol):
async def get_scan(self, submission_id: str) -> SkillScan | SubmissionMissing | None: ...
async def start_scan(
self, submission_id: str, scan: SkillScan
) -> SkillScanInput | SubmissionMissing | SubmissionConflict: ...
async def finish_scan(self, submission_id: str, scan: SkillScan) -> None: ...
async def submit(
self, upload: PluginPublication, submitted_by: str, base: PluginPublicationBase | None
) -> SkillSubmission | SubmissionLimitExceeded: ...
async def list(
self, submitted_by: str | None, status: SubmissionStatus | None, limit: int, offset: int
) -> SubmissionPage: ...
async def get(self, submission_id: str, submitted_by: str | None) -> SkillSubmission | None: ...
async def get_archive(self, submission_id: str, submitted_by: str | None) -> bytes | None: ...
async def review(
self,
submission_id: str,
reviewer: str,
approve: bool,
reason: str | None,
*,
publication_base_url: str | None = None,
) -> SubmissionReviewResult: ...

10
litellm/proxy/_lazy_features.py
View file

@ -48,6 +48,15 @@ def _mount_app(prefix: str, attr_name: str = "app") -> Callable[["FastAPI", obje
return _register
def _include_skill_hub(app: "FastAPI", module: object) -> None:
from litellm.proxy.discovery_endpoints.skill_hub_submissions import router
from litellm.proxy.discovery_endpoints.skill_scanning import router as scan_router
_include_router()(app, module)
app.include_router(router)
app.include_router(scan_router)
@dataclass(frozen=True)
class LazyFeature:
name: str
@ -287,6 +296,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
name="skill_hub",
module_path="litellm.proxy.discovery_endpoints.skill_hub_endpoints",
path_prefixes=("/skill-hub",),
register_fn=_include_skill_hub,
),
LazyFeature(
name="claude_code_marketplace",

2004
litellm/proxy/_lazy_openapi_snapshot.json
View file

File diff suppressed because it is too large Load diff

27
litellm/proxy/_types.py
View file

@ -88,6 +88,7 @@ from litellm.types.proxy.carried_budget_state import (
)
from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
from litellm.types.proxy.request_tag_mappings import RequestTagMapping
from litellm.types.proxy.skill_scanning import SkillScannerConfig
from litellm.types.proxy.spend_capture_rate import SpendCaptureRateCheckSettings
from litellm.types.router import AllowedModelRegion, RouterErrors, UpdateRouterConfig
from litellm.types.router_weights import validate_router_settings_dict
@ -449,15 +450,29 @@ class LiteLLMRoutes(enum.Enum):
"/v1/traces/{trace_id}/spans/{span_id}/error",
)
litellm_native_routes = [
skill_hub_routes = (
"/skill-hub/skills",
"/skill-hub/skills/upload",
"/skill-hub/skills/import",
"/skill-hub/skills/{name}",
"/skill-hub/skills/{name}/archive",
"/skill-hub/skills/{name}/content",
"/skill-hub/skills/{plugin_name}/downloads",
"/skill-hub/skills/{name}/versions",
"/skill-hub/skills/{name}/versions/{version_id}",
"/skill-hub/skills/{name}/versions/{version_id}/archive",
"/skill-hub/skills/{name}/publish",
"/skill-hub/skills/{name}/unpublish",
"/skill-hub/skills/{name}/upload",
"/skill-hub/submissions",
"/skill-hub/submissions/{submission_id}",
"/skill-hub/submissions/{submission_id}/archive",
"/skill-hub/submissions/{submission_id}/approve",
"/skill-hub/submissions/{submission_id}/reject",
"/skill-hub/submissions/{submission_id}/scan",
)
litellm_native_routes = [
"/rag/ingest",
"/v1/rag/ingest",
"/rag/query",
@ -476,6 +491,7 @@ class LiteLLMRoutes(enum.Enum):
"/claude-code/plugins",
"/claude-code/plugins/upload",
"/claude-code/plugins/{plugin_name}/archive/{digest}.zip",
"/claude-code/plugins/{plugin_name}/downloads",
"/claude-code/plugins/{plugin_name}",
]
@ -736,6 +752,7 @@ class LiteLLMRoutes(enum.Enum):
+ key_management_routes
+ mcp_management_routes
+ list(agent_management_routes)
+ list(skill_hub_routes)
)
spend_tracking_routes = [
@ -997,6 +1014,7 @@ class LiteLLMRoutes(enum.Enum):
# Agent registry - reads are role-scoped and writes are proxy-admin-gated
# inside agent_endpoints/endpoints.py
*agent_management_routes,
*skill_hub_routes,
] # routes that manage their own allowed/disallowed logic
## Org Admin Routes ##
@ -2538,8 +2556,8 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
"""
use_spend_logs_ledger: bool = Field(
False,
description="Apply request accounting from stored SpendLogs deliveries. Enable only after all accounting writers are upgraded and ledger activation gates pass.",
True,
description="Apply request accounting from stored SpendLogs deliveries. Set to false to roll back to the previous in-memory and Redis spend queues.",
)
request_tag_mappings: Sequence[RequestTagMapping] | None = Field(
@ -3076,6 +3094,9 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="If True, forwards client headers (e.g. Authorization) to the LLM API. Required for Claude Code with Max subscription.",
)
skill_scanner: SkillScannerConfig | None = Field(
None, description="Opt-in external scanner for manual admin review of skill submissions"
)
mcp_required_fields: list[str] | None = Field(
None,
description="List of MCP server fields that must be filled in for a submission to pass standards checks (e.g. ['description', 'source_url', 'alias']).",

160
litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py
View file

@ -27,14 +27,22 @@ from datetime import datetime, timezone
from types import MappingProxyType
from typing import Annotated, Final, Protocol, TypedDict, get_args
from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile
from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile
from fastapi.responses import JSONResponse, Response
from pydantic import Json, TypeAdapter
from typing_extensions import ReadOnly
from typing_extensions import ReadOnly, assert_never
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.url_utils import is_https_url_with_host
from litellm.models.plugin_upload import PluginUpload, has_agent_skill_metadata
from litellm.models.plugin_upload import (
PluginDownloader,
PluginDownloadResult,
PluginDownloadsResponse,
PluginPublication,
PluginPublicationBase,
PluginPublicationError,
has_agent_skill_metadata,
)
from litellm.proxy._types import CommonProxyErrors, ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import (
SkillVisibility,
@ -208,8 +216,8 @@ async def upload_plugin(
}
spec: Final = RegisterPluginRequest.model_validate({**metadata.model_dump(), "source": source})
_validate_installation_preference(spec)
plugin_id: Final = await uploads.create(
PluginUpload(
published: Final = await uploads.publish(
PluginPublication(
name=spec.name,
version=spec.version,
description=spec.description,
@ -218,15 +226,15 @@ async def upload_plugin(
created_by=user_api_key_dict.user_id,
)
)
if plugin_id is None:
raise _name_conflict_error(spec.name)
if isinstance(published, PluginPublicationError):
raise _publication_error(published)
return RegisterPluginResponse(
status="success",
action="created",
plugin=PluginResponse(
id=plugin_id,
id=published.id,
name=spec.name,
version=spec.version,
version=published.version,
description=spec.description,
source=source,
installation_preference=spec.installation_preference,
@ -235,6 +243,27 @@ async def upload_plugin(
)
def require_download_recorded(result: PluginDownloadResult) -> None:
match result:
case PluginDownloadResult.RECORDED:
return
case PluginDownloadResult.SKILL_NOT_FOUND:
raise _error_response(404, "Skill archive not found")
case _:
assert_never(result)
def plugin_downloader(caller: UserAPIKeyAuth | None) -> PluginDownloader:
return PluginDownloader(
authenticated=caller is not None,
user_id=caller.user_id if caller is not None else None,
user_email=caller.user_email if caller is not None else None,
key_alias=caller.key_alias if caller is not None else None,
team_id=caller.team_id if caller is not None else None,
team_alias=caller.team_alias if caller is not None else None,
)
@router.get(
"/claude-code/plugins/{plugin_name}/archive/{digest}.zip",
tags=["Claude Code Marketplace"],
@ -268,6 +297,7 @@ async def download_plugin_archive(
content: Final = await uploads.get_content(archive)
if content is None:
raise _error_response(404, "Skill archive not found")
require_download_recorded(await uploads.record_download(archive, digest, plugin_downloader(caller)))
return Response(
content=content,
media_type="application/zip",
@ -279,6 +309,35 @@ async def download_plugin_archive(
)
@router.get(
"/claude-code/plugins/{plugin_name}/downloads",
tags=["Claude Code Marketplace"],
response_model=PluginDownloadsResponse,
)
async def get_plugin_downloads(
plugin_name: str,
uploads: PluginUploadRepositoryDep,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
limit: Annotated[int, Query(ge=1, le=100, description="Maximum number of download events")] = 25,
cursor: Annotated[str | None, Query(pattern=r"^[1-9][0-9]{0,18}$", description="Previous next_cursor")] = None,
) -> PluginDownloadsResponse:
"""Admin-only ZIP download totals and authenticated downloader history.
Anonymous downloads contribute only to totals. Counts include repeat requests and client prefetches,
starting when tracking is deployed. External sources and installation completion are not observed.
"""
_require_proxy_admin(user_api_key_dict)
if uploads is None:
raise _error_response(500, CommonProxyErrors.db_not_connected_error.value)
before: Final = int(cursor) if cursor is not None else None
if before is not None and before > 9223372036854775807:
raise _error_response(422, "Invalid download cursor")
result: Final = await uploads.downloads(plugin_name, limit, before)
if result is None:
raise _error_response(404, "Plugin not found")
return result
@router.get(
"/claude-code/marketplace.json",
tags=["Claude Code Marketplace"],
@ -472,10 +531,14 @@ def _error_response(status_code: int, message: str) -> HTTPException:
return HTTPException(status_code=status_code, detail={"error": message})
def _name_conflict_error(name: str) -> HTTPException:
return _error_response(
409, f"A skill named '{name}' already exists. Update the existing skill instead of adding it again."
)
def _publication_error(error: PluginPublicationError) -> HTTPException:
match error.kind:
case "invalid_version":
return _error_response(400, error.message)
case "not_found":
return _error_response(404, error.message)
case "conflict":
return _error_response(409, error.message)
def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None:
@ -498,6 +561,7 @@ def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None:
)
async def register_plugin(
request: RegisterPluginRequest,
uploads: PluginUploadRepositoryDep,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
@ -540,12 +604,11 @@ async def register_plugin(
}'
```
"""
from prisma.errors import UniqueViolationError
try:
_require_proxy_admin(user_api_key_dict)
prisma_client: Final = await get_prisma_client()
if uploads is None:
raise _error_response(500, CommonProxyErrors.db_not_connected_error.value)
if not re.match(r"^[a-z0-9-]+$", request.name):
raise HTTPException(
@ -556,30 +619,19 @@ async def register_plugin(
_validate_plugin_source(request.source)
_validate_installation_preference(request)
existing: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique(
where={"name": request.name}
)
if existing:
raise _name_conflict_error(request.name)
manifest: Final[Mapping[str, object]] = _build_plugin_manifest(request.name, request)
try:
plugin: Final[_PluginRecord] = await ClaudeCodePluginRepository(prisma_client).table.create(
data={
"name": request.name,
"version": request.version,
"description": request.description,
"manifest_json": json.dumps(manifest),
"files_json": "{}",
"enabled": True,
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
"created_by": user_api_key_dict.user_id,
}
plugin: Final = await uploads.publish(
PluginPublication(
name=request.name,
version=request.version,
description=request.description,
manifest=manifest,
created_by=user_api_key_dict.user_id,
)
except UniqueViolationError:
raise _name_conflict_error(request.name)
)
if isinstance(plugin, PluginPublicationError):
raise _publication_error(plugin)
verbose_proxy_logger.info("Plugin %s created successfully", request.name)
@ -590,7 +642,7 @@ async def register_plugin(
id=plugin.id,
name=plugin.name,
version=plugin.version,
description=plugin.description,
description=request.description,
source=request.source,
installation_preference=request.installation_preference,
enabled=plugin.enabled,
@ -755,6 +807,7 @@ async def get_plugin(
async def update_plugin(
plugin_name: str,
request: UpdatePluginRequest,
uploads: PluginUploadRepositoryDep,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
@ -763,7 +816,8 @@ async def update_plugin(
The plugin is identified by its name in the path, which is the resource
identity and cannot be changed here. This is a full replace, not a merge:
the manifest is rebuilt from the request body, so any optional field left
out is reset to its default (e.g. an omitted version is cleared, not kept).
out is reset to its default, except version which automatically increments
the published patch version when omitted or null.
Send the full desired state.
Returns 404 if no plugin with the given name exists; use
@ -801,6 +855,8 @@ async def update_plugin(
try:
_require_proxy_admin(user_api_key_dict)
if uploads is None:
raise _error_response(500, CommonProxyErrors.db_not_connected_error.value)
prisma_client: Final = await get_prisma_client()
@ -808,6 +864,8 @@ async def update_plugin(
existing: Final[_PluginRecord | None] = await repository.table.find_unique(where={"name": plugin_name})
if not existing:
raise _error_response(404, f"Plugin '{plugin_name}' not found")
if request.expected_updated_at is not None and request.expected_updated_at != existing.updated_at:
raise _error_response(409, "The published skill changed. Reload it before publishing an update")
current_manifest: Final = MANIFEST_FIELDS.validate_json(existing.manifest_json or "{}")
if request.source != current_manifest.get("source"):
@ -823,18 +881,18 @@ async def update_plugin(
),
}
plugin: Final[_PluginRecord | None] = await repository.table.update(
where={"name": plugin_name},
data={
"version": request.version,
"description": request.description,
"manifest_json": json.dumps(manifest),
"files_json": "{}",
"updated_at": datetime.now(timezone.utc),
},
plugin: Final = await uploads.publish(
PluginPublication(
name=plugin_name,
version=request.version,
description=request.description,
manifest=manifest,
created_by=user_api_key_dict.user_id,
),
PluginPublicationBase(existing.id, existing.updated_at),
)
if plugin is None:
raise _error_response(404, f"Plugin '{plugin_name}' not found")
if isinstance(plugin, PluginPublicationError):
raise _publication_error(plugin)
verbose_proxy_logger.info("Plugin %s updated successfully", plugin_name)
@ -845,7 +903,7 @@ async def update_plugin(
id=plugin.id,
name=plugin.name,
version=plugin.version,
description=plugin.description,
description=request.description,
source=request.source,
installation_preference=request.installation_preference,
enabled=plugin.enabled,

5
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -909,13 +909,14 @@ class ResetBudgetJob:
db_span("reset_budget_cascade", "LiteLLM_BudgetTable"),
budget_cascade_unit_of_work(self._new_batch) as uow,
):
uow.row_locks.queue_row_locks(cascade.budget_ids)
_queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE)
uow.team_memberships.queue_total_spend_repair(cascade.budget_ids)
_queue_budget_linked_resets(uow.team_memberships, cascade)
_queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE)
_queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE)
_queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE)
_queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE)
_queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE)
_queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE)
for budget_id, budget_reset_at in cascade.budget_resets:
uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at)

266
litellm/proxy/db/db_spend_update_writer.py
View file

@ -328,13 +328,13 @@ _SpendTableName = Literal[
"agent_list_transactions",
]
_SPEND_TABLE_COMMIT_ORDER: Final[tuple[_SpendTableName, ...]] = (
"user_list_transactions",
"end_user_list_transactions",
"key_list_transactions",
"team_list_transactions",
"team_member_list_transactions",
"org_list_transactions",
"org_member_list_transactions",
"team_list_transactions",
"user_list_transactions",
"end_user_list_transactions",
"team_member_list_transactions",
"key_list_transactions",
"project_list_transactions",
"tag_list_transactions",
"model_access_group_list_transactions",
@ -563,6 +563,24 @@ SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend,
"""
def _org_members_in_key_order(spend_by_member_key: Mapping[str, float]) -> tuple[tuple[tuple[str, str], float], ...]:
"""Sorted by the decoded (organization_id, user_id), the order every other writer locks these rows in, since
the percent-encoded key does not sort the same way."""
decoded: Final = (
((unquote(key.split("::")[1]), unquote(key.split("::")[3])), cost) for key, cost in spend_by_member_key.items()
)
return tuple(sorted(decoded))
def _charged_team_ids(updates: DBSpendUpdateTransactions) -> tuple[str, ...]:
members: Final = updates.get("team_member_list_transactions") or {}
return tuple(
sorted(
{*(updates.get("team_list_transactions") or {}), *(key.split("::")[1] for key in members)},
)
)
async def _write_team_member_spend(transaction: _SpendTransaction, spend_by_member_key: Mapping[str, float]) -> None:
# key is "team_id::<value>::user_id::<value>"; locks are taken in sorted team_id order like the team endpoints
rows: Final = sorted((key.split("::")[1], key.split("::")[3], cost) for key, cost in spend_by_member_key.items())
@ -715,7 +733,7 @@ class DBSpendUpdateWriter:
proxy_settings: Final = cast( # cast-ok: proxy settings use string keys
Mapping[str, object], general_settings
)
use_ledger: Final = proxy_settings.get("use_spend_logs_ledger") is True
use_ledger: Final = proxy_settings.get("use_spend_logs_ledger", True) is not False
request_fields: Final = _SPEND_METADATA_ADAPTER.validate_python(kwargs or {})
if kwargs is not None:
kwargs[ACCOUNTING_LEDGER_KWARG] = use_ledger # rebind-ok: spend counters leave ledger windows to the worker
@ -866,8 +884,12 @@ class DBSpendUpdateWriter:
payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs
)
if "accounting_payload" in payload:
from litellm.proxy.utils import accounting_only_spend_log
await self._insert_spend_log_to_db(
payload={**payload, "accounting_log_disabled": disable_spend_logs},
payload=accounting_only_spend_log(payload)
if disable_spend_logs
else {**payload, "accounting_log_disabled": False},
prisma_client=prisma_client,
llm_router=llm_router,
)
@ -2301,6 +2323,10 @@ class DBSpendUpdateWriter:
transaction: _SpendTransaction,
updates: DBSpendUpdateTransactions,
) -> None:
"""Every team advisory lock is taken before the first row lock, the order the team roster writers use, so
applying a batch of charges cannot deadlock with a concurrent roster change."""
for team_id in _charged_team_ids(updates):
_ = await transaction.execute_raw(_TEAM_ADVISORY_LOCK_SQL, team_id)
await self._commit_spend_updates_to_db(
prisma_client, 0, proxy_logging_obj, updates, outer_transaction=transaction
)
@ -2320,6 +2346,104 @@ class DBSpendUpdateWriter:
"""
from litellm.proxy.utils import ProxyUpdateSpend
### UPDATE ORG TABLE ###
org_list_transactions: Final = db_spend_update_transactions["org_list_transactions"]
verbose_proxy_logger.debug("Org Spend transactions: %s", org_list_transactions)
if org_list_transactions is not None and len(org_list_transactions.keys()) > 0:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(
prisma_client, "LiteLLM_OrganizationTable", outer_transaction=outer_transaction
) as transaction:
async with transaction.batch_() as batcher:
# Sort by org_id for consistent lock ordering across pods to prevent deadlocks.
for org_id, response_cost in sorted(org_list_transactions.items()):
batcher.litellm_organizationtable.update_many( # 'update_many' prevents error from being raised if no row exists
where={"organization_id": org_id},
data={"spend": {"increment": response_cost}},
)
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
if on_table_committed is not None:
on_table_committed("org_list_transactions")
org_member_list_transactions: Final = db_spend_update_transactions.get("org_member_list_transactions")
verbose_proxy_logger.debug("Org Membership Spend transactions: %s", org_member_list_transactions)
if org_member_list_transactions is not None and len(org_member_list_transactions.keys()) > 0:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with (
_spend_update_tx(
prisma_client, "LiteLLM_OrganizationMembership", outer_transaction=outer_transaction
) as transaction,
transaction.batch_() as batcher,
):
for (org_id, member_id), response_cost in _org_members_in_key_order(
org_member_list_transactions
):
batcher.litellm_organizationmembership.update_many(
where={"organization_id": org_id, "user_id": member_id},
data={"spend": {"increment": response_cost}},
)
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
if on_table_committed is not None:
on_table_committed("org_member_list_transactions")
### UPDATE TEAM TABLE ###
team_list_transactions: Final = db_spend_update_transactions["team_list_transactions"]
verbose_proxy_logger.debug("Team Spend transactions: %s", team_list_transactions)
if team_list_transactions is not None and len(team_list_transactions.keys()) > 0:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(
prisma_client, "LiteLLM_TeamTable", outer_transaction=outer_transaction
) as transaction:
async with transaction.batch_() as batcher:
# Sort by team_id for consistent lock ordering across pods to prevent deadlocks.
for team_id, response_cost in sorted(team_list_transactions.items()):
verbose_proxy_logger.debug(
"Updating spend for team id=%s by %s", team_id, response_cost
)
batcher.litellm_teamtable.update_many( # 'update_many' prevents error from being raised if no row exists
where={"team_id": team_id},
data={"spend": {"increment": response_cost}},
)
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
if on_table_committed is not None:
on_table_committed("team_list_transactions")
### UPDATE USER TABLE ###
user_list_transactions: Final = db_spend_update_transactions["user_list_transactions"]
verbose_proxy_logger.debug("User Spend transactions: %s", user_list_transactions)
@ -2366,75 +2490,6 @@ class DBSpendUpdateWriter:
)
if on_table_committed is not None:
on_table_committed("end_user_list_transactions")
### UPDATE KEY TABLE ###
key_list_transactions: Final = db_spend_update_transactions["key_list_transactions"]
verbose_proxy_logger.debug("KEY Spend transactions: %s", key_list_transactions)
if key_list_transactions is not None and len(key_list_transactions.keys()) > 0:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(
prisma_client, "LiteLLM_VerificationToken", outer_transaction=outer_transaction
) as transaction:
async with transaction.batch_() as batcher:
# Sort by token for consistent lock ordering across pods to prevent deadlocks.
for token, response_cost in sorted(key_list_transactions.items()):
spend_increment: _SpendIncrement = {"increment": response_cost}
batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists
where={"token": token},
data={
"spend": spend_increment,
"total_spend": spend_increment,
"last_active": datetime.now(timezone.utc),
},
)
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
if on_table_committed is not None:
on_table_committed("key_list_transactions")
### UPDATE TEAM TABLE ###
team_list_transactions: Final = db_spend_update_transactions["team_list_transactions"]
verbose_proxy_logger.debug("Team Spend transactions: %s", team_list_transactions)
if team_list_transactions is not None and len(team_list_transactions.keys()) > 0:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(
prisma_client, "LiteLLM_TeamTable", outer_transaction=outer_transaction
) as transaction:
async with transaction.batch_() as batcher:
# Sort by team_id for consistent lock ordering across pods to prevent deadlocks.
for team_id, response_cost in sorted(team_list_transactions.items()):
verbose_proxy_logger.debug(
"Updating spend for team id=%s by %s", team_id, response_cost
)
batcher.litellm_teamtable.update_many( # 'update_many' prevents error from being raised if no row exists
where={"team_id": team_id},
data={"spend": {"increment": response_cost}},
)
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
if on_table_committed is not None:
on_table_committed("team_list_transactions")
### UPDATE TEAM Membership TABLE with spend ###
team_member_list_transactions: Final = db_spend_update_transactions["team_member_list_transactions"]
@ -2485,22 +2540,27 @@ class DBSpendUpdateWriter:
elif on_table_committed is not None:
on_table_committed("team_member_list_transactions")
### UPDATE ORG TABLE ###
org_list_transactions: Final = db_spend_update_transactions["org_list_transactions"]
verbose_proxy_logger.debug("Org Spend transactions: %s", org_list_transactions)
if org_list_transactions is not None and len(org_list_transactions.keys()) > 0:
### UPDATE KEY TABLE ###
key_list_transactions: Final = db_spend_update_transactions["key_list_transactions"]
verbose_proxy_logger.debug("KEY Spend transactions: %s", key_list_transactions)
if key_list_transactions is not None and len(key_list_transactions.keys()) > 0:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(
prisma_client, "LiteLLM_OrganizationTable", outer_transaction=outer_transaction
prisma_client, "LiteLLM_VerificationToken", outer_transaction=outer_transaction
) as transaction:
async with transaction.batch_() as batcher:
# Sort by org_id for consistent lock ordering across pods to prevent deadlocks.
for org_id, response_cost in sorted(org_list_transactions.items()):
batcher.litellm_organizationtable.update_many( # 'update_many' prevents error from being raised if no row exists
where={"organization_id": org_id},
data={"spend": {"increment": response_cost}},
# Sort by token for consistent lock ordering across pods to prevent deadlocks.
for token, response_cost in sorted(key_list_transactions.items()):
spend_increment: _SpendIncrement = {"increment": response_cost}
batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists
where={"token": token},
data={
"spend": spend_increment,
"total_spend": spend_increment,
"last_active": datetime.now(timezone.utc),
},
)
break
except Exception as e:
@ -2514,39 +2574,7 @@ class DBSpendUpdateWriter:
proxy_logging_obj=proxy_logging_obj,
)
if on_table_committed is not None:
on_table_committed("org_list_transactions")
org_member_list_transactions: Final = db_spend_update_transactions.get("org_member_list_transactions")
verbose_proxy_logger.debug("Org Membership Spend transactions: %s", org_member_list_transactions)
if org_member_list_transactions is not None and len(org_member_list_transactions.keys()) > 0:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with (
_spend_update_tx(
prisma_client, "LiteLLM_OrganizationMembership", outer_transaction=outer_transaction
) as transaction,
transaction.batch_() as batcher,
):
for key, response_cost in sorted(org_member_list_transactions.items()):
_, quoted_org_id, _, quoted_user_id = key.split("::")
batcher.litellm_organizationmembership.update_many(
where={"organization_id": unquote(quoted_org_id), "user_id": unquote(quoted_user_id)},
data={"spend": {"increment": response_cost}},
)
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
if on_table_committed is not None:
on_table_committed("org_member_list_transactions")
on_table_committed("key_list_transactions")
### UPDATE PROJECT TABLE ###
project_list_transactions: Final = db_spend_update_transactions.get("project_list_transactions")

11
litellm/proxy/deps/repositories.py
View file

@ -6,6 +6,7 @@ from litellm.models.deployment_warning import DeploymentWarningIgnoreStore
from litellm.models.gateway_activity import GatewayActivityStore
from litellm.models.issues import IssueActivityStore, IssueStore
from litellm.models.plugin_upload import PluginUploadStore
from litellm.models.skill_submission import SkillSubmissionStore
from litellm.models.usage_alerts import UsageAlertStore
from litellm.models.user_banner import UserBannerStore
from litellm.proxy.deps.database import DatabaseDep
@ -14,6 +15,7 @@ from litellm.repositories.gateway_activity_repository import GatewayActivityRepo
from litellm.repositories.issue_activity import IssueActivityRepository
from litellm.repositories.issues import IssuesRepository
from litellm.repositories.plugin_upload_repository import PluginUploadRepository
from litellm.repositories.skill_submission_repository import SkillSubmissionRepository
from litellm.repositories.usage_alerts import UsageAlertsRepository
from litellm.repositories.user_banner_repository import UserBannerRepository
@ -50,6 +52,15 @@ async def get_plugin_upload_repository(database: DatabaseDep) -> PluginUploadSto
PluginUploadRepositoryDep: TypeAlias = Annotated[PluginUploadStore | None, Depends(get_plugin_upload_repository)]
async def get_skill_submission_repository(database: DatabaseDep) -> SkillSubmissionStore | None:
return None if database is None else SkillSubmissionRepository(database.prisma)
SkillSubmissionRepositoryDep: TypeAlias = Annotated[
SkillSubmissionStore | None, Depends(get_skill_submission_repository)
]
async def get_usage_alerts_repository(database: DatabaseDep) -> UsageAlertStore | None:
return None if database is None else UsageAlertsRepository.from_prisma(database.prisma)

409
litellm/proxy/discovery_endpoints/skill_hub_endpoints.py
View file

@ -1,5 +1,8 @@
import asyncio
import hashlib
from collections.abc import Mapping
from dataclasses import dataclass, replace
from datetime import datetime
from typing import Annotated, Final
from urllib.parse import quote
@ -10,18 +13,31 @@ import litellm
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared HTTP factory has untyped optional params
)
from litellm.models.plugin_upload import PluginArchiveMetadata, PluginUpload, PluginUploadStore
from litellm.models.plugin_upload import (
PluginArchiveMetadata,
PluginDownloader,
PluginDownloadsResponse,
PluginPublication,
PluginSnapshot,
PluginUploadStore,
PluginVersion,
PublishedPlugin,
has_agent_skill_metadata,
)
from litellm.models.skill_submission import SkillSubmission, SkillSubmissionStore
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace import (
MANIFEST_FIELDS,
delete_plugin,
disable_plugin,
enable_plugin,
get_plugin_downloads,
header_key_caller,
list_plugins,
plugin_archive_source_matches,
plugin_archive_url,
update_plugin,
plugin_downloader,
require_download_recorded,
)
from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import resolve_skill_visibility
from litellm.proxy.anthropic_endpoints.claude_code_endpoints.skill_upload_archive import build_uploaded_skill_archive
@ -31,11 +47,17 @@ from litellm.proxy.common_utils.claude_code_marketplace_auth import (
marketplace_key_header_present,
)
from litellm.proxy.common_utils.resource_ownership import is_proxy_admin
from litellm.proxy.deps.repositories import PluginUploadRepositoryDep
from litellm.proxy.deps.repositories import PluginUploadRepositoryDep, SkillSubmissionRepositoryDep
from litellm.proxy.discovery_endpoints.agent_skills_archive import (
SkillArchive,
build_agent_skill_archive,
)
from litellm.proxy.discovery_endpoints.skill_hub_mutations import (
SkillMutationError,
archive_writer,
write_skill,
write_target,
)
from litellm.proxy.discovery_endpoints.skill_import import (
MAX_IMPORT_BYTES,
SkillImportError,
@ -43,18 +65,33 @@ from litellm.proxy.discovery_endpoints.skill_import import (
parse_skill_source,
select_imported_skill,
)
from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3
from litellm.types.proxy.claude_code_endpoints import PluginListItem, UpdatePluginRequest
from litellm.types.proxy.claude_code_endpoints import PluginListItem
from litellm.types.proxy.discovery_endpoints.agent_skills_endpoints import AgentSkillsIndexEntry
from litellm.types.proxy.skill_hub import (
ImportSkillRequest,
SkillHubList,
SkillHubMetadata,
SkillHubMutationResponse,
SkillHubPendingSkill,
SkillHubProposal,
SkillHubPublication,
SkillHubSkill,
SkillHubVersion,
SkillHubVersions,
)
router: Final = APIRouter(prefix="/skill-hub/skills", tags=["Skills Hub"])
_EDITABLE_METADATA: Final = frozenset(SkillHubMetadata.model_fields)
router.add_api_route(
"/{plugin_name}/downloads", get_plugin_downloads, methods=["GET"], response_model=PluginDownloadsResponse
)
@dataclass(frozen=True, slots=True)
class VisibleSkillArchive:
archive: SkillArchive
metadata: PluginArchiveMetadata
downloader: PluginDownloader
def _url(request: Request, route: str) -> str:
@ -67,7 +104,7 @@ def _download_url(request: Request, name: str) -> str:
return _url(request, f"skill-hub/skills/{quote(name, safe='')}/archive")
def _skill_response(request: Request, plugin: PluginListItem) -> SkillHubSkill:
def _skill_response(request: Request, plugin: PluginListItem, download_count: int | None = None) -> SkillHubSkill:
local_archive: Final = plugin_archive_source_matches(
request,
plugin.name,
@ -80,16 +117,22 @@ def _skill_response(request: Request, plugin: PluginListItem) -> SkillHubSkill:
name=plugin.name,
description=plugin.description,
version=plugin.version,
domain=plugin.domain,
namespace=plugin.namespace,
keywords=None if plugin.keywords is None else tuple(plugin.keywords),
updated_at=None if plugin.updated_at is None else datetime.fromisoformat(plugin.updated_at),
visibility="public" if plugin.enabled else "private",
source={"source": "archive", "url": _download_url(request, plugin.name)} if local_archive else plugin.source,
download_url=_download_url(request, plugin.name) if local_archive else None,
content_url=_url(request, f"skill-hub/skills/{quote(plugin.name, safe='')}/content") if local_archive else None,
download_count=download_count,
)
@router.get("", response_model=SkillHubList)
async def list_skills(
request: Request,
uploads: PluginUploadRepositoryDep,
caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
query: str = "",
offset: Annotated[int, Query(ge=0)] = 0,
@ -103,8 +146,14 @@ async def list_skills(
for plugin in sorted(catalog.plugins, key=lambda item: item.name)
if needle in " ".join((plugin.name, plugin.description or "", *(plugin.keywords or ()))).casefold()
)
page: Final = matched[offset : offset + limit]
counts: Final[Mapping[str, int]] = (
await _require_store(uploads).download_counts(tuple(plugin.id for plugin in page))
if is_proxy_admin(caller)
else {}
)
return SkillHubList(
skills=tuple(_skill_response(request, plugin) for plugin in matched[offset : offset + limit]),
skills=tuple(_skill_response(request, plugin, counts.get(plugin.id)) for plugin in page),
count=len(matched),
next_offset=offset + limit if offset + limit < len(matched) else None,
)
@ -122,22 +171,28 @@ async def get_skill(
metadata: Final = await store.get_metadata(name)
if metadata is None or not (await resolve_skill_visibility(caller)).allows(metadata):
raise HTTPException(404, "Skill not found")
return _metadata_response(request, metadata)
return skill_metadata_response(request, metadata)
def _metadata_response(request: Request, metadata: PluginArchiveMetadata) -> SkillHubSkill:
def skill_metadata_response(request: Request, metadata: PluginArchiveMetadata) -> SkillHubSkill:
manifest: Final = MANIFEST_FIELDS.validate_json(metadata.manifest_json or "{}")
return _skill_response(
request,
PluginListItem(
id=metadata.id or metadata.name,
name=metadata.name,
description=metadata.description,
version=metadata.version,
source=dict(metadata.source),
has_agent_skill_metadata=metadata.has_agent_skill_metadata,
enabled=metadata.enabled,
created_at=None,
updated_at=None,
PluginListItem.model_validate(
{
"id": metadata.id or metadata.name,
"name": metadata.name,
"description": metadata.description,
"version": metadata.version,
"source": dict(metadata.source),
"has_agent_skill_metadata": metadata.has_agent_skill_metadata,
"enabled": metadata.enabled,
"created_at": None,
"updated_at": metadata.updated_at.isoformat() if metadata.updated_at is not None else None,
"domain": manifest.get("domain"),
"namespace": manifest.get("namespace"),
"keywords": manifest.get("keywords"),
}
),
)
@ -153,57 +208,33 @@ def _require_admin(caller: UserAPIKeyAuth) -> None:
raise HTTPException(403, "Only proxy admins may publish skills")
async def skill_publisher(
caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> UserAPIKeyAuth:
from litellm.proxy.proxy_server import proxy_logging_obj
_require_admin(caller)
limiter: Final = proxy_logging_obj.get_proxy_hook("parallel_request_limiter")
if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3):
raise HTTPException(503, "Skills Hub publishing requires the default v3 rate limiter")
key_identity: Final = caller.api_key or caller.token
identity: Final = (
f"user:{caller.user_id}" if caller.user_id else f"key:{key_identity}" if key_identity else "local-admin"
)
result: Final = await limiter.atomic_check_and_increment_by_n(
descriptors=[
{
"key": "skill_hub_publish",
"value": hashlib.sha256(identity.encode()).hexdigest(),
"rate_limit": {"requests_per_unit": 10, "window_size": 60},
}
],
increments=[{"requests": 1}],
parent_otel_span=caller.parent_otel_span,
)
if result["overall_code"] == "OVER_LIMIT":
raise HTTPException(
429, "Skills Hub allows 10 uploads or imports per minute per user", headers={"Retry-After": "60"}
)
return caller
@router.put("/{name}", response_model=SkillHubSkill)
@router.put("/{name}", response_model=SkillHubMutationResponse)
async def replace_skill_metadata(
request: Request,
name: str,
body: SkillHubMetadata,
uploads: PluginUploadRepositoryDep,
submissions: SkillSubmissionRepositoryDep,
caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> SkillHubSkill:
"""Replace registry metadata, clearing omitted fields and preserving skill files."""
_require_admin(caller)
store: Final = _require_store(uploads)
metadata: Final = await store.get_metadata(name)
if metadata is None:
) -> SkillHubMutationResponse:
"""Replace registry metadata, preserving files. Members submit an immutable proposal for admin review."""
existing: Final = await write_target(name, True, uploads, caller, body.expected_updated_at)
if isinstance(existing, SkillMutationError):
raise HTTPException(existing.status, existing.message)
if existing is None:
raise HTTPException(404, "Skill not found")
manifest: Final = MANIFEST_FIELDS.validate_json(metadata.manifest_json or "{}")
updated: Final = UpdatePluginRequest.model_validate(
{**manifest, **body.model_dump(), "source": dict(metadata.source)}
publication: Final = PluginPublication(
name=name,
version=body.version,
description=body.description,
manifest={
**{key: value for key, value in existing.manifest.items() if key not in _EDITABLE_METADATA},
**body.model_dump(exclude_none=True),
"name": name,
},
created_by=caller.user_id,
)
await update_plugin(plugin_name=name, request=updated, user_api_key_dict=caller)
return await get_skill(request, name, uploads, caller)
return await _mutate(request, publication, existing, _require_store(uploads), submissions, caller)
@router.post("/{name}/publish", response_model=SkillHubSkill)
@ -242,14 +273,15 @@ async def remove_skill(name: str, caller: Annotated[UserAPIKeyAuth, Depends(user
return Response(status_code=204)
@router.post("/upload", response_model=SkillHubSkill)
@router.post("/upload", response_model=SkillHubMutationResponse)
async def upload_skill(
request: Request,
uploads: PluginUploadRepositoryDep,
caller: Annotated[UserAPIKeyAuth, Depends(skill_publisher)],
submissions: SkillSubmissionRepositoryDep,
caller: Annotated[UserAPIKeyAuth, Depends(archive_writer)],
file: Annotated[UploadFile, File(description="Agent Skills ZIP with SKILL.md and optional supporting files")],
version: Annotated[str | None, Form()] = None,
) -> SkillHubSkill:
) -> SkillHubMutationResponse:
"""Publish a portable skill while retaining Claude marketplace compatibility."""
store: Final = _require_store(uploads)
content: Final = await file.read(MAX_IMPORT_BYTES + 1)
@ -258,21 +290,22 @@ async def upload_skill(
archive: Final = await asyncio.to_thread(build_agent_skill_archive, content)
if archive is None:
raise HTTPException(422, "Upload a valid Agent Skill with name and description in SKILL.md")
return await _publish(request, store, caller, archive, version, None)
return await _publish(request, store, submissions, caller, archive, version, None)
def import_http_client() -> httpx.AsyncClient:
return get_async_httpx_client(llm_provider="skill_hub_import").client
@router.post("/import", response_model=SkillHubSkill)
@router.post("/import", response_model=SkillHubMutationResponse)
async def import_skill(
request: Request,
body: ImportSkillRequest,
uploads: PluginUploadRepositoryDep,
caller: Annotated[UserAPIKeyAuth, Depends(skill_publisher)],
submissions: SkillSubmissionRepositoryDep,
caller: Annotated[UserAPIKeyAuth, Depends(archive_writer)],
client: Annotated[httpx.AsyncClient, Depends(import_http_client)],
) -> SkillHubSkill:
) -> SkillHubMutationResponse:
"""Import a public skills.sh or GitHub skill as a portable snapshot."""
store: Final = _require_store(uploads)
source: Final = parse_skill_source(body.url, skill=body.skill, ref=body.ref, path=body.path)
@ -284,17 +317,122 @@ async def import_skill(
archive: Final = await asyncio.to_thread(select_imported_skill, content, source)
if isinstance(archive, SkillImportError):
raise HTTPException(archive.status, archive.message)
return await _publish(request, store, caller, archive, body.version, body.url)
return await _publish(request, store, submissions, caller, archive, body.version, body.url)
async def _publish(
request: Request,
store: PluginUploadStore,
submissions: SkillSubmissionStore | None,
caller: UserAPIKeyAuth,
archive: SkillArchive,
version: str | None,
origin: str | None,
) -> SkillHubMutationResponse:
existing: Final = await write_target(archive.declared_name or "", False, store, caller)
if isinstance(existing, SkillMutationError):
raise HTTPException(existing.status, existing.message)
publication: Final = await prepare_skill_publication(request, archive, version, origin, caller.user_id)
return await _mutate(request, publication, existing, store, submissions, caller)
async def _mutate(
request: Request,
publication: PluginPublication,
existing: PluginSnapshot | None,
store: PluginUploadStore,
submissions: SkillSubmissionStore | None,
caller: UserAPIKeyAuth,
) -> SkillHubMutationResponse:
result: Final = await write_skill(publication, existing, store, submissions, caller)
if isinstance(result, SkillMutationError):
raise HTTPException(result.status, result.message)
if isinstance(result, SkillSubmission):
return SkillHubPendingSkill(submission_id=result.id, proposal=skill_proposal_response(request, result))
return skill_publication_response(request, publication, result)
def skill_proposal_response(request: Request, submission: SkillSubmission) -> SkillHubProposal:
proposal: Final = submission.proposal
return SkillHubProposal.model_validate(
{
**proposal.manifest,
"name": proposal.name,
"description": proposal.description,
"version": proposal.version,
**(
{
"source": {
"source": "archive",
"url": _url(request, f"skill-hub/submissions/{quote(submission.id, safe='')}/archive"),
}
}
if submission.has_archive
else {}
),
}
)
@router.put("/{name}/upload", response_model=SkillHubMutationResponse)
async def replace_skill_archive(
request: Request,
name: str,
uploads: PluginUploadRepositoryDep,
submissions: SkillSubmissionRepositoryDep,
caller: Annotated[UserAPIKeyAuth, Depends(archive_writer)],
file: Annotated[UploadFile, File(description="Replacement Agent Skills ZIP with the same declared skill name")],
version: Annotated[str | None, Form()] = None,
expected_updated_at: Annotated[datetime | None, Form()] = None,
) -> SkillHubMutationResponse:
"""Replace skill files, retaining registry fields. Members propose an update against the original editor timestamp."""
existing: Final = await write_target(name, True, uploads, caller, expected_updated_at)
if isinstance(existing, SkillMutationError):
raise HTTPException(existing.status, existing.message)
if existing is None:
raise HTTPException(404, "Skill not found")
content: Final = await file.read(MAX_IMPORT_BYTES + 1)
if len(content) > MAX_IMPORT_BYTES:
raise HTTPException(413, "Skill ZIP must be 10 MiB or smaller")
archive: Final = await asyncio.to_thread(build_agent_skill_archive, content)
if archive is None or archive.declared_name != name:
raise HTTPException(422, "Replacement SKILL.md must declare the existing skill name and a description")
prepared: Final = await prepare_skill_publication(request, archive, version, None, caller.user_id)
publication: Final = replace(prepared, manifest={**existing.manifest, **prepared.manifest})
return await _mutate(request, publication, existing, _require_store(uploads), submissions, caller)
def skill_publication_response(
request: Request, publication: PluginPublication, published: PublishedPlugin
) -> SkillHubSkill:
return _skill_response(
request,
PluginListItem.model_validate(
{
"source": publication.manifest.get("source", {}),
"id": published.id,
"name": published.name,
"description": publication.description,
"version": published.version,
"enabled": published.enabled,
"has_agent_skill_metadata": has_agent_skill_metadata(publication.manifest),
"created_at": None,
"updated_at": published.updated_at.isoformat(),
"domain": publication.manifest.get("domain"),
"namespace": publication.manifest.get("namespace"),
"keywords": publication.manifest.get("keywords"),
}
),
)
async def prepare_skill_publication(
request: Request,
archive: SkillArchive,
version: str | None,
origin: str | None,
created_by: str | None,
) -> PluginPublication:
name: Final = archive.declared_name
if name is None:
raise HTTPException(422, "SKILL.md must declare a skill name")
@ -322,20 +460,119 @@ async def _publish(
"origin": origin,
},
}
identifier: Final = await store.create(
PluginUpload(name, version, archive.declared_description, manifest, plugin, caller.user_id)
return PluginPublication(
name=name,
version=version,
description=archive.declared_description,
manifest=manifest,
content=plugin,
created_by=created_by,
)
if identifier is None:
raise HTTPException(409, "A skill with this name already exists")
return _metadata_response(
request,
PluginArchiveMetadata(
name, True, source, None, identifier, version, archive.declared_description, has_agent_skill_metadata=True
async def _version_metadata(name: str, store: PluginUploadStore, caller: UserAPIKeyAuth) -> PluginArchiveMetadata:
metadata: Final = await store.get_metadata(name)
if metadata is None or metadata.id is None or not (await resolve_skill_visibility(caller)).allows(metadata):
raise HTTPException(404, "Skill not found")
return metadata
def _version_response(request: Request, name: str, version: PluginVersion) -> SkillHubVersion:
return SkillHubVersion(
id=version.id,
version=version.version,
description=version.description,
manifest=version.manifest,
published_at=version.published_at,
published_by=version.published_by,
content_kind="stored_archive" if version.archive_available else "source_reference",
download_url=(
_url(request, f"skill-hub/skills/{quote(name, safe='')}/versions/{quote(version.id, safe='')}/archive")
if version.archive_available
else None
),
)
async def _visible_archive(request: Request, name: str, uploads: PluginUploadStore | None) -> SkillArchive:
@router.get("/{name}/versions", response_model=SkillHubVersions)
async def list_skill_versions(
request: Request,
name: str,
uploads: PluginUploadRepositoryDep,
caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
limit: Annotated[int, Query(ge=1, le=100)] = 20,
offset: Annotated[int, Query(ge=0)] = 0,
) -> SkillHubVersions:
"""Read published metadata history. External references are not retained artifacts."""
store: Final = _require_store(uploads)
metadata: Final = await _version_metadata(name, store, caller)
versions: Final = await store.versions(metadata, limit + 1, offset)
return SkillHubVersions(
versions=tuple(_version_response(request, name, version) for version in versions[:limit]),
next_offset=offset + limit if len(versions) > limit else None,
)
@router.get("/{name}/versions/{version_id}", response_model=SkillHubVersion)
async def get_skill_version(
request: Request,
name: str,
version_id: str,
uploads: PluginUploadRepositoryDep,
caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> SkillHubVersion:
"""Read a retained publication under the current skill's visibility policy."""
store: Final = _require_store(uploads)
metadata: Final = await _version_metadata(name, store, caller)
version: Final = await store.version(metadata, version_id)
if version is None:
raise HTTPException(404, "Skill version not found")
return _version_response(request, name, version)
@router.get(
"/{name}/versions/{version_id}/archive",
response_class=Response,
responses={200: {"content": {"application/zip": {"schema": {"type": "string", "format": "binary"}}}}},
)
async def download_skill_version(
name: str,
version_id: str,
uploads: PluginUploadRepositoryDep,
caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> Response:
"""Download a retained publication as a portable ZIP with SKILL.md at its root."""
store: Final = _require_store(uploads)
metadata: Final = await _version_metadata(name, store, caller)
version: Final = await store.version(metadata, version_id)
if version is None:
raise HTTPException(404, "Skill version not found")
content: Final = await store.version_content(metadata, version_id)
archive: Final = (
await asyncio.to_thread(build_agent_skill_archive, content, plugin_name=name) if content is not None else None
)
if archive is None:
raise HTTPException(404, "This version has no retained Agent Skills archive")
require_download_recorded(
await store.record_download(
replace(metadata, version=version.version),
archive.digest.removeprefix("sha256:"),
plugin_downloader(caller),
)
)
return Response(
archive.content,
media_type="application/zip",
headers={
"Content-Disposition": f'attachment; filename="{name}.zip"',
"Cache-Control": "private, no-store",
"X-Content-Type-Options": "nosniff",
"ETag": f'"{archive.digest}"',
},
)
async def _visible_archive(request: Request, name: str, uploads: PluginUploadStore | None) -> VisibleSkillArchive:
anonymous_allowed: Final = litellm.public_skills_index is True and not marketplace_auth_required()
caller: Final = (
await header_key_caller(request) if marketplace_key_header_present(request) or not anonymous_allowed else None
@ -358,7 +595,7 @@ async def _visible_archive(request: Request, name: str, uploads: PluginUploadSto
)
if archive is None:
raise HTTPException(404, "This skill has no stored Agent Skills archive")
return archive
return VisibleSkillArchive(archive, metadata, plugin_downloader(caller))
@router.get(
@ -370,9 +607,15 @@ async def download_skill(
request: Request, name: str, uploads: PluginUploadRepositoryDep, digest: str | None = None
) -> Response:
"""Download the complete portable skill, with SKILL.md at the ZIP root."""
archive: Final = await _visible_archive(request, name, uploads)
visible: Final = await _visible_archive(request, name, uploads)
archive: Final = visible.archive
if digest is not None and digest != archive.digest:
raise HTTPException(404, "Skill archive version not found")
require_download_recorded(
await _require_store(uploads).record_download(
visible.metadata, archive.digest.removeprefix("sha256:"), visible.downloader
)
)
return Response(
archive.content,
media_type="application/zip",
@ -388,9 +631,9 @@ async def download_skill(
@router.get("/{name}/content", response_class=Response, responses={200: {"content": {"text/markdown": {}}}})
async def skill_content(request: Request, name: str, uploads: PluginUploadRepositoryDep) -> Response:
"""Read the original SKILL.md without downloading supporting files."""
archive: Final = await _visible_archive(request, name, uploads)
visible: Final = await _visible_archive(request, name, uploads)
return Response(
archive.manifest_content,
visible.archive.manifest_content,
media_type="text/markdown",
headers={"Cache-Control": "private, no-store", "X-Content-Type-Options": "nosniff"},
)

172
litellm/proxy/discovery_endpoints/skill_hub_mutations.py Normal file
View file

@ -0,0 +1,172 @@
import hashlib
from dataclasses import dataclass
from datetime import datetime
from typing import Annotated, Final
from fastapi import Depends, HTTPException
from typing_extensions import assert_never
from litellm.models.plugin_upload import (
PluginPublication,
PluginPublicationBase,
PluginPublicationError,
PluginSnapshot,
PluginUploadStore,
PublishedPlugin,
)
from litellm.models.skill_submission import (
SkillSubmission,
SkillSubmissionStore,
SubmissionConflict,
SubmissionLimitExceeded,
SubmissionMissing,
SubmissionReviewed,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import resolve_skill_visibility
from litellm.proxy.auth.auth_exception_handler import DB_UNAVAILABLE_FALLBACK_USER_ID
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.resource_ownership import is_proxy_admin
from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3
from litellm.repositories.plugin_publication import resolve_plugin_version
from litellm.types.proxy.auth.user_roles import LitellmUserRoles
@dataclass(frozen=True, slots=True)
class SkillMutationError:
status: int
message: str
def writer_error(caller: UserAPIKeyAuth) -> SkillMutationError | None:
if is_proxy_admin(caller):
return None
if (
not caller.user_id
or caller.user_id == DB_UNAVAILABLE_FALLBACK_USER_ID
or caller.user_role in (None, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY)
):
return SkillMutationError(403, "Submitting a skill requires a writable user identity.")
return None
def submission_owner(caller: UserAPIKeyAuth) -> str | None | SkillMutationError:
if is_proxy_admin(caller) or caller.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY:
return None
if not caller.user_id or caller.user_id == DB_UNAVAILABLE_FALLBACK_USER_ID:
return SkillMutationError(403, "Reading submissions requires a user identity.")
return caller.user_id
async def archive_writer(caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)]) -> UserAPIKeyAuth:
from litellm.proxy.proxy_server import proxy_logging_obj
error: Final = writer_error(caller)
if error is not None:
raise HTTPException(error.status, error.message)
limiter: Final = proxy_logging_obj.get_proxy_hook("parallel_request_limiter")
if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3):
raise HTTPException(503, "Skills Hub publishing requires the default v3 rate limiter")
key_identity: Final = caller.api_key or caller.token
identity: Final = (
f"user:{caller.user_id}" if caller.user_id else f"key:{key_identity}" if key_identity else "local-admin"
)
result: Final = await limiter.atomic_check_and_increment_by_n(
descriptors=[
{
"key": "skill_hub_publish",
"value": hashlib.sha256(identity.encode()).hexdigest(),
"rate_limit": {"requests_per_unit": 10, "window_size": 60},
}
],
increments=[{"requests": 1}],
parent_otel_span=caller.parent_otel_span,
)
if result["overall_code"] == "OVER_LIMIT":
raise HTTPException(
429, "Skills Hub allows 10 uploads or imports per minute per user", headers={"Retry-After": "60"}
)
return caller
async def write_target(
name: str,
updating: bool,
uploads: PluginUploadStore | None,
caller: UserAPIKeyAuth,
expected_updated_at: datetime | None = None,
) -> PluginSnapshot | SkillMutationError | None:
error: Final = writer_error(caller)
if error is not None:
return error
if uploads is None:
return SkillMutationError(503, "Skills Hub requires a connected database")
existing: Final = await uploads.get_plugin(name)
if updating and (existing is None or not (await resolve_skill_visibility(caller)).allows(existing)):
return SkillMutationError(404, "Skill not found")
if existing is not None and expected_updated_at is not None and existing.updated_at != expected_updated_at:
return SkillMutationError(409, "The published skill changed. Reload it before submitting.")
if not updating and existing is not None:
return SkillMutationError(409, f"A skill named '{name}' already exists. Update the existing skill instead.")
return existing
def publication_error(error: PluginPublicationError) -> SkillMutationError:
return SkillMutationError({"invalid_version": 400, "conflict": 409, "not_found": 404}[error.kind], error.message)
async def write_skill(
publication: PluginPublication,
existing: PluginSnapshot | None,
uploads: PluginUploadStore,
submissions: SkillSubmissionStore | None,
caller: UserAPIKeyAuth,
) -> PublishedPlugin | SkillSubmission | SkillMutationError:
error: Final = writer_error(caller)
if error is not None:
return error
validation: Final = resolve_plugin_version(
publication.version, None if existing is None else existing.version, updating=existing is not None
)
if isinstance(validation, PluginPublicationError):
return publication_error(validation)
expected: Final = None if existing is None else PluginPublicationBase(existing.id, existing.updated_at)
if is_proxy_admin(caller):
published: Final = await uploads.publish(publication, expected)
return publication_error(published) if isinstance(published, PluginPublicationError) else published
if submissions is None or not caller.user_id:
return SkillMutationError(503, "Skills Hub requires a connected database")
submission: Final = await submissions.submit(publication, caller.user_id, expected)
return (
SkillMutationError(429, submission.message) if isinstance(submission, SubmissionLimitExceeded) else submission
)
async def review_skill(
submission_id: str,
submissions: SkillSubmissionStore | None,
caller: UserAPIKeyAuth,
approve: bool,
reason: str | None,
publication_base_url: str | None = None,
) -> SubmissionReviewed | SkillMutationError:
if not is_proxy_admin(caller):
return SkillMutationError(403, "Only proxy admins may review skill submissions.")
if submissions is None:
return SkillMutationError(503, "Skills Hub requires a connected database")
result: Final = await submissions.review(
submission_id, caller.user_id or "proxy-admin", approve, reason, publication_base_url=publication_base_url
)
match result:
case SubmissionMissing():
return SkillMutationError(404, "Skill submission not found")
case SubmissionConflict(message=message):
return SkillMutationError(409, message)
case PluginPublicationError(kind="invalid_version") as error:
return publication_error(error)
case PluginPublicationError(message=message):
return SkillMutationError(409, message)
case SubmissionReviewed():
return result
case _:
assert_never(result)

192
litellm/proxy/discovery_endpoints/skill_hub_submissions.py Normal file
View file

@ -0,0 +1,192 @@
import asyncio
from typing import Annotated, Final
from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response
from litellm.models.plugin_upload import PluginPublication
from litellm.models.skill_submission import SkillSubmission, SkillSubmissionStore, SubmissionStatus
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import resolve_skill_visibility
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.deps.repositories import PluginUploadRepositoryDep, SkillSubmissionRepositoryDep
from litellm.proxy.discovery_endpoints.agent_skills_archive import build_agent_skill_archive
from litellm.proxy.discovery_endpoints.skill_hub_endpoints import (
skill_metadata_response,
skill_proposal_response,
skill_publication_response,
)
from litellm.proxy.discovery_endpoints.skill_hub_mutations import SkillMutationError, review_skill, submission_owner
from litellm.types.proxy.skill_hub import (
SkillHubRejection,
SkillHubReviewResponse,
SkillHubSubmission,
SkillHubSubmissionDetail,
SkillHubSubmissionList,
)
router: Final = APIRouter(prefix="/skill-hub/submissions", tags=["Skills Hub"])
def _store(submissions: SkillSubmissionStore | None) -> SkillSubmissionStore:
if submissions is None:
raise HTTPException(503, "Skills Hub requires a connected database")
return submissions
def _owner(caller: UserAPIKeyAuth) -> str | None:
owner: Final = submission_owner(caller)
if isinstance(owner, SkillMutationError):
raise HTTPException(owner.status, owner.message)
return owner
def _response(request: Request, submission: SkillSubmission) -> SkillHubSubmission:
return SkillHubSubmission(
id=submission.id,
skill_name=submission.plugin_name,
operation=submission.operation,
status=submission.status,
proposal=skill_proposal_response(request, submission),
has_archive=submission.has_archive,
submitted_by=submission.submitted_by,
created_at=submission.created_at,
base_skill_id=submission.base_plugin_id,
base_updated_at=submission.base_updated_at,
reviewed_by=submission.reviewed_by,
reviewed_at=submission.reviewed_at,
rejection_reason=submission.rejection_reason,
)
@router.get("", response_model=SkillHubSubmissionList)
async def list_submissions(
request: Request,
submissions: SkillSubmissionRepositoryDep,
caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
status: SubmissionStatus | None = None,
limit: Annotated[int, Query(ge=1, le=100)] = 20,
offset: Annotated[int, Query(ge=0)] = 0,
) -> SkillHubSubmissionList:
"""Read a page of owned proposals. Proxy admins and admin viewers can read every owner's proposals."""
owner: Final = _owner(caller)
page: Final = await _store(submissions).list(owner, status, limit, offset)
return SkillHubSubmissionList(
submissions=tuple(_response(request, row) for row in page.submissions),
count=len(page.submissions),
has_more=page.has_more,
)
@router.get("/{submission_id}", response_model=SkillHubSubmissionDetail)
async def get_submission(
request: Request,
submission_id: str,
submissions: SkillSubmissionRepositoryDep,
uploads: PluginUploadRepositoryDep,
caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> SkillHubSubmissionDetail:
"""Read an immutable proposal and its currently visible published skill."""
owner: Final = _owner(caller)
submission: Final = await _store(submissions).get(submission_id, owner)
if submission is None:
raise HTTPException(404, "Skill submission not found")
if uploads is None:
raise HTTPException(503, "Skills Hub requires a connected database")
current: Final = await uploads.get_metadata(submission.plugin_name)
visible: Final = current is not None and (owner is None or (await resolve_skill_visibility(caller)).allows(current))
return SkillHubSubmissionDetail(
submission=_response(request, submission),
current_skill=skill_metadata_response(request, current) if current is not None and visible else None,
)
@router.get(
"/{submission_id}/archive",
response_class=Response,
responses={200: {"content": {"application/zip": {"schema": {"type": "string", "format": "binary"}}}}},
)
async def download_submission(
submission_id: str,
submissions: SkillSubmissionRepositoryDep,
caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> Response:
"""Download an owned proposal as a private portable ZIP. Legacy nonportable archives return 404."""
owner: Final = _owner(caller)
store: Final = _store(submissions)
submission: Final = await store.get(submission_id, owner)
if submission is None:
raise HTTPException(404, "Skill submission not found")
content: Final = await store.get_archive(submission_id, owner)
archive: Final = (
await asyncio.to_thread(build_agent_skill_archive, content, plugin_name=submission.plugin_name)
if content is not None
else None
)
if archive is None:
raise HTTPException(404, "This submission has no portable Agent Skills archive")
return Response(
archive.content,
media_type="application/zip",
headers={
"Content-Disposition": f'attachment; filename="{submission.plugin_name}.zip"',
"Cache-Control": "private, no-store",
"X-Content-Type-Options": "nosniff",
"ETag": f'"{archive.digest}"',
},
)
async def _review(
request: Request,
submission_id: str,
submissions: SkillSubmissionStore | None,
caller: UserAPIKeyAuth,
approve: bool,
reason: str | None,
) -> SkillHubReviewResponse:
from litellm.proxy.utils import get_custom_url
result: Final = await review_skill(
submission_id, submissions, caller, approve, reason, get_custom_url(str(request.base_url))
)
if isinstance(result, SkillMutationError):
raise HTTPException(result.status, result.message)
proposal: Final = result.submission.proposal
publication: Final = PluginPublication(
name=proposal.name,
version=proposal.version,
description=proposal.description,
manifest={
**proposal.manifest,
**({"source": dict(result.published_source)} if result.published_source else {}),
},
created_by=result.submission.submitted_by,
)
return SkillHubReviewResponse(
action="approved" if approve else "rejected",
submission=_response(request, result.submission),
skill=None if result.plugin is None else skill_publication_response(request, publication, result.plugin),
)
@router.put("/{submission_id}/approve", response_model=SkillHubReviewResponse)
async def approve_submission(
request: Request,
submission_id: str,
submissions: SkillSubmissionRepositoryDep,
caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> SkillHubReviewResponse:
"""Publish a pending proposal and record its approval atomically. Only proxy admins may review."""
return await _review(request, submission_id, submissions, caller, True, None)
@router.put("/{submission_id}/reject", response_model=SkillHubReviewResponse)
async def reject_submission(
request: Request,
submission_id: str,
submissions: SkillSubmissionRepositoryDep,
caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
body: SkillHubRejection | None = None,
) -> SkillHubReviewResponse:
"""Reject a pending proposal without changing the published skill. Only proxy admins may review."""
return await _review(request, submission_id, submissions, caller, False, None if body is None else body.reason)

186
litellm/proxy/discovery_endpoints/skill_scanning.py Normal file
View file

@ -0,0 +1,186 @@
import asyncio
import io
import json
from datetime import datetime, timedelta, timezone
from typing import Annotated, Final
from uuid import uuid4
import httpx
from fastapi import APIRouter, BackgroundTasks, Body, Depends, HTTPException
from pydantic import ValidationError
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared HTTP factory has untyped optional params
)
from litellm.models.skill_submission import SkillScanInput, SkillSubmissionStore, SubmissionConflict, SubmissionMissing
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.resource_ownership import is_proxy_admin
from litellm.proxy.deps.repositories import SkillSubmissionRepositoryDep
from litellm.secret_managers.main import get_secret_str
from litellm.types.proxy.skill_scanning import (
SkillScan,
SkillScannerConfig,
SkillScanReport,
SkillScanRequest,
SkillScanResponse,
)
router: Final = APIRouter()
_MAX_REPORT_BYTES: Final = 1024 * 1024
def scanner_config() -> SkillScannerConfig | str | None:
from litellm.proxy.proxy_server import general_settings_view
value: Final = general_settings_view().get("skill_scanner")
if value is None:
return None
try:
return SkillScannerConfig.model_validate(value)
except ValidationError:
return "The skill scanner configuration is invalid"
async def scan_report(
client: httpx.AsyncClient, config: SkillScannerConfig, api_key: str, target: SkillScanInput
) -> SkillScanReport | str:
try:
async with client.stream(
"POST",
config.endpoint,
headers={"Authorization": f"Bearer {api_key}", "Accept-Encoding": "identity"},
data={
"source": json.dumps(dict(target.source)),
"use_llm": str(config.use_llm).lower(),
"expires_at": target.scan.expires_at.isoformat(),
},
files=None if target.content is None else {"file": ("skill.zip", target.content, "application/zip")},
timeout=httpx.Timeout(config.timeout_seconds, connect=10),
follow_redirects=False,
) as response:
if response.status_code != 200:
return f"Scanner returned HTTP {response.status_code}. Check the scanner service and authentication."
output: Final = io.BytesIO()
async for chunk in response.aiter_bytes(chunk_size=65536):
if output.tell() + len(chunk) > _MAX_REPORT_BYTES:
return "Scanner report exceeded the 1 MiB limit."
output.write(chunk)
report: Final = SkillScanReport.model_validate_json(output.getvalue())
if config.use_llm and report.scan_mode != "static+llm":
return report.model_copy(
update={
"complete": False,
"limitations": (*report.limitations[:99], "Requested semantic analysis did not run."),
}
)
return report
except (httpx.HTTPError, ValidationError):
return "The scanner was unreachable or returned an invalid report."
async def complete_scan(
submissions: SkillSubmissionStore,
submission_id: str,
target: SkillScanInput,
config: SkillScannerConfig,
api_key: str,
) -> None:
try:
client: Final = get_async_httpx_client(
"skill_scanner", params={"ssl_verify": True, "follow_redirects": False}
).client
report: Final = await asyncio.wait_for(
scan_report(client, config, api_key, target),
timeout=(target.scan.expires_at - datetime.now(timezone.utc)).total_seconds(),
)
completed: Final = target.scan.model_copy(
update={
"status": "completed" if isinstance(report, SkillScanReport) else "failed",
"report": report if isinstance(report, SkillScanReport) else None,
"error": report if isinstance(report, str) else None,
"completed_at": datetime.now(timezone.utc),
}
)
await submissions.finish_scan(submission_id, completed)
except (TimeoutError, asyncio.CancelledError) as error:
await submissions.finish_scan(
submission_id,
target.scan.model_copy(
update={
"status": "failed",
"error": "Scan timed out or was interrupted. Try scanning again.",
"completed_at": datetime.now(timezone.utc),
}
),
)
if isinstance(error, asyncio.CancelledError):
raise
@router.get(
"/skill-hub/submissions/{submission_id}/scan",
tags=["Skills Hub"],
response_model=SkillScanResponse,
)
async def get_skill_scan(
submission_id: str,
submissions: SkillSubmissionRepositoryDep,
caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> SkillScanResponse:
"""Read the latest manual security scan. Only proxy admins may read reports."""
if not is_proxy_admin(caller):
raise HTTPException(403, "Proxy admin access is required to review security scans")
if submissions is None:
raise HTTPException(503, "Database is not connected")
scan: Final = await submissions.get_scan(submission_id)
if isinstance(scan, SubmissionMissing):
raise HTTPException(404, "Skill submission not found")
return SkillScanResponse(
configured=isinstance(scanner_config(), SkillScannerConfig),
scan=None if scan is None else scan.at(datetime.now(timezone.utc)),
)
@router.post(
"/skill-hub/submissions/{submission_id}/scan",
tags=["Skills Hub"],
response_model=SkillScanResponse,
)
async def start_skill_scan(
submission_id: str,
background_tasks: BackgroundTasks,
submissions: SkillSubmissionRepositoryDep,
caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
_request: Annotated[SkillScanRequest | None, Body()] = None,
) -> SkillScanResponse:
"""Start an optional scan of a pending submission. Scans never change approval decisions."""
if not is_proxy_admin(caller):
raise HTTPException(403, "Proxy admin access is required to scan skills")
if submissions is None:
raise HTTPException(503, "Database is not connected")
config: Final = scanner_config()
if config is None:
raise HTTPException(503, "Security scanning is not configured")
if isinstance(config, str):
raise HTTPException(503, config)
api_key: Final = get_secret_str(config.api_key_env)
if not api_key:
raise HTTPException(503, "The scanner authentication secret is not configured")
now: Final = datetime.now(timezone.utc)
target: Final = await submissions.start_scan(
submission_id,
SkillScan(
id=str(uuid4()),
status="running",
requested_by=caller.user_id or "proxy-admin",
started_at=now,
expires_at=now + timedelta(seconds=config.timeout_seconds),
),
)
if isinstance(target, SubmissionMissing):
raise HTTPException(404, "Skill submission not found")
if isinstance(target, SubmissionConflict):
raise HTTPException(409, target.message)
background_tasks.add_task(complete_scan, submissions, submission_id, target, config, api_key)
return SkillScanResponse(configured=True, scan=target.scan)

45
litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py
View file

@ -4,7 +4,8 @@
#
# +-------------------------------------------------------------+
import os
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from itertools import chain, dropwhile, takewhile
from typing import TYPE_CHECKING, Final, Literal, Optional
from fastapi import HTTPException
@ -14,6 +15,10 @@ from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_scan_only_tool_results_for_guardrail,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
@ -218,15 +223,14 @@ class ZscalerAIGuard(CustomGuardrail):
zscaler_ai_guard_result = None
direction: Final = "OUT" if input_type == "response" else "IN"
verbose_proxy_logger.debug("direction: %s", direction)
# Concatenate all texts and send to Zscaler AI Guard
if texts:
concatenated_text: Final = " ".join(texts)
content: Final = self._content_to_scan(inputs, input_type)
if content:
zscaler_ai_guard_result = await self.make_zscaler_ai_guard_api_call(
zscaler_ai_guard_url=self.zscaler_ai_guard_url,
api_key=self.api_key,
policy_id=policy_id,
direction=direction,
content=concatenated_text,
content=content,
user=user_for_payload,
**kwargs,
)
@ -244,6 +248,37 @@ class ZscalerAIGuard(CustomGuardrail):
verbose_proxy_logger.debug("ZscalerAIGuard: Successfully applied guardrail.")
return inputs
def _content_to_scan(
self,
inputs: GenericGuardrailAPIInputs,
input_type: Literal["request", "response"],
) -> str | None:
structured_messages: Final = inputs.get("structured_messages") or []
scans_last_user_turn: Final = (
input_type == "request"
and not effective_scan_only_tool_results_for_guardrail(self)
and any("content" in message for message in structured_messages)
)
texts: Final = (
self._last_user_turn_texts(structured_messages) if scans_last_user_turn else inputs.get("texts", [])
)
return " ".join(texts) if texts else None
@staticmethod
def _last_user_turn_texts(messages: Sequence[Mapping[str, object]]) -> tuple[str, ...]:
from_last_user_message: Final = dropwhile(lambda message: message.get("role") != "user", reversed(messages))
last_user_turn: Final = tuple(takewhile(lambda message: message.get("role") == "user", from_last_user_message))
texts: Final = tuple(chain.from_iterable(map(ZscalerAIGuard._content_texts, reversed(last_user_turn))))
return texts if any(text.strip() for text in texts) else ()
@staticmethod
def _content_texts(message: Mapping[str, object]) -> tuple[str, ...]:
content: Final = message.get("content")
if isinstance(content, str):
return (content,)
parts: Final = (as_str_mapping(part) for part in as_object_list(content) or ())
return tuple(text for part in parts if part is not None and isinstance(text := part.get("text"), str))
def extract_blocking_info(self, response):
"""
Extracts transaction ID and blocking detector details from a response.

91
litellm/proxy/management_endpoints/access_group_endpoints.py
View file

@ -347,47 +347,45 @@ async def _cache_access_group_record(record: _AccessGroupRecord) -> None:
# ---------------------------------------------------------------------------
async def _sync_add_access_group_to_teams(tx: _AccessGroupTx, team_ids: list[str], access_group_id: str) -> None:
"""Add access_group_id to each team's access_group_ids (idempotent)."""
for team_id in team_ids:
async def _sync_access_group_on_teams(
tx: _AccessGroupTx, add: Sequence[str], remove: Sequence[str], access_group_id: str
) -> None:
"""Rows are locked in sorted team_id order, the order every spend writer locks team rows in, so this
cannot deadlock with them."""
to_add: Final = frozenset(add)
for team_id in sorted(to_add | frozenset(remove)):
team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id})
if team is not None and access_group_id not in (team.access_group_ids or []):
if team is None:
continue
current = list(team.access_group_ids or [])
if team_id in to_add and access_group_id not in current:
await tx.litellm_teamtable.update(
where={"team_id": team_id}, data={"access_group_ids": [*current, access_group_id]}
)
elif team_id not in to_add and access_group_id in current:
await tx.litellm_teamtable.update(
where={"team_id": team_id},
data={"access_group_ids": list(team.access_group_ids or []) + [access_group_id]},
data={"access_group_ids": list(_without_access_group(current, access_group_id))},
)
async def _sync_remove_access_group_from_teams(tx: _AccessGroupTx, team_ids: list[str], access_group_id: str) -> None:
"""Remove access_group_id from each team's access_group_ids (idempotent)."""
for team_id in team_ids:
team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id})
if team is not None and access_group_id in (team.access_group_ids or []):
await tx.litellm_teamtable.update(
where={"team_id": team_id},
data={"access_group_ids": [ag for ag in (team.access_group_ids or ()) if ag != access_group_id]},
)
async def _sync_add_access_group_to_keys(tx: _AccessGroupTx, key_tokens: list[str], access_group_id: str) -> None:
"""Add access_group_id to each key's access_group_ids (idempotent)."""
for token in key_tokens:
async def _sync_access_group_on_keys(
tx: _AccessGroupTx, add: Sequence[str], remove: Sequence[str], access_group_id: str
) -> None:
"""Rows are locked in sorted token order, the order every spend writer locks key rows in."""
to_add: Final = frozenset(add)
for token in sorted(to_add | frozenset(remove)):
key = await tx.litellm_verificationtoken.find_unique(where={"token": token})
if key is not None and access_group_id not in (key.access_group_ids or []):
if key is None:
continue
current = list(key.access_group_ids or [])
if token in to_add and access_group_id not in current:
await tx.litellm_verificationtoken.update(
where={"token": token},
data={"access_group_ids": list(key.access_group_ids or []) + [access_group_id]},
where={"token": token}, data={"access_group_ids": [*current, access_group_id]}
)
async def _sync_remove_access_group_from_keys(tx: _AccessGroupTx, key_tokens: list[str], access_group_id: str) -> None:
"""Remove access_group_id from each key's access_group_ids (idempotent)."""
for token in key_tokens:
key = await tx.litellm_verificationtoken.find_unique(where={"token": token})
if key is not None and access_group_id in (key.access_group_ids or []):
elif token not in to_add and access_group_id in current:
await tx.litellm_verificationtoken.update(
where={"token": token},
data={"access_group_ids": [ag for ag in (key.access_group_ids or ()) if ag != access_group_id]},
where={"token": token}, data={"access_group_ids": list(_without_access_group(current, access_group_id))}
)
@ -582,8 +580,8 @@ async def create_access_group(
)
# Sync team and key tables to reference the new access group
await _sync_add_access_group_to_teams(tx, data.assigned_team_ids or [], record.access_group_id)
await _sync_add_access_group_to_keys(tx, data.assigned_key_ids or [], record.access_group_id)
await _sync_access_group_on_teams(tx, data.assigned_team_ids or [], (), record.access_group_id)
await _sync_access_group_on_keys(tx, data.assigned_key_ids or [], (), record.access_group_id)
except HTTPException:
raise
except Exception as e:
@ -718,10 +716,8 @@ async def update_access_group(
data=update_data,
)
await _sync_add_access_group_to_teams(tx, teams_to_add, access_group_id)
await _sync_remove_access_group_from_teams(tx, teams_to_remove, access_group_id)
await _sync_add_access_group_to_keys(tx, keys_to_add, access_group_id)
await _sync_remove_access_group_from_keys(tx, keys_to_remove, access_group_id)
await _sync_access_group_on_teams(tx, teams_to_add, teams_to_remove, access_group_id)
await _sync_access_group_on_keys(tx, keys_to_add, keys_to_remove, access_group_id)
except HTTPException:
raise
except Exception as e:
@ -790,25 +786,8 @@ async def delete_access_group(
)
affected_key_tokens = list(all_affected_key_tokens)
# Update teams returned by find_many directly — we already have their data.
for team in teams_with_group:
await tx.litellm_teamtable.update(
where={"team_id": team.team_id},
data={"access_group_ids": [ag for ag in (team.access_group_ids or ()) if ag != access_group_id]},
)
# Use _sync_remove only for out-of-sync teams not found by the hasSome query.
out_of_sync_team_ids: Final = set(existing.assigned_team_ids or []) - {t.team_id for t in teams_with_group}
await _sync_remove_access_group_from_teams(tx, list(out_of_sync_team_ids), access_group_id)
# Update keys returned by find_many directly — we already have their data.
for key in keys_with_group:
await tx.litellm_verificationtoken.update(
where={"token": key.token},
data={"access_group_ids": [ag for ag in (key.access_group_ids or ()) if ag != access_group_id]},
)
# Use _sync_remove only for out-of-sync keys not found by the hasSome query.
out_of_sync_key_tokens: Final = set(existing.assigned_key_ids or []) - {k.token for k in keys_with_group}
await _sync_remove_access_group_from_keys(tx, list(out_of_sync_key_tokens), access_group_id)
await _sync_access_group_on_teams(tx, (), affected_team_ids, access_group_id)
await _sync_access_group_on_keys(tx, (), affected_key_tokens, access_group_id)
detached_agent_ids: Final = await _detach_access_group_from_agents(tx, access_group_id)

26
litellm/proxy/management_endpoints/common_daily_activity.py
View file

@ -287,12 +287,12 @@ def compute_tag_metadata_totals(records: Sequence[DailySpendRecord]) -> SpendMet
def _entity_metadata(
entity_metadata_field: Mapping[str, dict[str, object]] | None,
entity_metadata_field: Mapping[str, Mapping[str, object]] | None,
entity_id: str,
) -> dict[str, object]:
"""The metadata payload for one entity breakdown bucket, empty when the caller passed none."""
stored: Final = entity_metadata_field.get(entity_id) if entity_metadata_field else None
return stored if stored is not None else {}
return dict(stored) if stored is not None else {}
def provider_bucket_metadata(provider: str, base: Mapping[str, object]) -> Mapping[str, object]:
@ -1043,7 +1043,7 @@ def _fold_entity_rollups_sync(
results: Sequence[DailySpendData],
entity_rows: Sequence[EntityRollupRow],
api_key_metadata: Mapping[str, KeyMetadataRow],
entity_metadata_field: Mapping[str, dict[str, object]] | None, # mutable-ok: shared field shape
entity_metadata_field: Mapping[str, Mapping[str, object]] | None,
) -> None:
"""Write breakdown.entities onto the already-built per-day results."""
by_date: Final = {day.date.strftime("%Y-%m-%d"): day for day in results}
@ -1117,6 +1117,19 @@ async def _guardrail_spend_or_empty(
return ()
async def _user_entity_metadata(
repository: DailyActivityRepository, entity_rows: Sequence[EntityRollupRow]
) -> Mapping[str, Mapping[str, object]]:
user_ids: Final = frozenset(row.entity_id for row in entity_rows if row.entity_id)
users: Final = await repository.user_metadata(user_ids)
return MappingProxyType(
{
user_id: MappingProxyType({"user_email": user.user_email, "user_alias": user.user_alias})
for user_id, user in users.items()
}
)
async def get_daily_activity_aggregated(
repository: DailyActivityRepository,
scope: DailyActivityScope,
@ -1161,12 +1174,17 @@ async def get_daily_activity_aggregated(
if entity_api_keys
else MappingProxyType({})
)
user_entity_metadata: Final[Mapping[str, Mapping[str, object]]] = (
await _user_entity_metadata(repository, entity_records)
if scope.table is DailyActivityTable.USER
else MappingProxyType({})
)
await asyncio.to_thread(
_fold_entity_rollups_sync,
results=aggregated["results"],
entity_rows=entity_records,
api_key_metadata=entity_key_metadata,
entity_metadata_field=entity_metadata_field,
entity_metadata_field=MappingProxyType({**user_entity_metadata, **(entity_metadata_field or {})}),
)
results: Final = (
list(await asyncio.to_thread(_with_guardrail_breakdown, aggregated["results"], guardrail_rows))

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