scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-20 00:11:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 75

Merge branches 'main' and 'litellm_fix_integration_conftest_import' of github.com:BerriAI/litellm into litellm_fix_integration_conftest_import

Browse source 
# Conflicts:
#	tests/integration/_support/client.py
#	tests/integration/management/test_partial_update_sequences.py
This commit is contained in:
Yuneng Jiang 2026-09-17 10:33:31 -07:00
commit 7013204ded
No known key found for this signature in database
712 changed files with 42070 additions and 19693 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

32
.circleci/config.yml
View file

@ -257,7 +257,7 @@ commands:
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
- v3-integration-uv-cache-{{ checksum "uv.lock" }}
- run:
name: Install Dependencies
command: |
@ -266,7 +266,7 @@ commands:
- save_cache:
paths:
- ~/.cache/uv
key: v1-uv-cache-{{ checksum "uv.lock" }}
key: v3-integration-uv-cache-{{ checksum "uv.lock" }}
jobs:
# Add Windows testing job
@ -2955,6 +2955,32 @@ jobs:
working_directory: ~/project
steps:
- setup_litellm_test_deps
- when:
condition:
equal: [browser, << parameters.suite >>]
steps:
- install_node
- restore_cache:
keys:
- integration-ui-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
- run:
name: Install locked browser dependencies
command: |
cd ui/litellm-dashboard
npm ci
cd ../../tests/e2e/ui
npm ci
sudo env PATH="$PATH" DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=l \
timeout --signal=TERM --kill-after=20s 6m node node_modules/@playwright/test/cli.js install-deps chromium
timeout --signal=TERM --kill-after=20s 3m node node_modules/@playwright/test/cli.js install chromium
- save_cache:
key: integration-ui-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
paths:
- ~/.npm
- ~/.cache/ms-playwright
- run:
name: Build the candidate dashboard
command: cd ui/litellm-dashboard && NEXT_TELEMETRY_DISABLED=1 npm run build
- start_postgres:
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
- start_redis
@ -2983,7 +3009,7 @@ workflows:
name: integration-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, providers]
suite: [management, accounting, database, providers, extensions, browser]
filters:
branches:
only:

13
.circleci/scripts/classify_changes.sh
View file

@ -1,12 +1,14 @@
#!/usr/bin/env bash
set -uo pipefail
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness>}"
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness|cost-map-only>}"
has_client=false
has_backend=false
has_ci=false
has_provider_harness=false
has_cost_map=false
outside_cost_map_set=false
while IFS= read -r file || [ -n "$file" ]; do
[ -n "$file" ] || continue
case "$file" in
@ -20,9 +22,18 @@ while IFS= read -r file || [ -n "$file" ]; do
.github/* | .circleci/*) has_ci=true; has_backend=true ;;
*) has_backend=true ;;
esac
case "$file" in
model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json | model_prices_and_context_window.schema.json)
has_cost_map=true ;;
tests/test_litellm/* | tests/proxy_unit_tests/*) : ;;
*) outside_cost_map_set=true ;;
esac
done
case "$category" in
cost-map-only)
{ [ "$has_cost_map" = true ] && [ "$outside_cost_map_set" = false ]; } && echo run || echo skip
;;
provider-harness)
[ "$has_provider_harness" = true ] && echo run || echo skip
;;

29
.circleci/scripts/run_integration.sh
View file

@ -1,6 +1,11 @@
#!/usr/bin/env bash
set -euo pipefail
if [ "${GITHUB_ACTIONS:-}" = true ]; then
echo "Integration contracts are owned by CircleCI" >&2
exit 1
fi
suite="${1:?integration suite required}"
results="test-results/integration-${suite}"
mkdir -p "$results"
@ -65,7 +70,13 @@ export INTEGRATION_PROXY_URL=http://127.0.0.1:4000
export INTEGRATION_PEER_URL=""
export INTEGRATION_UPSTREAM_URL=http://127.0.0.1:8190
export INTEGRATION_MASTER_KEY="$LITELLM_MASTER_KEY"
export INTEGRATION_SEED="$((16#$(git rev-parse --short=8 HEAD)))"
export LITELLM_UI_PATH="$PWD/litellm/proxy/_experimental/out"
if [ "$suite" = browser ]; then
export LITELLM_UI_PATH="$PWD/ui/litellm-dashboard/out"
test -f "$LITELLM_UI_PATH/index.html"
fi
export INTEGRATION_SEED="$(.venv/bin/python -c 'import hashlib,os; print(int(hashlib.sha256((os.environ.get("CIRCLE_SHA1", "local") + os.environ.get("CIRCLE_WORKFLOW_ID", "local")).encode()).hexdigest()[:8],16))')"
export INTEGRATION_ORDER_SEED="$INTEGRATION_SEED"
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma > "$results/prisma-generate.log" 2>&1
@ -102,7 +113,7 @@ start_proxy() {
local log_name="$2"
setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" \
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \
LITELLM_MODE=PRODUCTION LITELLM_LOCAL_MODEL_COST_MAP=True STORE_MODEL_IN_DB=True \
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
.venv/bin/python -m integration._support.proxy --config tests/integration/proxy_config.yaml \
@ -131,6 +142,19 @@ if [ "$suite" = providers ]; then
--junitxml="$results/replay-controls.xml"
fi
if [ "$suite" = browser ]; then
export E2E_UI_BASE_URL="$INTEGRATION_PROXY_URL" E2E_UI_ARTIFACT_DIR="$PWD/$results"
export INTEGRATION_PYTHON="$PWD/.venv/bin/python"
timeout --signal=TERM --kill-after=20s 3m env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
INTEGRATION_RUN_ID="$integration_identity" DATABASE_URL="$DATABASE_URL" \
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" INTEGRATION_PYTHON="$INTEGRATION_PYTHON" \
E2E_UI_BASE_URL="$E2E_UI_BASE_URL" E2E_UI_ARTIFACT_DIR="$E2E_UI_ARTIFACT_DIR" \
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" CI=true \
node tests/e2e/ui/node_modules/@playwright/test/cli.js test --config tests/e2e/ui/integration.config.ts
.venv/bin/python .circleci/scripts/verify_integration_browser.py "$results/browser-results.json"
exit 0
fi
timeout --signal=TERM --kill-after=20s 11m env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
INTEGRATION_RUN_ID="$integration_identity" \
DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
@ -138,5 +162,6 @@ timeout --signal=TERM --kill-after=20s 11m env -i PATH="$PATH" HOME="$HOME" PYTH
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
INTEGRATION_MASTER_KEY="$INTEGRATION_MASTER_KEY" LITELLM_MODE=PRODUCTION \
INTEGRATION_SEED="$INTEGRATION_SEED" \
INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \
LITELLM_LOCAL_MODEL_COST_MAP=True AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
.venv/bin/python tests/integration/run.py "$suite" --results "$results"

60
.circleci/scripts/verify_integration_browser.py Normal file
View file

@ -0,0 +1,60 @@
import json
import sys
from pathlib import Path
from typing import Final
from pydantic import TypeAdapter
from typing_extensions import NotRequired, ReadOnly, TypedDict
class BrowserAttempt(TypedDict):
status: ReadOnly[str]
retry: ReadOnly[int]
class BrowserTest(TypedDict):
results: ReadOnly[list[BrowserAttempt]]
class BrowserSpec(TypedDict):
file: ReadOnly[str]
title: ReadOnly[str]
tests: ReadOnly[list[BrowserTest]]
class BrowserSuite(TypedDict):
specs: NotRequired[ReadOnly[list[BrowserSpec]]]
suites: NotRequired[ReadOnly[list["BrowserSuite"]]]
def main() -> None:
result: Final = json.loads(Path(sys.argv[1]).read_text())
assert not result.get("errors"), result.get("errors")
expected: Final = json.loads(
(Path(__file__).resolve().parents[2] / "tests/integration/contracts.json").read_text()
)["browser"]
assert expected and result["stats"]["expected"] == len(expected)
assert all(result["stats"][name] == 0 for name in ("unexpected", "flaky", "skipped"))
def cases(suite: BrowserSuite) -> tuple[BrowserSpec, ...]:
return tuple(suite.get("specs", ())) + tuple(spec for child in suite.get("suites", ()) for spec in cases(child))
suites: Final = TypeAdapter(list[BrowserSuite]).validate_python(result["suites"], strict=True)
specs: Final = tuple(spec for suite in suites for spec in cases(suite))
repository: Final = Path(__file__).resolve().parents[2]
report_root: Final = Path(result["config"]["rootDir"])
assert report_root.is_absolute(), "Playwright rootDir must be explicit"
observed: Final = tuple(
str((report_root / spec["file"]).resolve().relative_to(repository)) + "::" + spec["title"] for spec in specs
)
assert sorted(observed) == sorted(expected)
for spec in specs:
tests: Final = spec["tests"]
assert len(tests) == 1 and len(tests[0]["results"]) == 1
assert tests[0]["results"][0]["status"] == "passed" and tests[0]["results"][0]["retry"] == 0
sys.stdout.write("One canonical browser contract passed once without skips or retries\n")
if __name__ == "__main__":
main()

4
.github/CODEOWNERS vendored
View file

@ -4,7 +4,7 @@
/ui/nginx.conf
/ui/litellm-dashboard/src/lib/http/schema.d.ts
/ui/litellm-dashboard/tsconfig.tsbuildinfo
/model_prices_and_context_window.json @mateo-berri
/litellm/model_prices_and_context_window_backup.json @mateo-berri
/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri
/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri
/litellm-proxy-extras/litellm_proxy_extras/migrations/ @yuneng-berri @ryan-crabbe-berri
/.github/CODEOWNERS @yuneng-berri

39
.github/scripts/assert_ci_coverage.py vendored
View file

@ -505,6 +505,7 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
return frozenset(), ()
entries: Final = json.loads(manifest.read_text())
paths: Final = frozenset(node.split("::", 1)[0] for node in entries["tests"])
browser_paths: Final = frozenset(node.split("::", 1)[0] for node in entries.get("browser", {}))
circle_path: Final = repo_root / ".circleci/config.yml"
circle: Final = yaml.safe_load(circle_path.read_text()) if circle_path.exists() else {}
steps: Final = circle.get("jobs", {}).get("integration_contracts", {}).get("steps", ())
@ -523,7 +524,7 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
.get("suite", (job["integration_contracts"].get("suite"),))
if isinstance(suite, str)
)
required: Final = frozenset(
required: Final = (frozenset({"browser"}) if browser_paths else frozenset()) | frozenset(
group
for group, folders in entries["groups"].items()
if any(any(path.startswith(f"tests/integration/{folder}/") for folder in folders) for path in paths)
@ -551,6 +552,40 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
for path in paths
if not (repo_root / path).is_file()
)
browser_commands: Final = tuple(
scalar.value
for path in (repo_root / ".github/workflows").glob("*.y*ml")
for scalar in _scalars(yaml.safe_load(path.read_text()), path.name)
if scalar.key in {"run", "command"}
)
browser_findings: Final = tuple(
Finding(path, "browser integration contract is explicitly selected by GitHub Actions")
for path in browser_paths
if any(
path in command
or pathlib.Path(path).name in command
or "integrationCritical" in command
or "integration.config.ts" in command
or ("run_integration.sh" in command and "browser" in command)
for command in browser_commands
)
) + tuple(
Finding(path, "canonical browser integration file is missing")
for path in browser_paths
if not (repo_root / path).is_file()
)
default_browser: Final = repo_root / "tests/e2e/ui/playwright.config.ts"
exclusion_findings: Final = (
(
Finding(
str(default_browser.relative_to(repo_root)),
"default Playwright selection must exclude integrationCritical",
),
)
if browser_paths
and (not default_browser.exists() or "**/integrationCritical/**" not in default_browser.read_text())
else ()
)
group_findings: Final = tuple(
Finding(group, "canonical integration group is not scheduled by CircleCI")
for group in sorted(required - scheduled)
@ -559,7 +594,7 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
return frozenset(), findings + (
Finding(str(manifest.relative_to(repo_root)), "dedicated CircleCI runner is missing"),
)
return paths, findings + group_findings
return paths | browser_paths, findings + group_findings + browser_findings + exclusion_findings
def main() -> int:

465
.github/scripts/auto_merge_price_sync.py vendored Normal file
View file

@ -0,0 +1,465 @@
"""Auto-merge the provider-info-sync bot's cost-map pull requests.
Evaluates every gate (author allowlist, cost-map-only diff, required and
non-required checks, Greptile confidence, Bugbot review, human reviews) and
merges with a merge commit when all of them hold. Every hold reason is
logged; the process exits 0 on hold and 1 only on API or programming errors.
``DRY_RUN=1`` prints the verdict without calling the merge endpoint.
"""
from __future__ import annotations
import json
import os
import re
import subprocess
import sys
import time
import urllib.error
import urllib.request
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Final
REPO_ROOT: Final = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
CLASSIFY_SCRIPT: Final = os.path.join(REPO_ROOT, ".circleci", "scripts", "classify_changes.sh")
API_ROOT: Final = "https://api.github.com"
CHANGED_FILE_CEILING: Final = 3000
OK_CHECK_CONCLUSIONS: Final = frozenset({"success", "skipped", "neutral"})
GREPTILE_LOGIN: Final = "greptile-apps[bot]"
BUGBOT_LOGIN: Final = "cursor[bot]"
GREPTILE_SCORE_RE: Final = re.compile(r"Confidence Score:\s*(\d)/5")
BUGBOT_REVIEW_MARKER: Final = "<!-- BUGBOT_REVIEW -->"
BUGBOT_STALE_MARKER: Final = "<!-- BUGBOT_REVIEW_STALE -->"
BUGBOT_CLEAN: Final = "found no new issues"
@dataclass(frozen=True, slots=True)
class PullRequest:
number: int
title: str
author_login: str
state: str
draft: bool
mergeable: bool | None
mergeable_state: str
head_sha: str
@dataclass(frozen=True, slots=True)
class CheckRun:
name: str
status: str
conclusion: str | None
@dataclass(frozen=True, slots=True)
class CommitStatus:
context: str
state: str
@dataclass(frozen=True, slots=True)
class IssueComment:
author_login: str
body: str
updated_at: datetime
@dataclass(frozen=True, slots=True)
class Review:
author_login: str
state: str
body: str
commit_id: str
submitted_at: datetime
@dataclass(frozen=True, slots=True)
class Verdict:
merge: bool
reasons: tuple[str, ...]
@dataclass(frozen=True, slots=True)
class EvaluationInputs:
pr: PullRequest
changed_files: tuple[str, ...]
required_contexts: frozenset[str]
check_runs: tuple[CheckRun, ...]
statuses: tuple[CommitStatus, ...]
comments: tuple[IssueComment, ...]
reviews: tuple[Review, ...]
head_commit_date: datetime
self_check_name: str
author_allowlist: frozenset[str]
def _is_bot_login(login: str) -> bool:
return login.lower().endswith("[bot]")
def _classify(changed_files: Sequence[str]) -> str:
result: Final = subprocess.run(
["bash", CLASSIFY_SCRIPT, "cost-map-only"],
input="\n".join(changed_files),
capture_output=True,
text=True,
check=False,
)
if result.returncode != 0:
return "error"
return result.stdout.strip()
def evaluate(
inputs: EvaluationInputs,
*,
classify: Callable[[Sequence[str]], str] = _classify,
) -> Verdict:
pr: Final = inputs.pr
reasons: list[str] = []
if pr.author_login.lower() not in {login.lower() for login in inputs.author_allowlist}:
reasons.append(f"author {pr.author_login!r} not in allowlist")
if pr.state != "open":
reasons.append("pr not open")
if pr.draft:
reasons.append("pr is a draft")
if pr.mergeable is None:
reasons.append("mergeability unknown")
elif not pr.mergeable:
reasons.append("pr not mergeable")
if pr.mergeable_state == "dirty":
reasons.append("pr has merge conflicts")
if len(inputs.changed_files) > CHANGED_FILE_CEILING:
reasons.append(f"changed file count {len(inputs.changed_files)} over {CHANGED_FILE_CEILING} ceiling")
else:
decision: Final = classify(inputs.changed_files)
if decision != "run":
reasons.append("changed files outside the cost-map-only set")
green_runs: Final = frozenset(run.name for run in inputs.check_runs if run.conclusion in OK_CHECK_CONCLUSIONS)
green_statuses: Final = frozenset(status.context for status in inputs.statuses if status.state == "success")
for context in sorted(inputs.required_contexts):
if context not in green_runs and context not in green_statuses:
reasons.append(f"required check {context!r} not green")
for run in inputs.check_runs:
if run.name == inputs.self_check_name:
continue
if run.status != "completed" or run.conclusion not in OK_CHECK_CONCLUSIONS:
reasons.append(f"check run {run.name!r} is {run.status}/{run.conclusion}")
for status in inputs.statuses:
if status.state != "success":
reasons.append(f"commit status {status.context!r} is {status.state}")
greptile: Final = tuple(
comment
for comment in inputs.comments
if comment.author_login == GREPTILE_LOGIN and GREPTILE_SCORE_RE.search(comment.body)
)
if not greptile:
reasons.append("greptile score not available")
else:
latest: Final = max(greptile, key=lambda comment: comment.updated_at)
match: Final = GREPTILE_SCORE_RE.search(latest.body)
score: Final = int(match.group(1)) if match else 0
if latest.updated_at < inputs.head_commit_date:
reasons.append("greptile score older than head commit")
elif score != 5:
reasons.append(f"greptile score {score}/5 below 5")
bugbot: Final = tuple(
review
for review in inputs.reviews
if review.author_login == BUGBOT_LOGIN
and BUGBOT_REVIEW_MARKER in review.body
and BUGBOT_STALE_MARKER not in review.body
and review.commit_id == pr.head_sha
)
if not bugbot:
reasons.append("bugbot review not available")
else:
latest_review: Final = max(bugbot, key=lambda review: review.submitted_at)
if BUGBOT_CLEAN not in latest_review.body:
reasons.append("bugbot reported issues")
latest_state_by_reviewer: Final[dict[str, str]] = {}
for review in sorted(inputs.reviews, key=lambda review: review.submitted_at):
if _is_bot_login(review.author_login):
continue
latest_state_by_reviewer[review.author_login] = review.state
for reviewer, state in latest_state_by_reviewer.items():
if state == "CHANGES_REQUESTED":
reasons.append(f"changes requested by {reviewer}")
return Verdict(merge=not reasons, reasons=tuple(reasons))
def _request(token: str, method: str, path: str, body: Mapping[str, object] | None = None) -> object:
url: Final = path if path.startswith("http") else f"{API_ROOT}{path}"
data: Final = None if body is None else json.dumps(body).encode("utf-8")
request: Final = urllib.request.Request(
url,
data=data,
method=method,
headers={
"Accept": "application/vnd.github+json",
"Authorization": f"Bearer {token}",
"X-GitHub-Api-Version": "2022-11-28",
},
)
with urllib.request.urlopen(request) as response:
return json.loads(response.read().decode("utf-8"))
def _request_allow_fail(
token: str, method: str, path: str, body: Mapping[str, object] | None = None
) -> tuple[int, object | None]:
url: Final = path if path.startswith("http") else f"{API_ROOT}{path}"
data: Final = None if body is None else json.dumps(body).encode("utf-8")
request: Final = urllib.request.Request(
url,
data=data,
method=method,
headers={
"Accept": "application/vnd.github+json",
"Authorization": f"Bearer {token}",
"X-GitHub-Api-Version": "2022-11-28",
},
)
try:
with urllib.request.urlopen(request) as response:
return response.status, json.loads(response.read().decode("utf-8"))
except urllib.error.HTTPError as exc:
return exc.code, None
def _items(payload: object, key: str | None = None) -> tuple[object, ...]:
source: Final = payload.get(key) if key and isinstance(payload, Mapping) else payload
if not isinstance(source, list):
return ()
return tuple(source)
def _paginate(token: str, path: str, key: str | None = None) -> list[object]:
separator: Final = "&" if "?" in path else "?"
results: list[object] = []
for page in range(1, 10_000):
batch: Final = _items(_request(token, "GET", f"{path}{separator}per_page=100&page={page}"), key)
results.extend(batch)
if len(batch) < 100:
return results
return results
def _text(value: object) -> str:
return value if isinstance(value, str) else ""
def _int(value: object) -> int:
return value if isinstance(value, int) else 0
def _bool(value: object) -> bool:
return value is True
def _nested(value: object, *keys: str) -> object:
current: object = value
for key in keys:
if not isinstance(current, Mapping):
return None
current = current.get(key)
return current
def _parse_time(value: object) -> datetime:
text: Final = _text(value)
if not text:
return datetime.min.replace(tzinfo=timezone.utc)
return datetime.fromisoformat(text.replace("Z", "+00:00"))
def _load_pr(token: str, repo: str, number: int) -> PullRequest:
data: Final = _request(token, "GET", f"/repos/{repo}/pulls/{number}")
if not isinstance(data, Mapping):
raise RuntimeError(f"unexpected pull payload for #{number}")
return PullRequest(
number=number,
title=_text(data.get("title")),
author_login=_text(_nested(data, "user", "login")),
state=_text(data.get("state")),
draft=_bool(data.get("draft")),
mergeable=data.get("mergeable") if isinstance(data.get("mergeable"), bool) else None,
mergeable_state=_text(data.get("mergeable_state")),
head_sha=_text(_nested(data, "head", "sha")),
)
def _list_candidate_prs(token: str, repo: str, base: str, allowlist: frozenset[str]) -> list[int]:
candidates: Final = _paginate(token, f"/repos/{repo}/pulls?state=open&base={base}")
return [
_int(item.get("number"))
for item in candidates
if isinstance(item, Mapping) and _text(_nested(item, "user", "login")).lower() in allowlist
]
def _changed_files(token: str, repo: str, number: int) -> tuple[str, ...]:
files: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/files")
return tuple(_text(item.get("filename")) for item in files if isinstance(item, Mapping))
def _required_contexts(token: str, repo: str, base: str) -> frozenset[str]:
payload: Final = _request(token, "GET", f"/repos/{repo}/rules/branches/{base}")
contexts: set[str] = set()
for rule in _items(payload):
if not isinstance(rule, Mapping) or rule.get("type") != "required_status_checks":
continue
checks: Final = _nested(rule, "parameters", "required_status_checks")
for check in _items(checks):
if isinstance(check, Mapping):
context: Final = _text(check.get("context"))
if context:
contexts.add(context)
return frozenset(contexts)
def _check_runs(token: str, repo: str, sha: str) -> tuple[CheckRun, ...]:
runs: Final = _paginate(token, f"/repos/{repo}/commits/{sha}/check-runs", key="check_runs")
return tuple(
CheckRun(
name=_text(item.get("name")),
status=_text(item.get("status")),
conclusion=item.get("conclusion") if isinstance(item.get("conclusion"), str) else None,
)
for item in runs
if isinstance(item, Mapping)
)
def _statuses(token: str, repo: str, sha: str) -> tuple[CommitStatus, ...]:
payload: Final = _request(token, "GET", f"/repos/{repo}/commits/{sha}/status")
return tuple(
CommitStatus(context=_text(item.get("context")), state=_text(item.get("state")))
for item in _items(payload, "statuses")
if isinstance(item, Mapping)
)
def _comments(token: str, repo: str, number: int) -> tuple[IssueComment, ...]:
comments: Final = _paginate(token, f"/repos/{repo}/issues/{number}/comments")
return tuple(
IssueComment(
author_login=_text(_nested(item, "user", "login")),
body=_text(item.get("body")),
updated_at=_parse_time(item.get("updated_at")),
)
for item in comments
if isinstance(item, Mapping)
)
def _reviews(token: str, repo: str, number: int) -> tuple[Review, ...]:
reviews: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/reviews")
return tuple(
Review(
author_login=_text(_nested(item, "user", "login")),
state=_text(item.get("state")),
body=_text(item.get("body")),
commit_id=_text(item.get("commit_id")),
submitted_at=_parse_time(item.get("submitted_at")),
)
for item in reviews
if isinstance(item, Mapping)
)
def _head_commit_date(token: str, repo: str, number: int) -> datetime:
commits: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/commits")
if not commits:
return datetime.min.replace(tzinfo=timezone.utc)
last: Final = commits[-1]
if not isinstance(last, Mapping):
return datetime.min.replace(tzinfo=timezone.utc)
return _parse_time(_nested(last, "commit", "committer", "date"))
def _mergeable_or_refetch(token: str, repo: str, pr: PullRequest) -> PullRequest:
if pr.mergeable is not None:
return pr
time.sleep(5)
return _load_pr(token, repo, pr.number)
def _gather_inputs(
token: str,
repo: str,
number: int,
base: str,
self_check_name: str,
allowlist: frozenset[str],
) -> EvaluationInputs:
pr: Final = _mergeable_or_refetch(token, repo, _load_pr(token, repo, number))
return EvaluationInputs(
pr=pr,
changed_files=_changed_files(token, repo, number),
required_contexts=_required_contexts(token, repo, base),
check_runs=_check_runs(token, repo, pr.head_sha),
statuses=_statuses(token, repo, pr.head_sha),
comments=_comments(token, repo, number),
reviews=_reviews(token, repo, number),
head_commit_date=_head_commit_date(token, repo, number),
self_check_name=self_check_name,
author_allowlist=allowlist,
)
def merge_request_body(pr: PullRequest) -> dict[str, str]:
return {"merge_method": "merge", "commit_title": f"{pr.title} (#{pr.number})", "sha": pr.head_sha}
def _merge(token: str, repo: str, pr: PullRequest) -> None:
status, _ = _request_allow_fail(token, "PUT", f"/repos/{repo}/pulls/{pr.number}/merge", merge_request_body(pr))
if status in (200, 405, 409):
print(f"auto-merge-price-sync: PR #{pr.number} merge call returned {status}")
return
raise RuntimeError(f"merge call for PR #{pr.number} returned {status}")
def main() -> int:
token: Final = os.environ.get("GH_TOKEN", "")
repo: Final = os.environ.get("REPO", "")
base: Final = os.environ.get("BASE_BRANCH", "main")
dry_run: Final = os.environ.get("DRY_RUN", "") != ""
self_check_name: Final = os.environ.get("SELF_CHECK_NAME", "auto-merge-price-sync")
allowlist: Final = frozenset(login.lower() for login in os.environ.get("PR_AUTHOR_ALLOWLIST", "").split() if login)
if not token:
print("auto-merge-price-sync: app credentials not configured")
return 0
if not repo:
print("auto-merge-price-sync: REPO not set", file=sys.stderr)
return 1
pr_number_env: Final = os.environ.get("PR_NUMBER", "")
candidates: Final = [int(pr_number_env)] if pr_number_env else _list_candidate_prs(token, repo, base, allowlist)
for number in candidates:
inputs: Final = _gather_inputs(token, repo, number, base, self_check_name, allowlist)
verdict: Final = evaluate(inputs)
for reason in verdict.reasons:
print(f"auto-merge-price-sync: PR #{number} hold: {reason}")
if not verdict.merge:
continue
print(f"auto-merge-price-sync: PR #{number} all gates green")
if dry_run:
print(f"auto-merge-price-sync: DRY_RUN merge suppressed for PR #{number}")
continue
_merge(token, repo, inputs.pr)
return 0
if __name__ == "__main__":
sys.exit(main())

73
.github/workflows/ai-gateway-image.yml vendored
View file

@ -1,73 +0,0 @@
name: ai-gateway image
on:
push:
paths:
- "litellm-rust/**"
- "litellm/**"
- "enterprise/**"
- "litellm-proxy-extras/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/workflows/ai-gateway-image.yml"
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
paths:
- "litellm-rust/**"
- "litellm/**"
- "enterprise/**"
- "litellm-proxy-extras/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/workflows/ai-gateway-image.yml"
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
ai-gateway-image:
name: ai-gateway release image
runs-on: ubuntu-latest
timeout-minutes: 60
permissions:
contents: read
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Build the release image
run: docker build -f litellm-rust/crates/ai-gateway/Dockerfile -t litellm-ai-gateway:${{ github.sha }} .
- name: Start the gateway and wait for readiness
env:
IMAGE: litellm-ai-gateway:${{ github.sha }}
run: |
docker run -d --name ai-gateway -p 4001:4001 \
-e LITELLM_MASTER_KEY=sk-ci-not-a-real-key \
-e OPENAI_API_KEY=sk-ci-not-a-real-key \
"$IMAGE"
for _ in $(seq 1 60); do
if curl -fsS http://127.0.0.1:4001/health/readiness; then
echo "gateway is serving readiness"
exit 0
fi
sleep 2
done
echo "gateway never became ready" >&2
docker logs ai-gateway >&2
exit 1
- name: Assert the gateway loaded the baked config
run: |
docker logs ai-gateway 2>&1 | tee gateway.log
grep 'via python config reader' gateway.log
- name: Stop the gateway
if: always()
run: docker rm -f ai-gateway || true

61
.github/workflows/auto-merge-price-sync.yml vendored Normal file
View file

@ -0,0 +1,61 @@
name: auto-merge-price-sync
on:
issue_comment:
types: [created, edited]
check_suite:
types: [completed]
status: {}
schedule:
- cron: "*/30 * * * *"
workflow_dispatch:
inputs:
pr-number:
description: "Evaluate only this PR number (empty = scan all open sync-bot PRs)"
required: false
default: ""
permissions:
contents: read
pull-requests: read
checks: read
statuses: read
concurrency:
group: auto-merge-price-sync
cancel-in-progress: false
jobs:
auto-merge-price-sync:
runs-on: ubuntu-latest
timeout-minutes: 15
env:
PROVIDER_INFO_SYNC_APP_ID: ${{ secrets.PROVIDER_INFO_SYNC_APP_ID }}
PROVIDER_INFO_SYNC_APP_PRIVATE_KEY: ${{ secrets.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Mint app token
id: app-token
if: ${{ env.PROVIDER_INFO_SYNC_APP_ID != '' && env.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY != '' }}
uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0
with:
app-id: ${{ secrets.PROVIDER_INFO_SYNC_APP_ID }}
private-key: ${{ secrets.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY }}
- name: Auto-merge eligible sync PRs
env:
GH_TOKEN: ${{ steps.app-token.outputs.token }}
REPO: ${{ github.repository }}
PR_NUMBER: ${{ (github.event.issue.pull_request && github.event.issue.number) || github.event.inputs.pr-number || '' }}
BASE_BRANCH: main
PR_AUTHOR_ALLOWLIST: "berriai-litellm-provider-info-sync[bot]"
SELF_CHECK_NAME: auto-merge-price-sync
run: python3 .github/scripts/auto_merge_price_sync.py

70
.github/workflows/test-rust.yml vendored
View file

@ -70,7 +70,7 @@ env:
jobs:
rust-lint:
runs-on: ubuntu-latest
timeout-minutes: 10
timeout-minutes: 15
defaults:
run:
working-directory: litellm-rust
@ -81,28 +81,48 @@ jobs:
- run: rustup toolchain install --no-self-update
- run: cargo fmt --check
- run: cargo fmt --all --check
- uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
- uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
with:
path: |
~/.cargo/registry
~/.cargo/git
litellm-rust/target
key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-${{ github.job }}-
workspaces: litellm-rust
cache-on-failure: true
- run: cargo clippy --workspace --all-targets --locked -- -D warnings
- run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings
- run: cargo clippy -p litellm-ai-gateway --all-targets --all-features --locked -- -D warnings
rust-test:
runs-on: ubuntu-latest
timeout-minutes: 30
timeout-minutes: 20
defaults:
run:
working-directory: litellm-rust
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- run: rustup toolchain install --no-self-update
- uses: taiki-e/install-action@d438492cf8a250514fa2d34b30bc3c0dc37c65ff # v2.87.8
with:
tool: cargo-nextest@0.9.143
- uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
with:
workspaces: litellm-rust
cache-on-failure: true
- run: cargo nextest run --workspace --locked
- run: cargo test --workspace --doc --locked
rust-wheel:
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
@ -118,24 +138,10 @@ jobs:
- run: rustup toolchain install --no-self-update
- uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
- uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
with:
path: |
~/.cargo/registry
~/.cargo/git
litellm-rust/target
key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-${{ github.job }}-
- run: cargo test --workspace --locked
working-directory: litellm-rust
- run: cargo test -p litellm-core --features bedrock-auth --locked
working-directory: litellm-rust
- run: cargo test -p litellm-ai-gateway --features server --locked
working-directory: litellm-rust
workspaces: litellm-rust
cache-on-failure: true
- run: uv build --wheel --out-dir dist

3
ARCHITECTURE.md
View file

@ -45,7 +45,7 @@ sequenceDiagram
ProxyServer->>Auth: user_api_key_auth()
Auth->>Redis: Check API key cache
Redis-->>Auth: Key info + spend limits
ProxyServer->>Hooks: max_budget_limiter, parallel_request_limiter
ProxyServer->>Hooks: parallel_request_limiter, cache_control_check
Hooks->>Redis: Check/increment rate limit counters
ProxyServer->>Router: route_request()
Router->>Main: litellm.acompletion()
@ -145,7 +145,6 @@ graph TD
| Hook | File | Purpose |
|------|------|---------|
| `max_budget_limiter` | `proxy/hooks/max_budget_limiter.py` | Enforce budget limits |
| `parallel_request_limiter` | `proxy/hooks/parallel_request_limiter_v3.py` | Rate limiting per key/user |
| `cache_control_check` | `proxy/hooks/cache_control_check.py` | Cache validation |
| `responses_id_security` | `proxy/hooks/responses_id_security.py` | Response ID validation |

5
litellm-proxy-extras/litellm_proxy_extras/migrations/20260916000000_add_key_total_spend/migration.sql Normal file
View file

@ -0,0 +1,5 @@
-- AlterTable
ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;

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

@ -426,6 +426,7 @@ model LiteLLM_VerificationToken {
key_alias String?
soft_budget_cooldown Boolean @default(false) // key-level state on if budget alerts need to be cooled down
spend Float @default(0.0)
total_spend Float @default(0.0)
expires DateTime?
models String[]
aliases Json @default("{}")
@ -528,6 +529,7 @@ model LiteLLM_DeletedVerificationToken {
key_alias String?
soft_budget_cooldown Boolean @default(false)
spend Float @default(0.0)
total_spend Float @default(0.0)
expires DateTime?
models String[]
aliases Json @default("{}")

340
litellm-rust/Cargo.lock generated
View file

@ -70,6 +70,12 @@ dependencies = [
"rustversion",
]
[[package]]
name = "arcstr"
version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d"
[[package]]
name = "async-compression"
version = "0.4.46"
@ -262,6 +268,17 @@ dependencies = [
"tokio",
]
[[package]]
name = "aws-smithy-eventstream"
version = "0.61.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5a9381123ab62d20c13082b151f30f962a3b112b727345394536dfa39a482944"
dependencies = [
"aws-smithy-types",
"bytes",
"crc32fast",
]
[[package]]
name = "aws-smithy-http"
version = "0.64.0"
@ -462,64 +479,6 @@ dependencies = [
"tracing",
]
[[package]]
name = "axum"
version = "0.7.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f"
dependencies = [
"async-trait",
"axum-core",
"base64 0.22.1",
"bytes",
"futures-util",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"hyper 1.10.1",
"hyper-util",
"itoa",
"matchit",
"memchr",
"mime",
"percent-encoding",
"pin-project-lite",
"rustversion",
"serde",
"serde_json",
"serde_path_to_error",
"serde_urlencoded",
"sha1",
"sync_wrapper",
"tokio",
"tokio-tungstenite",
"tower",
"tower-layer",
"tower-service",
"tracing",
]
[[package]]
name = "axum-core"
version = "0.4.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "09f2bd6146b97ae3359fa0cc6d6b376d9539582c7b4220f041a33ec24c226199"
dependencies = [
"async-trait",
"bytes",
"futures-util",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"mime",
"pin-project-lite",
"rustversion",
"sync_wrapper",
"tower-layer",
"tower-service",
"tracing",
]
[[package]]
name = "azure_core"
version = "1.1.0"
@ -1582,7 +1541,6 @@ dependencies = [
"http 1.4.2",
"http-body 1.1.0",
"httparse",
"httpdate",
"itoa",
"pin-project-lite",
"smallvec",
@ -1890,12 +1848,6 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "lazy_static"
version = "1.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
[[package]]
name = "libc"
version = "0.2.186"
@ -1903,40 +1855,24 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
[[package]]
name = "litellm-ai-gateway"
name = "linux-raw-sys"
version = "0.12.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53"
[[package]]
name = "litellm-auth"
version = "0.1.0"
dependencies = [
"axum",
"base64 0.22.1",
"futures-channel",
"futures-util",
"litellm-config",
"litellm-core",
"reqwest 0.12.28",
"rustls 0.23.42",
"rustls-native-certs",
"serde",
"serde_json",
"sha2 0.10.9",
"subtle",
"tokio",
"tokio-tungstenite",
"tower",
"tracing",
]
[[package]]
name = "litellm-config"
version = "0.1.0"
dependencies = [
"litellm-core",
"pyo3",
"serde_json",
"thiserror 2.0.19",
"tokio",
"veil",
]
[[package]]
name = "litellm-core"
name = "litellm-auth-aws"
version = "0.1.0"
dependencies = [
"aws-config",
@ -1945,13 +1881,86 @@ dependencies = [
"aws-sigv4",
"aws-smithy-runtime-api",
"aws-types",
"litellm-auth",
"moka",
"reqwest 0.12.28",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
]
[[package]]
name = "litellm-auth-azure"
version = "0.1.0"
dependencies = [
"azure_core",
"azure_identity",
"litellm-auth",
"moka",
"serde_json",
"sha2 0.10.9",
"strum",
"tokio",
"url",
]
[[package]]
name = "litellm-auth-gcp"
version = "0.1.0"
dependencies = [
"gcp_auth",
"litellm-auth",
"moka",
"serde_json",
"sha2 0.10.9",
"tokio",
]
[[package]]
name = "litellm-cache"
version = "0.1.0"
dependencies = [
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-cache-memory"
version = "0.1.0"
dependencies = [
"litellm-cache",
"rstest",
"serde_json",
"tokio",
]
[[package]]
name = "litellm-cache-redis"
version = "0.1.0"
dependencies = [
"litellm-cache",
"redis",
"redis-test",
"serde_json",
"tokio",
]
[[package]]
name = "litellm-core"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"bytes",
"data-url",
"futures-util",
"gcp_auth",
"litellm-auth",
"litellm-auth-aws",
"litellm-auth-azure",
"litellm-auth-gcp",
"mime_guess",
"moka",
"rand 0.8.7",
@ -1968,18 +1977,32 @@ dependencies = [
"thiserror 2.0.19",
"tokio",
"tokio-tungstenite",
"tracing",
"tracing-subscriber",
"url",
"veil",
]
[[package]]
name = "litellm-framing"
version = "0.1.0"
dependencies = [
"aws-smithy-eventstream",
"aws-smithy-types",
"bytes",
"futures-util",
"rstest",
"sse-stream",
"thiserror 2.0.19",
"tokio",
]
[[package]]
name = "litellm-python-bridge"
version = "0.1.0"
dependencies = [
"bytes",
"criterion",
"futures-util",
"litellm-auth",
"litellm-core",
"litellm-python-interop",
"litellm-token-counter",
@ -1990,7 +2013,6 @@ dependencies = [
"serde_json",
"tokio",
"tokio-tungstenite",
"tracing",
]
[[package]]
@ -2065,12 +2087,6 @@ version = "0.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c"
[[package]]
name = "matchit"
version = "0.7.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94"
[[package]]
name = "memchr"
version = "2.8.3"
@ -2172,6 +2188,16 @@ dependencies = [
"minimal-lexical",
]
[[package]]
name = "num-bigint"
version = "0.5.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0"
dependencies = [
"num-integer",
"num-traits",
]
[[package]]
name = "num-conv"
version = "0.2.2"
@ -2688,6 +2714,36 @@ dependencies = [
"crossbeam-utils",
]
[[package]]
name = "redis"
version = "1.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed"
dependencies = [
"arcstr",
"combine",
"itoa",
"num-bigint",
"percent-encoding",
"ryu",
"sha1_smol",
"socket2 0.6.5",
"url",
"xxhash-rust",
]
[[package]]
name = "redis-test"
version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "804d36862e4323b69f96440cbb13c9894fc90176abdeaf91264e21d5d77f6aca"
dependencies = [
"rand 0.9.5",
"redis",
"socket2 0.6.5",
"tempfile",
]
[[package]]
name = "redox_syscall"
version = "0.5.18"
@ -2878,6 +2934,19 @@ dependencies = [
"semver",
]
[[package]]
name = "rustix"
version = "1.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "891efababe418670775f199f0d233d84843c227a0949a883ce15b37c78d6629d"
dependencies = [
"bitflags",
"errno",
"libc",
"linux-raw-sys",
"windows-sys 0.61.2",
]
[[package]]
name = "rustls"
version = "0.21.12"
@ -3128,6 +3197,12 @@ dependencies = [
"digest 0.10.7",
]
[[package]]
name = "sha1_smol"
version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d"
[[package]]
name = "sha2"
version = "0.10.9"
@ -3150,15 +3225,6 @@ dependencies = [
"digest 0.11.3",
]
[[package]]
name = "sharded-slab"
version = "0.1.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6"
dependencies = [
"lazy_static",
]
[[package]]
name = "shlex"
version = "2.0.1"
@ -3241,6 +3307,19 @@ dependencies = [
"unicode-segmentation",
]
[[package]]
name = "sse-stream"
version = "0.2.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c25ac7aff0abd1dbc474536e40416e1102c7dd9bfba0b9861c6d357f835dcfb4"
dependencies = [
"bytes",
"futures-util",
"http-body 1.1.0",
"http-body-util",
"pin-project-lite",
]
[[package]]
name = "stable_deref_trait"
version = "1.2.1"
@ -3340,6 +3419,19 @@ version = "0.13.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
[[package]]
name = "tempfile"
version = "3.27.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
dependencies = [
"fastrand",
"getrandom 0.4.3",
"once_cell",
"rustix",
"windows-sys 0.61.2",
]
[[package]]
name = "thiserror"
version = "1.0.69"
@ -3380,15 +3472,6 @@ dependencies = [
"syn 3.0.0",
]
[[package]]
name = "thread_local"
version = "1.1.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070"
dependencies = [
"cfg-if",
]
[[package]]
name = "time"
version = "0.3.53"
@ -3606,7 +3689,6 @@ dependencies = [
"tokio",
"tower-layer",
"tower-service",
"tracing",
]
[[package]]
@ -3650,7 +3732,6 @@ version = "0.1.44"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100"
dependencies = [
"log",
"pin-project-lite",
"tracing-attributes",
"tracing-core",
@ -3686,17 +3767,6 @@ dependencies = [
"tracing",
]
[[package]]
name = "tracing-subscriber"
version = "0.3.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319"
dependencies = [
"sharded-slab",
"thread_local",
"tracing-core",
]
[[package]]
name = "try-lock"
version = "0.2.5"
@ -4245,6 +4315,12 @@ version = "0.13.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4"
[[package]]
name = "xxhash-rust"
version = "0.8.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "aee1b19627c7c60102ab80d3a9cbe18de90bfe03bfa6c3715447681f0e8c8af6"
[[package]]
name = "yoke"
version = "0.8.3"

23
litellm-rust/Cargo.toml
View file

@ -1,12 +1,5 @@
[workspace]
members = [
"crates/core",
"crates/token-counter",
"crates/config",
"crates/ai-gateway",
"crates/python-interop",
"crates/python-bridge",
]
members = ["crates/*"]
resolver = "2"
[workspace.package]
@ -17,14 +10,15 @@ repository = "https://github.com/BerriAI/litellm"
[workspace.dependencies]
bytes = "1"
tracing = "0.1"
tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }
litellm-core = { path = "crates/core" }
litellm-auth = { path = "crates/auth" }
litellm-auth-aws = { path = "crates/auth-aws" }
litellm-auth-azure = { path = "crates/auth-azure" }
litellm-auth-gcp = { path = "crates/auth-gcp" }
litellm-cache = { path = "crates/cache" }
litellm-cache-memory = { path = "crates/cache-memory" }
litellm-token-counter = { path = "crates/token-counter" }
litellm-config = { path = "crates/config" }
litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
litellm-python-interop = { path = "crates/python-interop" }
axum = "0.7"
pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
pythonize = "0.29.0"
@ -42,9 +36,6 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"]
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
base64 = "0.22"
gcp_auth = "0.12.7"
azure_core = "1.0.0"
azure_identity = { version = "1.0.0", features = ["tokio"] }
moka = { version = "0.12.16", features = ["future"] }
strum = { version = "0.28.0", features = ["derive"] }
url = "2.5.8"

56
litellm-rust/crates/ai-gateway/Cargo.toml
View file

@ -1,56 +0,0 @@
[package]
name = "litellm-ai-gateway"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[lib]
name = "litellm_ai_gateway"
[[bin]]
name = "litellm-ai-gateway"
path = "src/main.rs"
required-features = ["server"]
[[bin]]
name = "trace-parity-gateway"
path = "src/bin/trace_parity_gateway.rs"
required-features = ["trace-parity"]
[dependencies]
tracing.workspace = true
litellm-core = { workspace = true, features = ["bedrock-auth"] }
litellm-config.workspace = true
# reqwest (rustls + json) is used by io/ocr and ships realtime logs to the
# Python proxy callbacks API.
reqwest.workspace = true
# rustls and its root store are direct dependencies so `io::tls` can build the
# one TLS config the outbound dials use; see that module for why it has to.
rustls.workspace = true
rustls-native-certs.workspace = true
# `sync` powers the bounded mpsc channel the realtime logger drains.
tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time", "sync"] }
tokio-tungstenite.workspace = true
futures-util.workspace = true
serde_json.workspace = true
base64.workspace = true
axum = { workspace = true, features = ["ws"], optional = true }
serde.workspace = true
subtle = { workspace = true, optional = true }
# sha2 hashes the master key into user_api_key_hash (matches the proxy's
# SHA-256 hash_token) so the plaintext credential never enters a log payload.
sha2 = { workspace = true, optional = true }
tower = { version = "0.5.3", features = ["util"], optional = true }
[features]
default = []
server = ["dep:axum", "dep:subtle", "dep:sha2"]
# Build the gateway's config from the proxy YAML via an embedded Python
# interpreter (links libpython; requires `litellm` importable at runtime).
python-config = ["litellm-config/python"]
trace-parity = ["server", "dep:tower", "litellm-core/observability"]
[dev-dependencies]
futures-channel = "0.3"
tower = { version = "0.5.3", features = ["util"] }

109
litellm-rust/crates/ai-gateway/Dockerfile
View file

@ -1,109 +0,0 @@
# Multi-stage build for the LiteLLM Rust AI Gateway (realtime WebSocket proxy).
#
# Build context is the **repo root** so we can install `litellm` from this repo's
# source (the gateway loads its model_list via litellm.proxy.read_model_list,
# which is not in any PyPI release yet) AND build the rust workspace under
# litellm-rust/.
#
# docker build -f litellm-rust/crates/ai-gateway/Dockerfile -t litellm-ai-gateway .
#
# No secrets live in this file. Runtime config (LITELLM_MASTER_KEY,
# OPENAI_API_KEY referenced by config.yaml, etc.) is injected as environment
# variables at deploy time.
# ---- Chef -------------------------------------------------------------------
# cargo-chef caches the dependency build so only the gateway crate recompiles on
# a source-only change. python3-dev is present in every rust stage because the
# `python-config` feature links libpython via pyo3 (even in the cook step), and
# python3-pip builds the litellm wheel in the builder stage.
FROM rust:1.98-slim-bookworm AS chef
ENV PYO3_PYTHON=python3.11
# rustup reads rust-toolchain.toml from any parent of the working directory, so
# copying it in is what keeps every cargo call below on the repo's pinned
# channel rather than on whatever the base image happens to ship.
COPY rust-toolchain.toml /build/rust-toolchain.toml
WORKDIR /build/litellm-rust
RUN apt-get update \
&& apt-get install -y --no-install-recommends \
python3 python3-dev python3-pip pkg-config libssl-dev clang \
&& rm -rf /var/lib/apt/lists/* \
&& cargo install cargo-chef --locked --version 0.1.77
# ---- Planner ----------------------------------------------------------------
# Produce the dependency recipe from the rust workspace manifests + Cargo.lock.
FROM chef AS planner
COPY litellm-rust/ .
RUN cargo chef prepare --recipe-path recipe.json
# ---- Builder ----------------------------------------------------------------
FROM chef AS builder
# Cook (compile) just the dependencies first — this layer is cached and reused
# whenever only gateway source changes.
COPY --from=planner /build/litellm-rust/recipe.json recipe.json
RUN cargo chef cook --locked --release \
-p litellm-ai-gateway --features server,python-config \
--recipe-path recipe.json
# Now copy the real sources and build the gateway binary. Deps are already cooked
# above, so this step only recompiles the gateway crate.
COPY litellm-rust/ .
RUN cargo build --locked --release -p litellm-ai-gateway --bin litellm-ai-gateway --features server,python-config
# The root pyproject builds with maturin against litellm-rust/crates/python-bridge,
# so the wheel is built here, next to the crate sources and the cargo toolchain,
# and the runtime stage installs the artifact instead of compiling anything.
# litellm[proxy] pins litellm-enterprise and litellm-proxy-extras to the versions
# in this repo, and those hit PyPI hours after every version bump merges, so both
# wheels are built from the repo too instead of being resolved from PyPI.
COPY pyproject.toml README.md LICENSE /build/
COPY litellm/ /build/litellm/
COPY enterprise/ /build/enterprise/
COPY litellm-proxy-extras/ /build/litellm-proxy-extras/
RUN pip3 wheel --no-cache-dir --no-deps --wheel-dir /build/dist \
/build /build/enterprise /build/litellm-proxy-extras
# ---- Runtime ----------------------------------------------------------------
# python:3.11-slim-bookworm ships libpython3.11, matching the builder's PyO3
# 3.11 ABI so the embedded interpreter links and imports cleanly.
FROM python:3.11-slim-bookworm AS runtime
# CA certificates for outbound TLS to the OpenAI realtime endpoint.
RUN apt-get update \
&& apt-get install -y --no-install-recommends ca-certificates \
&& rm -rf /var/lib/apt/lists/*
WORKDIR /app
# Install litellm (with proxy extras) FROM THIS REPO'S SOURCE so
# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. The two
# sibling wheels come from the builder as well, so the pins in litellm[proxy]
# resolve against them and never wait on a PyPI publish.
COPY --from=builder /build/dist/*.whl /tmp/wheels/
RUN wheel="$(ls /tmp/wheels/litellm-*.whl)" \
&& pip install --no-cache-dir \
/tmp/wheels/litellm_enterprise-*.whl \
/tmp/wheels/litellm_proxy_extras-*.whl \
"${wheel}[proxy]" \
&& rm -rf /tmp/wheels
# The compiled gateway binary (pure-Rust realtime hot path; Python is load-time
# only).
COPY --from=builder /build/litellm-rust/target/release/litellm-ai-gateway /usr/local/bin/litellm-ai-gateway
# Default config.yaml. A real deploy can override this (e.g. mount a Render
# secret file at the same path) — never bake secrets into the image.
COPY litellm-rust/crates/ai-gateway/config.yaml /app/config.yaml
# Bind to all interfaces (Render routes to 0.0.0.0:$PORT) and load the model_list
# from config.yaml via the embedded python config reader.
ENV HOST=0.0.0.0 \
LITELLM_CONFIG_PATH=/app/config.yaml
# Drop to a non-root user. The realtime hot path needs no root privileges, so
# running unprivileged limits blast radius if the process is ever compromised.
# The binary in /usr/local/bin is world-executable (COPY default mode 755); we
# only need /app (and the config.yaml it reads) owned by the unprivileged user.
RUN useradd --system --no-create-home --uid 10001 appuser \
&& chown -R appuser:appuser /app
USER appuser
ENTRYPOINT ["/usr/local/bin/litellm-ai-gateway"]

54
litellm-rust/crates/ai-gateway/Dockerfile.dockerignore
View file

@ -1,54 +0,0 @@
# Dockerfile-specific ignore-file for the Rust AI Gateway build.
#
# The build context is the repo root (so the image can pip install litellm from
# source AND build the rust workspace). BuildKit honors `<Dockerfile>.dockerignore`
# next to the Dockerfile and it takes precedence over the repo-root `.dockerignore`,
# so this file shrinks the (large) repo-root context for THIS build only without
# touching the root `.dockerignore` used by the main litellm images.
#
# Strategy: ignore everything, then re-include only what the build needs:
# - litellm/ (pip install . needs the full package + proxy reader)
# - litellm-rust/ (the rust workspace; Cargo.lock + crate sources)
# - enterprise/ (litellm/proxy/enterprise symlinks into it; maturin walks it)
# - litellm-proxy-extras/ (built into a wheel alongside enterprise/ for litellm[proxy])
# - pyproject.toml / README.md / LICENSE (packaging metadata for the wheel build)
# - rust-toolchain.toml (the pinned channel every cargo call in the build uses)
*
# --- re-include the build inputs ---
!litellm/
!litellm-rust/
!enterprise/
!litellm-proxy-extras/
!pyproject.toml
!rust-toolchain.toml
!README.md
!LICENSE
# --- prune heavy / irrelevant subpaths back out of the re-included trees ---
# Rust build artifacts (huge; regenerated in the builder).
**/target/
# Committed python distribution artifacts; the wheel build does not read them.
enterprise/dist/
litellm-proxy-extras/dist/
# Python caches and compiled bytecode.
**/__pycache__/
**/*.pyc
**/*.pyo
**/.pytest_cache/
**/.ruff_cache/
**/.mypy_cache/
# Node / UI build output bundled under the python package (not needed to import
# litellm.proxy.read_model_list).
**/node_modules/
litellm/proxy/_experimental/out/
# Tests, logs, and local scratch.
**/tests/
**/test/
*.log
log.txt
*.tgz
# VCS / editor / CI metadata that may live under re-included trees.
**/.git/
.git/
**/.DS_Store

206
litellm-rust/crates/ai-gateway/README.md
View file

@ -1,206 +0,0 @@
# LiteLLM Rust AI Gateway
A minimal Axum service that fronts OpenAI's realtime API. Clients open a
WebSocket to `GET /v1/realtime`; the gateway authenticates, selects a deployment,
dials OpenAI upstream, and splices the two sockets frame-by-frame.
## Crates
`litellm-rust` has six crates. A crate is a layer or shared foundation, not a route:
| Crate | Role |
|-------|------|
| litellm-core | The LiteLLM SDK in Rust — per-route entrypoints (`messages::messages()`) that resolve the provider, transform, and make the call; plus types, provider transforms, and the router. |
| litellm-token-counter | Standalone input token counting shared by host integrations without pulling in the full SDK. |
| litellm-config | Config-loading boundary. Returns resolved deployments and optionally delegates loading to Python. |
| litellm-ai-gateway | The Axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. |
| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. |
Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers, token counter, and Python interop.
- **Client endpoint:** `wss://<host>/v1/realtime?model=<model>` (WebSocket)
- **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset)
- **Health:** `GET /health/readiness`, `GET /health/liveness`
- **Request logs:** POSTed to a LiteLLM proxy at `/v1/rust_control_plane/logs` (see [Request logging](#request-logging))
> **Realtime serving is pure Rust.** Python is used at **load time only** — to
> read the config once at boot. The realtime hot path never touches Python.
The former `/health/gil` route and its acquisition counter were removed. They
only observed the single startup config load and did not prove that every GIL
acquisition was instrumented
## Configuration (config.yaml)
The gateway loads its `model_list` from a **config.yaml**, the same as the
LiteLLM proxy. Point `LITELLM_CONFIG_PATH` at the file:
```yaml
# config.yaml
model_list:
- model_name: gpt-realtime
litellm_params:
model: openai/gpt-realtime
api_key: os.environ/OPENAI_API_KEY
```
```bash
LITELLM_CONFIG_PATH=./config.yaml ./litellm-ai-gateway
```
At boot `litellm-config` calls into `litellm.proxy.read_model_list` and returns
resolved deployments to the gateway, which constructs the router. The Python
backend still reuses the **real proxy config reader** (`ProxyConfig.get_config`),
so everything the proxy supports in config.yaml works here too:
- `include:` to merge in other config files,
- `os.environ/VAR` secret references (resolved via the secret manager, never
inlined),
- DB-stored models (when a database is configured).
Secrets stay out of the config — reference them with `os.environ/...` and set
the env var at deploy time. The shipped Docker image is built with the
`python-config` feature and **bundles litellm**, so config loading works out of
the box; the default baked config lives at `/app/config.yaml` and can be
overridden at deploy time (e.g. a Render secret file mounted at the same path).
### Environment variables
| Var | Required | Default | Purpose |
|---|---|---|---|
| `LITELLM_CONFIG_PATH` | yes (config mode) | — | Path to the config.yaml the gateway loads its `model_list` from. The Docker image defaults this to `/app/config.yaml`. |
| `LITELLM_MASTER_KEY` | yes | — | Bearer token clients must send. Unset ⇒ all `/v1/realtime` requests are rejected (fail closed). |
| `OPENAI_API_KEY` | yes | — | Upstream OpenAI key. Referenced by config.yaml as `os.environ/OPENAI_API_KEY` for the gateway→OpenAI dial. |
| `HOST` | no | `127.0.0.1` | **Set to `0.0.0.0` in any container/deploy** or external traffic is refused. |
| `PORT` | no | `4001` | Listen port. Render and most PaaS inject this automatically. |
| `LITELLM_PROXY_BASE_URL` | no | `http://localhost:4000` | LiteLLM proxy that request logs are POSTed to. See [Request logging](#request-logging). |
> Secrets (`LITELLM_MASTER_KEY`, `OPENAI_API_KEY`) are never baked into the image
> or `render.yaml` — inject them at deploy time only.
### Lean env stand-in (fallback)
If the binary is built **without** `python-config` (default features), or
`LITELLM_CONFIG_PATH` is unset, the gateway falls back to a single-deployment
stand-in built from the environment:
| Var | Default | Purpose |
|---|---|---|
| `OPENAI_REALTIME_MODEL` | `gpt-realtime` | The single deployment's model name (also the `?model=` clients pass). |
The default workspace build links no libpython and needs no config file. This
fallback mode only supports one hard-coded OpenAI deployment. **config.yaml is the recommended path** — use the
stand-in only for the leanest possible build.
## Request logging
The gateway runs no spend logic. When a session ends it builds one
`StandardLoggingPayload` and POSTs it to `{LITELLM_PROXY_BASE_URL}/v1/rust_control_plane/logs`
(admin-only, bearer = `LITELLM_MASTER_KEY`), and the proxy replays it through its
normal callbacks (spend logs, Langfuse, etc.). The POST is non-blocking: a bounded
channel drained by a background worker, dropping with a counter if the proxy is
down. It sends one payload per session. Both env vars are in the table above.
Worker tuning, rarely needed: `LITELLM_LOG_CHANNEL_CAPACITY` (4096),
`LITELLM_LOG_BATCH_SIZE` (256), `LITELLM_LOG_FLUSH_INTERVAL_MS` (500).
## Build & run with Docker
The image is built `--features server,python-config` and installs litellm **from this
repo's source** (the config reader is newer than any PyPI release), so the build
**context is the repo root**:
```bash
# from the repo root
docker build -f litellm-rust/crates/ai-gateway/Dockerfile -t litellm-ai-gateway .
docker run --rm -p 4001:4001 \
-e HOST=0.0.0.0 -e PORT=4001 \
-e LITELLM_MASTER_KEY=sk-local \
-e OPENAI_API_KEY=$OPENAI_API_KEY \
litellm-ai-gateway # LITELLM_CONFIG_PATH defaults to /app/config.yaml
# smoke test
curl -s -o /dev/null -w '%{http_code}\n' localhost:4001/health/readiness # -> 200
curl -s -o /dev/null -w '%{http_code}\n' localhost:4001/v1/realtime # -> 401 (auth fails closed)
```
On boot you should see `loaded model_list from /app/config.yaml via python
config reader` — that confirms the config path (not the env stand-in fallback).
To use your own config, mount it over the default:
```bash
docker run --rm -p 4001:4001 \
-e HOST=0.0.0.0 -e LITELLM_MASTER_KEY=sk-local -e OPENAI_API_KEY=$OPENAI_API_KEY \
-v $(pwd)/my-config.yaml:/app/config.yaml:ro \
litellm-ai-gateway
```
### Cargo-only (no Docker)
```bash
# config.yaml mode — needs litellm importable in the active python env
LITELLM_CONFIG_PATH=./crates/ai-gateway/config.yaml \
cargo run --release -p litellm-ai-gateway --features server,python-config
# env stand-in mode — no python, no config
cargo run --release -p litellm-ai-gateway --features server
```
## Deploy on Render
The service is a Docker **web service**; Render terminates TLS and supports
WebSockets, so the public endpoint is `wss://<service>.onrender.com/v1/realtime`.
### Option A — Blueprint (`render.yaml`)
`crates/ai-gateway/render.yaml` describes the service (Docker runtime,
`healthCheckPath: /health/readiness`, repo-root `dockerContext: .`,
`dockerfilePath: ./litellm-rust/crates/ai-gateway/Dockerfile`,
`LITELLM_CONFIG_PATH: /app/config.yaml`). `LITELLM_MASTER_KEY` and
`OPENAI_API_KEY` are `sync: false` — set them in the dashboard after the first
deploy. To use a non-default model_list, mount a **Render Secret File** at
`/app/config.yaml`. Point a Render Blueprint at this repo/branch and apply.
### Option B — Render API
```bash
# create a Docker web service from this repo+branch, then set env vars:
curl -X POST https://api.render.com/v1/services \
-H "Authorization: Bearer $RENDER_API_KEY" -H "Content-Type: application/json" \
-d '{
"type": "web_service", "name": "litellm-rust-ai-gateway",
"ownerId": "<owner-id>", "repo": "https://github.com/BerriAI/litellm",
"branch": "<branch-with-this-dockerfile>",
"serviceDetails": {
"env": "docker",
"envSpecificDetails": {
"dockerfilePath": "./litellm-rust/crates/ai-gateway/Dockerfile",
"dockerContext": "."
},
"healthCheckPath": "/health/readiness"
}
}'
# then set env vars LITELLM_MASTER_KEY, OPENAI_API_KEY, HOST=0.0.0.0,
# LITELLM_CONFIG_PATH=/app/config.yaml
```
Health check path **must** be `/health/readiness`. `autoDeploy` is off by default
in the blueprint — trigger deploys manually (or flip it on) to pick up new commits.
## Scaling
Concurrency is what matters, not total connections: each in-flight session holds
one client socket + one upstream socket. To scale, raise the instance count /
enable autoscaling on the Render service (e.g. baseline 10, max 100). Each
instance needs file descriptors for `2 × peak_concurrent_sessions` — raise
`ulimit -n` if you push very high concurrency.
## Latency note
The gateway adds the cost of one extra hop: client→gateway, then a fresh
gateway→OpenAI realtime handshake (TLS + WS upgrade + `session.created`). In
benchmarks this is ~100–150 ms of added session-establishment time; first-audio
and steady-state streaming add no measurable overhead. To minimize it, deploy the
gateway in the Render region with the lowest RTT to OpenAI's realtime endpoint.

13
litellm-rust/crates/ai-gateway/config.yaml
View file

@ -1,13 +0,0 @@
# Sample realtime config for the LiteLLM Rust AI Gateway.
#
# litellm-config resolves this model_list at boot through the Python config
# reader (litellm.proxy.read_model_list), then the gateway builds its router.
# Includes, environment secrets, and database-stored models still work.
#
# Secrets are referenced (never inlined) via os.environ/. A real deploy can
# override this file (e.g. mount a Render secret file at LITELLM_CONFIG_PATH).
model_list:
- model_name: gpt-realtime
litellm_params:
model: openai/gpt-realtime
api_key: os.environ/OPENAI_API_KEY

35
litellm-rust/crates/ai-gateway/render.yaml
View file

@ -1,35 +0,0 @@
# Render blueprint for the LiteLLM Rust AI Gateway (realtime WebSocket proxy).
#
# Single instance for now (no autoscaling). The public endpoint is a
# WebSocket served over TLS: wss://<service>.onrender.com/v1/realtime
#
# Paths are relative to the **repo root** (Render's convention). The build
# context is the repo root so the image can install litellm from source — the
# gateway loads its model_list via litellm.proxy.read_model_list at boot.
#
# Secrets (LITELLM_MASTER_KEY, OPENAI_API_KEY) are marked sync: false — set
# them in the Render dashboard or via the API, never inline here.
services:
- type: web
name: litellm-rust-ai-gateway
runtime: docker
plan: standard
dockerfilePath: ./litellm-rust/crates/ai-gateway/Dockerfile
dockerContext: .
healthCheckPath: /health/readiness
numInstances: 1
envVars:
# The gateway loads its model_list from this config.yaml via the embedded
# python config reader. The image bakes a default config at /app/config.yaml;
# a real deploy can override it by mounting a Render secret file at this
# same path (Dashboard → Environment → Secret Files) — never inline secrets.
- key: LITELLM_CONFIG_PATH
value: /app/config.yaml
- key: HOST
value: 0.0.0.0
# Bearer token clients must send on /v1/realtime (fail closed if unset).
- key: LITELLM_MASTER_KEY
sync: false
# Referenced by config.yaml as os.environ/OPENAI_API_KEY for the upstream dial.
- key: OPENAI_API_KEY
sync: false

288
litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs
View file

@ -1,288 +0,0 @@
use litellm_core::audio_transcription::{
AudioTranscriptionRequest as CoreAudioTranscriptionRequest, ProviderAudioTranscriptionRequest,
prepare_audio_transcription_provider_call,
};
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
use litellm_core::error::Error;
use serde_json::{Map, Value, json};
use std::future::Future;
use std::pin::Pin;
use super::types::PreparedAudioTranscriptionRequest;
use crate::integrations::custom_guardrail::{
CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
};
use crate::integrations::custom_logger::{
CallType, CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails,
};
use crate::integrations::types::{
RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload,
};
pub(crate) struct AudioTranscriptionLifecycleHooks {
logger_runner: CustomLoggerRunner,
guardrail_runner: CustomGuardrailRunner,
request_metadata: RequestMetadata,
}
type AudioFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
type AudioLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
impl AudioTranscriptionLifecycleHooks {
pub(crate) fn new(
logger_runner: CustomLoggerRunner,
guardrail_runner: CustomGuardrailRunner,
request_metadata: RequestMetadata,
) -> Self {
Self {
logger_runner,
guardrail_runner,
request_metadata,
}
}
async fn run_pre_call_guardrails(
&self,
request: PreparedAudioTranscriptionRequest,
) -> Result<PreparedAudioTranscriptionRequest, Error> {
if self.guardrail_runner.is_empty() {
return Ok(request);
}
let (guardrail_request, _) = self
.guardrail_runner
.run_pre_call(
&guardrail_context(&self.request_metadata),
GuardrailRequest::new(json!({
"model": request.model,
"custom_llm_provider": request.custom_llm_provider,
"audio": request.audio,
"optional_params": request.optional_params,
})),
)
.await
.map_err(guardrail_error_to_core_error)?;
let Value::Object(mut data) = guardrail_request.data else {
return Err(Error::InvalidRequest(
"audio transcription pre_call guardrail must return an object".to_string(),
));
};
let audio = data.remove("audio").ok_or_else(|| {
Error::InvalidRequest("audio transcription guardrail removed audio".to_string())
})?;
let optional_params = match data.remove("optional_params") {
Some(Value::Object(value)) => value,
Some(_) => {
return Err(Error::InvalidRequest(
"audio transcription optional_params must be an object".to_string(),
));
}
None => Map::new(),
};
Ok(PreparedAudioTranscriptionRequest {
audio,
optional_params,
..request
})
}
async fn prepare_provider_request(
&self,
request: PreparedAudioTranscriptionRequest,
) -> Result<ProviderAudioTranscriptionRequest, Error> {
let PreparedAudioTranscriptionRequest {
model,
custom_llm_provider,
audio,
api_key,
api_base,
extra_headers,
optional_params,
timeout,
..
} = request;
let provider_request =
prepare_audio_transcription_provider_call(CoreAudioTranscriptionRequest {
model: &model,
audio,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: Some(&custom_llm_provider),
extra_headers,
optional_params,
timeout,
})?;
self.run_during_call_guardrails(provider_request).await
}
async fn run_during_call_guardrails(
&self,
request: ProviderAudioTranscriptionRequest,
) -> Result<ProviderAudioTranscriptionRequest, Error> {
if self.guardrail_runner.is_empty() {
return Ok(request);
}
let (guardrail_request, _) = self
.guardrail_runner
.run_during_call(
&guardrail_context(&self.request_metadata),
GuardrailRequest::new(json!({
"model": request.model(),
"custom_llm_provider": request.custom_llm_provider(),
"url": request.url(),
"body": request.body(),
})),
)
.await
.map_err(guardrail_error_to_core_error)?;
let Value::Object(mut data) = guardrail_request.data else {
return Err(Error::InvalidRequest(
"audio transcription during_call guardrail must return an object".to_string(),
));
};
let body = data.remove("body").ok_or_else(|| {
Error::InvalidRequest("audio transcription guardrail removed body".to_string())
})?;
Ok(request.with_body(body))
}
fn logging_payload(
&self,
context: &CallLifecycleContext,
timing: &CallLifecycleTiming,
) -> StandardLoggingPayload {
StandardLoggingPayload {
id: context.litellm_call_id.clone(),
litellm_call_id: context.litellm_call_id.clone(),
call_type: context.call_type.clone(),
model: context.model.clone(),
custom_llm_provider: context.custom_llm_provider.clone(),
response_cost: 0.0,
prompt_tokens: 0,
completion_tokens: 0,
total_tokens: 0,
start_time: timing.start_time,
end_time: timing.end_time,
stream: false,
metadata: StandardLoggingMetadata {
user_api_key_hash: self.request_metadata.user_api_key_hash.clone(),
user_api_key_user_id: self.request_metadata.user_api_key_user_id.clone(),
user_api_key_team_id: self.request_metadata.user_api_key_team_id.clone(),
..Default::default()
},
messages: None,
}
}
}
impl CallLifecycleHooks<PreparedAudioTranscriptionRequest, ProviderAudioTranscriptionRequest, Value>
for AudioTranscriptionLifecycleHooks
{
type PreCallFuture<'a> = AudioFuture<'a, PreparedAudioTranscriptionRequest>;
type DuringCallFuture<'a> = AudioFuture<'a, ProviderAudioTranscriptionRequest>;
type SuccessFuture<'a> = AudioLogFuture<'a>;
type FailureFuture<'a> = AudioLogFuture<'a>;
fn async_pre_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: PreparedAudioTranscriptionRequest,
) -> Self::PreCallFuture<'a> {
Box::pin(async move { self.run_pre_call_guardrails(request).await })
}
fn async_during_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: PreparedAudioTranscriptionRequest,
) -> Self::DuringCallFuture<'a> {
Box::pin(async move { self.prepare_provider_request(request).await })
}
fn async_log_success_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
response: &'a Value,
timing: &'a CallLifecycleTiming,
) -> Self::SuccessFuture<'a> {
Box::pin(async move {
if self.logger_runner.is_empty() {
return;
}
self.logger_runner
.async_log_success_event(
&ModelCallDetails::from_standard_logging_payload(
self.logging_payload(context, timing),
),
&CallbackValue::new("audio_transcription", response.clone()),
CallbackTiming::new(timing.start_time, timing.end_time),
)
.await;
})
}
fn async_log_failure_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
error: &'a Error,
timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a> {
Box::pin(async move {
if self.logger_runner.is_empty() {
return;
}
let logging_error = LoggingError {
message: error.to_string(),
kind: core_error_kind(error).to_string(),
};
self.logger_runner
.async_log_failure_event(
&ModelCallDetails::from_standard_logging_payload(
self.logging_payload(context, timing),
)
.with_failure_error(logging_error.clone()),
Some(&CallbackValue::new(
"error",
json!({"message": logging_error.message, "kind": logging_error.kind}),
)),
CallbackTiming::new(timing.start_time, timing.end_time),
)
.await;
})
}
}
fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext {
GuardrailContext {
call_type: CallType::Other("audio_transcription".to_string()),
selected_guardrails: Vec::new(),
metadata: std::collections::HashMap::new(),
user_api_key_hash: metadata.user_api_key_hash.clone(),
user_api_key_user_id: metadata.user_api_key_user_id.clone(),
user_api_key_team_id: metadata.user_api_key_team_id.clone(),
trace_parent: None,
}
}
fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
Error::InvalidRequest(format!("{}: {}", error.kind, error.message))
}
fn core_error_kind(error: &Error) -> &'static str {
match error {
Error::Auth(_)
| Error::MissingApiKey { .. }
| Error::MissingAzureAiCredentials
| Error::MissingAzureDocumentIntelligenceCredentials
| Error::MissingReductoApiKey => "AuthError",
Error::InvalidProvider(_) => "InvalidProvider",
Error::InvalidRequest(_) => "InvalidRequest",
Error::InvalidType { .. } => "InvalidType",
Error::MissingField(_) | Error::MissingDocumentUrl => "MissingField",
Error::Http { .. } => "HttpError",
Error::InvalidResponse(_) => "InvalidResponse",
Error::Network(_) => "NetworkError",
Error::Connect(_) => "ConnectError",
Error::Routing(_) => "RoutingError",
Error::Unsupported(_) => "UnsupportedRequest",
}
}

23
litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs
View file

@ -1,23 +0,0 @@
use litellm_core::Error;
use litellm_core::audio_transcription::execute_audio_transcription_provider_call;
use litellm_core::call_lifecycle::CallLifecycle;
use serde_json::Value;
mod hooks;
mod prepare;
mod types;
pub use types::AudioTranscriptionRequest;
use prepare::{PreparedAudioTranscriptionCall, prepare_audio_transcription_call};
pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
let PreparedAudioTranscriptionCall { request, hooks } =
prepare_audio_transcription_call(request);
CallLifecycle::default()
.run_request(request, &hooks, execute_audio_transcription_provider_call)
.await
}
#[cfg(test)]
mod tests;

55
litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs
View file

@ -1,55 +0,0 @@
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use super::hooks::AudioTranscriptionLifecycleHooks;
use super::types::{AudioTranscriptionRequest, PreparedAudioTranscriptionRequest};
use crate::integrations::custom_guardrail::CustomGuardrailRunner;
use crate::integrations::custom_logger::CustomLoggerRunner;
pub(crate) struct PreparedAudioTranscriptionCall {
pub(crate) request: PreparedAudioTranscriptionRequest,
pub(crate) hooks: AudioTranscriptionLifecycleHooks,
}
pub(crate) fn prepare_audio_transcription_call(
request: AudioTranscriptionRequest<'_>,
) -> PreparedAudioTranscriptionCall {
let call_id = request
.litellm_call_id
.map(str::to_string)
.unwrap_or_else(new_audio_transcription_call_id);
let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
.unwrap_or(CustomLlmProvider {
model: request.model,
custom_llm_provider: "bedrock",
});
PreparedAudioTranscriptionCall {
request: PreparedAudioTranscriptionRequest {
model: provider_info.model.to_string(),
custom_llm_provider: provider_info.custom_llm_provider.to_string(),
litellm_call_id: call_id,
audio: request.audio,
api_key: request.api_key.map(str::to_string),
api_base: request.api_base.map(str::to_string),
extra_headers: request.extra_headers,
optional_params: request.optional_params,
timeout: request.timeout,
},
hooks: AudioTranscriptionLifecycleHooks::new(
CustomLoggerRunner::new(request.callbacks),
CustomGuardrailRunner::new(request.guardrails),
request.request_metadata,
),
}
}
fn new_audio_transcription_call_id() -> String {
static COUNTER: AtomicU64 = AtomicU64::new(1);
let sequence = COUNTER.fetch_add(1, Ordering::Relaxed);
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |duration| duration.as_nanos());
format!("audio-transcription-{timestamp}-{sequence}")
}

53
litellm-rust/crates/ai-gateway/src/audio_transcription/tests.rs
View file

@ -1,53 +0,0 @@
use std::io::{Read, Write};
use std::net::TcpListener;
use std::thread;
use serde_json::{Map, json};
use super::{AudioTranscriptionRequest, audio_transcription};
#[tokio::test]
async fn bedrock_request_is_signed_and_contains_audio() {
let listener = TcpListener::bind("127.0.0.1:0").expect("listener");
let address = listener.local_addr().expect("address");
let server = thread::spawn(move || {
let (mut stream, _) = listener.accept().expect("connection");
let mut request = Vec::new();
let mut buffer = [0_u8; 16_384];
let count = stream.read(&mut buffer).expect("request");
request.extend_from_slice(&buffer[..count]);
let request = String::from_utf8_lossy(&request);
assert!(request.contains("POST /model/mistral.voxtral-mini-3b-2507/converse"));
assert!(request.contains("authorization: AWS4-HMAC-SHA256"));
assert!(request.contains("x-amz-date:"));
assert!(request.contains("\"bytes\":\"AQI=\""));
assert!(request.contains("Transcribe the audio. Respond with only the transcript."));
let response = b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 53\r\nConnection: close\r\n\r\n{\"output\":{\"message\":{\"content\":[{\"text\":\"hello\"}]}}}";
stream.write_all(response).expect("response");
});
let optional_params = Map::from_iter([
("aws_access_key_id".to_string(), json!("access-key")),
("aws_secret_access_key".to_string(), json!("secret-key")),
("aws_region_name".to_string(), json!("us-east-1")),
]);
let api_base = format!("http://{address}");
let response = audio_transcription(AudioTranscriptionRequest {
model: "mistral.voxtral-mini-3b-2507",
audio: json!({"data": "AQI=", "format": "wav", "filename": "audio.wav"}),
api_key: None,
api_base: Some(&api_base),
custom_llm_provider: Some("bedrock"),
extra_headers: None,
optional_params,
timeout: None,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: Default::default(),
litellm_call_id: None,
})
.await
.expect("transcription");
assert_eq!(response, json!({"text": "hello"}));
server.join().expect("server");
}

47
litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs
View file

@ -1,47 +0,0 @@
use std::sync::Arc;
use std::time::Duration;
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest};
use serde_json::{Map, Value};
use crate::integrations::custom_guardrail::CustomGuardrail;
use crate::integrations::custom_logger::CustomLogger;
use crate::integrations::types::RequestMetadata;
pub struct AudioTranscriptionRequest<'a> {
pub model: &'a str,
pub audio: Value,
pub api_key: Option<&'a str>,
pub api_base: Option<&'a str>,
pub custom_llm_provider: Option<&'a str>,
pub extra_headers: Option<Map<String, Value>>,
pub optional_params: Map<String, Value>,
pub timeout: Option<Duration>,
pub callbacks: Vec<Arc<dyn CustomLogger>>,
pub guardrails: Vec<Arc<dyn CustomGuardrail>>,
pub request_metadata: RequestMetadata,
pub litellm_call_id: Option<&'a str>,
}
pub(crate) struct PreparedAudioTranscriptionRequest {
pub(crate) model: String,
pub(crate) custom_llm_provider: String,
pub(crate) litellm_call_id: String,
pub(crate) audio: Value,
pub(crate) api_key: Option<String>,
pub(crate) api_base: Option<String>,
pub(crate) extra_headers: Option<Map<String, Value>>,
pub(crate) optional_params: Map<String, Value>,
pub(crate) timeout: Option<Duration>,
}
impl CallLifecycleRequest for PreparedAudioTranscriptionRequest {
fn lifecycle_context(&self) -> CallLifecycleContext {
CallLifecycleContext::new(
"audio_transcription",
self.model.clone(),
self.custom_llm_provider.clone(),
self.litellm_call_id.clone(),
)
}
}

93
litellm-rust/crates/ai-gateway/src/auth/mod.rs
View file

@ -1,93 +0,0 @@
//! Gateway authentication, as an axum **extractor** (the idiomatic pattern —
//! keeps handlers clean and auth testable).
//!
//! For now this is a single **master key**: any caller presenting it as
//! `Authorization: Bearer <key>` may invoke the gateway. Per-key auth, budgets,
//! and rate limits are delegated to the Python proxy in a later phase.
//!
//! A handler opts in by adding [`RequireMasterKey`] to its arguments; auth then
//! runs during extraction, before the handler body. Routes never re-implement it.
use axum::extract::FromRequestParts;
use axum::http::StatusCode;
use axum::http::header::AUTHORIZATION;
use axum::http::request::Parts;
use sha2::{Digest, Sha256};
use subtle::ConstantTimeEq;
use crate::state::AppState;
/// SHA-256 hex digest of a token — the exact transform the Python proxy applies
/// (`litellm.proxy.utils.hash_token`).
///
/// STRICT REQUIREMENT: a raw key (`LITELLM_MASTER_KEY`, a virtual key, …) must
/// **never** leave this gateway in a log payload. Spend logs and every callback
/// integration receive `user_api_key_hash`, so that field must be this hash, not
/// the credential. Hashing here also means the value matches the key's hash in
/// `LiteLLM_SpendLogs.api_key`, so realtime spend joins with the rest of LiteLLM.
pub fn hash_token(token: &str) -> String {
let digest = Sha256::digest(token.as_bytes());
let mut hex = String::with_capacity(digest.len() * 2);
for byte in digest {
use std::fmt::Write;
let _ = write!(hex, "{byte:02x}");
}
hex
}
/// Extractor that requires the configured master key as a bearer token.
///
/// Rejections: `500` when no master key is configured (permanent
/// misconfiguration, not a transient outage); `401` on a missing/incorrect
/// token. The comparison is constant-time.
pub struct RequireMasterKey;
#[axum::async_trait]
impl FromRequestParts<AppState> for RequireMasterKey {
type Rejection = (StatusCode, String);
async fn from_request_parts(
parts: &mut Parts,
state: &AppState,
) -> Result<Self, Self::Rejection> {
let Some(expected) = state.master_key.as_deref() else {
return Err((
StatusCode::INTERNAL_SERVER_ERROR,
"gateway auth not configured (set LITELLM_MASTER_KEY)".to_string(),
));
};
let provided = parts
.headers
.get(AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.strip_prefix("Bearer "))
.map(str::trim);
match provided {
Some(token) if bool::from(token.as_bytes().ct_eq(expected.as_bytes())) => Ok(Self),
_ => Err((
StatusCode::UNAUTHORIZED,
"missing or invalid bearer token".to_string(),
)),
}
}
}
#[cfg(test)]
mod tests {
use super::hash_token;
#[test]
fn hash_token_matches_python_sha256_hexdigest() {
// Must equal hashlib.sha256("sk-1234".encode()).hexdigest() — the value
// the proxy stores in LiteLLM_SpendLogs.api_key.
assert_eq!(
hash_token("sk-1234"),
"88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"
);
// 64 lowercase hex chars, and never the raw input.
let h = hash_token("sk-secret");
assert_eq!(h.len(), 64);
assert!(h.chars().all(|c| c.is_ascii_hexdigit()));
assert_ne!(h, "sk-secret");
}
}

42
litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs
View file

@ -1,42 +0,0 @@
use std::io::Read;
use serde::Deserialize;
use serde_json::Value;
#[derive(Deserialize)]
struct Input {
path: String,
model_alias: String,
provider_model: String,
api_base: String,
body: Value,
}
#[tokio::main]
async fn main() {
let mut input = String::new();
if let Err(error) = std::io::stdin().read_to_string(&mut input) {
fail(error);
}
let input: Input = match serde_json::from_str(&input) {
Ok(input) => input,
Err(error) => fail(error),
};
let result = litellm_ai_gateway::trace_parity::traced_request(
input.path,
input.model_alias,
input.provider_model,
input.api_base,
input.body,
)
.await;
match serde_json::to_string(&result) {
Ok(result) => println!("{result}"),
Err(error) => fail(error),
}
}
fn fail(error: impl std::fmt::Display) -> ! {
eprintln!("{error}");
std::process::exit(1)
}

14
litellm-rust/crates/ai-gateway/src/client.rs
View file

@ -1,14 +0,0 @@
use std::sync::OnceLock;
use std::time::Duration;
const HTTP_CLIENT_TIMEOUT_SECS: u64 = 600;
pub(crate) fn http_client() -> &'static reqwest::Client {
static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
CLIENT.get_or_init(|| {
reqwest::Client::builder()
.timeout(Duration::from_secs(HTTP_CLIENT_TIMEOUT_SECS))
.build()
.expect("failed to build reqwest client")
})
}

42
litellm-rust/crates/ai-gateway/src/constants.rs
View file

@ -1,42 +0,0 @@
//! Crate-level constants for the ai-gateway.
//!
//! Per `litellm-rust/CLAUDE.md`, magic numbers and fixed strings live here
//! (the Rust mirror of Python's `litellm/constants.py`), not inline in feature
//! modules. Env-overridable tunables keep their `DEFAULT_*` value here; the env
//! read + fallback happens at the host/config layer.
/// Default LiteLLM control-plane base URL for request-log egress when
/// `LITELLM_PROXY_BASE_URL` is unset.
pub(crate) const DEFAULT_PROXY_BASE_URL: &str = "http://localhost:4000";
/// The logs ingest path appended to the proxy base. Not a tunable; it is the
/// proxy's API contract (the rust-control-plane router on the Python proxy).
pub(crate) const RUST_CONTROL_PLANE_LOGS_PATH: &str = "/v1/rust_control_plane/logs";
/// Default bounded channel depth for the log-egress worker.
/// Override: `LITELLM_LOG_CHANNEL_CAPACITY`.
pub(crate) const DEFAULT_CHANNEL_CAPACITY: usize = 4096;
/// Default max records POSTed per request to the control plane.
/// Override: `LITELLM_LOG_BATCH_SIZE`.
pub(crate) const DEFAULT_MAX_BATCH_SIZE: usize = 256;
/// Default partial-batch flush cadence, in ms.
/// Override: `LITELLM_LOG_FLUSH_INTERVAL_MS`.
pub(crate) const DEFAULT_FLUSH_INTERVAL_MS: u64 = 500;
/// Provider attributed to realtime sessions in the logging payload.
#[cfg(feature = "server")]
pub(crate) const DEFAULT_PROVIDER: &str = "openai";
pub(crate) const DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS: u64 = 10;
pub(crate) const DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS: u64 = 300;
/// HTTP path for the non-streaming Anthropic Messages route.
#[cfg(feature = "server")]
pub(crate) const MESSAGES_ROUTE_PATH: &str = "/v1/messages";
/// Request headers owned by the gateway and never forwarded upstream.
#[cfg(feature = "server")]
pub(crate) const MESSAGES_HEADERS_NOT_FORWARDED: &[&str] =
&["authorization", "connection", "content-length", "host"];

127
litellm-rust/crates/ai-gateway/src/integrations/README.md
View file

@ -1,127 +0,0 @@
# LiteLLM Rust integrations
This directory contains Rust-native equivalents of LiteLLM integration hooks.
The first supported surfaces are terminal custom loggers and pre/during-call
custom guardrails.
## File layout
Every integration is a folder:
- `mod.rs` contains the implementation, trait, runner, or adapter
- `types.rs` contains the integration-local request, response, error, and future
types
Do not add new flat integration files such as `custom_logger.rs`. Shared wire
contracts that are used by multiple integrations can stay in
`integrations/types.rs`.
Call ordering and lifecycle timing live in `litellm-core/src/call_lifecycle`.
Call-type modules, such as OCR, adapt their request and response shapes into
that generic lifecycle runner.
## CustomLogger
Implement `CustomLogger` when Rust code needs to observe terminal success or
failure events. Method names intentionally match Python `CustomLogger` names.
```rust
use litellm_ai_gateway::integrations::custom_logger::{
CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails,
};
struct RecordingLogger;
impl CustomLogger for RecordingLogger {
fn async_log_success_event<'a>(
&'a self,
model_call_details: &'a ModelCallDetails,
response_obj: &'a CallbackValue,
timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async move {
let model = &model_call_details.model;
let provider = &model_call_details.custom_llm_provider;
let call_type = model_call_details.call_type.to_string();
let request_id = model_call_details.request_id.as_deref();
let response_object = &response_obj.object;
let duration = timing.end_time - timing.start_time;
let standard_payload = model_call_details.standard_logging_payload.as_ref();
Ok(())
})
}
fn async_log_failure_event<'a>(
&'a self,
model_call_details: &'a ModelCallDetails,
response_obj: Option<&'a CallbackValue>,
timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async move {
let error = model_call_details.failure_error.as_ref();
let response_object = response_obj.map(|value| value.object.as_str());
let duration = timing.end_time - timing.start_time;
Ok(())
})
}
}
```
Use `CustomLoggerRunner` to fan out terminal events to configured loggers. The
runner is a no-op when no loggers are configured, which is the expected fast
path for requests without callbacks.
## CustomGuardrail
Implement `CustomGuardrail` when Rust code needs to run pre-call or native
during-call checks. Method names intentionally match Python `CustomGuardrail`
entrypoints inherited from Python `CustomLogger`.
```rust
use litellm_ai_gateway::integrations::custom_guardrail::{
CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailEventHook,
GuardrailFuture, GuardrailRequest,
};
struct BlocklistedPromptGuardrail;
impl CustomGuardrail for BlocklistedPromptGuardrail {
fn guardrail_name(&self) -> &str {
"blocklisted-prompt"
}
fn supported_event_hooks(&self) -> &[GuardrailEventHook] {
&[GuardrailEventHook::PreCall]
}
fn async_pre_call_hook<'a>(
&'a self,
_context: &'a GuardrailContext,
request: GuardrailRequest,
) -> GuardrailFuture<'a> {
Box::pin(async move {
if request.data.to_string().contains("blocked phrase") {
return Ok(GuardrailDecision::Block(
litellm_ai_gateway::integrations::custom_guardrail::GuardrailError::blocked(
"blocked phrase detected",
),
));
}
Ok(GuardrailDecision::Allow(request))
})
}
}
```
Use `CustomGuardrailRunner::run_pre_call` for `pre_call` guardrails and
`CustomGuardrailRunner::run_during_call` for `during_call` guardrails. A
`GuardrailDecision::Mask` continues with modified request data.
`GuardrailDecision::Block` short-circuits the provider call.
## Current boundary
These are Rust-only primitives. Python callback and guardrail adapters are a
separate layer that should implement these Rust traits instead of changing the
runner interfaces.

468
litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/mod.rs
View file

@ -1,468 +0,0 @@
//! Rust mirror of Python `CustomGuardrail` entrypoints used by the proxy.
//!
//! This module is intentionally Rust-only: Python/PyO3 adapters are a later
//! layer that should implement this trait rather than changing the runner.
use std::future::Future;
use std::sync::Arc;
use crate::integrations::custom_logger::{
CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails,
};
pub mod types;
pub use types::{
GuardrailContext, GuardrailDecision, GuardrailDispatchReport, GuardrailError,
GuardrailEventHook, GuardrailFuture, GuardrailRequest,
};
pub trait CustomGuardrail: Send + Sync {
fn guardrail_name(&self) -> &str;
fn supported_event_hooks(&self) -> &[GuardrailEventHook];
/// Python 1:1 name: `async_pre_call_hook(user_api_key_dict, cache, data, call_type)`.
fn async_pre_call_hook<'a>(
&'a self,
_context: &'a GuardrailContext,
request: GuardrailRequest,
) -> GuardrailFuture<'a> {
Box::pin(async move { Ok(GuardrailDecision::Allow(request)) })
}
/// Python 1:1 name: `async_moderation_hook(data, user_api_key_dict, call_type)`.
fn async_moderation_hook<'a>(
&'a self,
_context: &'a GuardrailContext,
request: GuardrailRequest,
) -> GuardrailFuture<'a> {
Box::pin(async move { Ok(GuardrailDecision::Allow(request)) })
}
}
pub struct CustomGuardrailRunner {
guardrails: Vec<Arc<dyn CustomGuardrail>>,
}
impl CustomGuardrailRunner {
pub fn new(guardrails: Vec<Arc<dyn CustomGuardrail>>) -> Self {
Self { guardrails }
}
pub fn is_empty(&self) -> bool {
self.guardrails.is_empty()
}
pub async fn run_pre_call(
&self,
context: &GuardrailContext,
request: GuardrailRequest,
) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> {
self.run_hook(GuardrailEventHook::PreCall, context, request)
.await
}
pub async fn run_during_call(
&self,
context: &GuardrailContext,
request: GuardrailRequest,
) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> {
self.run_hook(GuardrailEventHook::DuringCall, context, request)
.await
}
pub async fn run_before_provider<F, Fut, T>(
&self,
event_hook: GuardrailEventHook,
context: &GuardrailContext,
request: GuardrailRequest,
provider: F,
) -> Result<T, GuardrailError>
where
F: FnOnce(GuardrailRequest) -> Fut,
Fut: Future<Output = Result<T, GuardrailError>>,
{
let (request, _) = self.run_hook(event_hook, context, request).await?;
provider(request).await
}
pub async fn run_pre_call_with_failure_logging(
&self,
context: &GuardrailContext,
request: GuardrailRequest,
logger_runner: &CustomLoggerRunner,
model_call_details: &ModelCallDetails,
timing: CallbackTiming,
) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> {
match self.run_pre_call(context, request).await {
Ok(result) => Ok(result),
Err(error) => {
let failure_details = model_call_details.clone().with_failure_error(LoggingError {
message: error.message.clone(),
kind: error.kind.clone(),
});
let response_obj = CallbackValue::new(
"guardrail_error",
serde_json::json!({
"message": error.message,
"kind": error.kind,
}),
);
logger_runner
.async_log_failure_event(&failure_details, Some(&response_obj), timing)
.await;
Err(error)
}
}
}
async fn run_hook(
&self,
event_hook: GuardrailEventHook,
context: &GuardrailContext,
mut request: GuardrailRequest,
) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> {
if self.guardrails.is_empty() {
return Ok((request, GuardrailDispatchReport::default()));
}
let mut report = GuardrailDispatchReport::default();
for guardrail in &self.guardrails {
if !self.should_run(guardrail.as_ref(), event_hook, context) {
continue;
}
report.invoked += 1;
let decision = match event_hook {
GuardrailEventHook::PreCall => {
guardrail
.async_pre_call_hook(context, request.clone())
.await?
}
GuardrailEventHook::DuringCall => {
guardrail
.async_moderation_hook(context, request.clone())
.await?
}
};
match decision.into_request() {
Ok(next_request) => request = next_request,
Err(error) => return Err(error),
}
}
Ok((request, report))
}
fn should_run(
&self,
guardrail: &dyn CustomGuardrail,
event_hook: GuardrailEventHook,
context: &GuardrailContext,
) -> bool {
let supports_hook = guardrail.supported_event_hooks().contains(&event_hook);
let selected = context.selected_guardrails.is_empty()
|| context
.selected_guardrails
.iter()
.any(|name| name == guardrail.guardrail_name());
supports_hook && selected
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::integrations::custom_logger::{CallType, CallbackValue, CustomLogger, LogFuture};
use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload};
use serde_json::json;
use std::sync::Mutex;
#[derive(Clone)]
enum TestDecision {
Allow,
Mask,
Block,
}
struct RecordingCustomGuardrail {
name: String,
hooks: Vec<GuardrailEventHook>,
decision: TestDecision,
calls: Mutex<Vec<&'static str>>,
}
impl RecordingCustomGuardrail {
fn new(name: &str, hooks: Vec<GuardrailEventHook>, decision: TestDecision) -> Self {
Self {
name: name.to_string(),
hooks,
decision,
calls: Mutex::new(Vec::new()),
}
}
fn calls(&self) -> Vec<&'static str> {
self.calls.lock().unwrap().clone()
}
fn decision(&self, mut request: GuardrailRequest) -> GuardrailDecision {
match self.decision {
TestDecision::Allow => GuardrailDecision::Allow(request),
TestDecision::Mask => {
request.data["masked"] = json!(true);
GuardrailDecision::Mask(request)
}
TestDecision::Block => {
GuardrailDecision::Block(GuardrailError::blocked("blocked by guardrail"))
}
}
}
}
impl CustomGuardrail for RecordingCustomGuardrail {
fn guardrail_name(&self) -> &str {
&self.name
}
fn supported_event_hooks(&self) -> &[GuardrailEventHook] {
&self.hooks
}
fn async_pre_call_hook<'a>(
&'a self,
_context: &'a GuardrailContext,
request: GuardrailRequest,
) -> GuardrailFuture<'a> {
Box::pin(async move {
self.calls.lock().unwrap().push("async_pre_call_hook");
Ok(self.decision(request))
})
}
fn async_moderation_hook<'a>(
&'a self,
_context: &'a GuardrailContext,
request: GuardrailRequest,
) -> GuardrailFuture<'a> {
Box::pin(async move {
self.calls.lock().unwrap().push("async_moderation_hook");
Ok(self.decision(request))
})
}
}
#[tokio::test]
async fn pre_call_dispatches_to_async_pre_call_hook() {
let guardrail = Arc::new(RecordingCustomGuardrail::new(
"pre",
vec![GuardrailEventHook::PreCall],
TestDecision::Allow,
));
let runner = CustomGuardrailRunner::new(vec![guardrail.clone()]);
let context =
GuardrailContext::new(CallType::Ocr).with_selected_guardrails(vec!["pre".to_string()]);
let request = GuardrailRequest::new(json!({"messages": ["hello"]}));
let (result, report) = runner
.run_pre_call(&context, request)
.await
.expect("guardrail allows request");
assert_eq!(report.invoked, 1);
assert_eq!(result.data["messages"], json!(["hello"]));
assert_eq!(guardrail.calls(), vec!["async_pre_call_hook"]);
}
#[tokio::test]
async fn during_call_dispatches_to_async_moderation_hook() {
let guardrail = Arc::new(RecordingCustomGuardrail::new(
"during",
vec![GuardrailEventHook::DuringCall],
TestDecision::Allow,
));
let runner = CustomGuardrailRunner::new(vec![guardrail.clone()]);
let context = GuardrailContext::new(CallType::Completion)
.with_selected_guardrails(vec!["during".to_string()]);
let request = GuardrailRequest::new(json!({"prompt": "hello"}));
let (_result, report) = runner
.run_during_call(&context, request)
.await
.expect("guardrail allows request");
assert_eq!(report.invoked, 1);
assert_eq!(guardrail.calls(), vec!["async_moderation_hook"]);
}
#[tokio::test]
async fn mask_decision_continues_with_updated_request() {
let guardrail = Arc::new(RecordingCustomGuardrail::new(
"masker",
vec![GuardrailEventHook::PreCall],
TestDecision::Mask,
));
let runner = CustomGuardrailRunner::new(vec![guardrail]);
let context = GuardrailContext::new(CallType::Ocr);
let request = GuardrailRequest::new(json!({"document": "secret"}));
let (result, report) = runner
.run_pre_call(&context, request)
.await
.expect("mask continues");
assert_eq!(report.invoked, 1);
assert_eq!(result.data["masked"], json!(true));
}
#[tokio::test]
async fn block_decision_short_circuits_and_logs_failure() {
struct RecordingFailureLogger {
errors: Mutex<Vec<String>>,
}
impl CustomLogger for RecordingFailureLogger {
fn async_log_failure_event<'a>(
&'a self,
model_call_details: &'a ModelCallDetails,
_response_obj: Option<&'a CallbackValue>,
_timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async move {
self.errors.lock().unwrap().push(
model_call_details
.failure_error
.as_ref()
.map(|error| error.kind.clone())
.unwrap_or_default(),
);
Ok(())
})
}
}
let guardrail = Arc::new(RecordingCustomGuardrail::new(
"blocker",
vec![GuardrailEventHook::PreCall],
TestDecision::Block,
));
let guardrail_runner = CustomGuardrailRunner::new(vec![guardrail]);
let logger = Arc::new(RecordingFailureLogger {
errors: Mutex::new(Vec::new()),
});
let logger_runner = CustomLoggerRunner::new(vec![logger.clone()]);
let context = GuardrailContext::new(CallType::Ocr);
let details = ModelCallDetails::from_standard_logging_payload(StandardLoggingPayload {
id: "req_ocr".to_string(),
litellm_call_id: "req_ocr".to_string(),
call_type: "ocr".to_string(),
model: "mistral-ocr-latest".to_string(),
custom_llm_provider: "mistral".to_string(),
response_cost: 0.0,
prompt_tokens: 0,
completion_tokens: 0,
total_tokens: 0,
start_time: 1.0,
end_time: 1.0,
stream: false,
metadata: StandardLoggingMetadata::default(),
messages: None,
});
let err = guardrail_runner
.run_pre_call_with_failure_logging(
&context,
GuardrailRequest::new(json!({"document": "bad"})),
&logger_runner,
&details,
CallbackTiming::new(1.0, 2.0),
)
.await
.expect_err("guardrail blocks request");
assert_eq!(err.kind, "GuardrailBlocked");
assert_eq!(
logger.errors.lock().unwrap().as_slice(),
["GuardrailBlocked"]
);
}
#[tokio::test]
async fn block_decision_short_circuits_later_guardrails_and_provider_work() {
let blocking_guardrail = Arc::new(RecordingCustomGuardrail::new(
"blocker",
vec![GuardrailEventHook::PreCall],
TestDecision::Block,
));
let later_guardrail = Arc::new(RecordingCustomGuardrail::new(
"later",
vec![GuardrailEventHook::PreCall],
TestDecision::Allow,
));
let runner =
CustomGuardrailRunner::new(vec![blocking_guardrail.clone(), later_guardrail.clone()]);
let provider_called = Arc::new(Mutex::new(false));
let provider_called_for_closure = provider_called.clone();
let result = runner
.run_before_provider(
GuardrailEventHook::PreCall,
&GuardrailContext::new(CallType::Completion),
GuardrailRequest::new(json!({"prompt": "blocked"})),
move |_request| async move {
*provider_called_for_closure.lock().unwrap() = true;
Ok("provider response")
},
)
.await;
assert!(result.is_err());
assert_eq!(blocking_guardrail.calls(), vec!["async_pre_call_hook"]);
assert_eq!(later_guardrail.calls(), Vec::<&'static str>::new());
assert!(!*provider_called.lock().unwrap());
}
#[tokio::test]
async fn run_before_provider_returns_provider_guardrail_error_directly() {
let guardrail = Arc::new(RecordingCustomGuardrail::new(
"allow",
vec![GuardrailEventHook::PreCall],
TestDecision::Allow,
));
let runner = CustomGuardrailRunner::new(vec![guardrail]);
let result = runner
.run_before_provider(
GuardrailEventHook::PreCall,
&GuardrailContext::new(CallType::Completion),
GuardrailRequest::new(json!({"prompt": "allowed"})),
|_request| async move {
Err::<&'static str, GuardrailError>(GuardrailError::blocked(
"provider-side guardrail error",
))
},
)
.await;
let err = result.expect_err("provider error is returned directly");
assert_eq!(err.kind, "GuardrailBlocked");
assert_eq!(err.message, "provider-side guardrail error");
}
#[tokio::test]
async fn no_guardrails_fast_path_dispatches_nothing() {
let runner = CustomGuardrailRunner::new(Vec::new());
let context = GuardrailContext::new(CallType::Ocr);
let request = GuardrailRequest::new(json!({"document": "ok"}));
let (result, report) = runner
.run_pre_call(&context, request)
.await
.expect("no guardrails allow request");
assert!(runner.is_empty());
assert_eq!(report, GuardrailDispatchReport::default());
assert_eq!(result.data["document"], json!("ok"));
}
}

110
litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/types.rs
View file

@ -1,110 +0,0 @@
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use serde_json::Value;
use crate::integrations::custom_logger::CallType;
pub type GuardrailFuture<'a> =
Pin<Box<dyn Future<Output = Result<GuardrailDecision, GuardrailError>> + Send + 'a>>;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum GuardrailEventHook {
PreCall,
DuringCall,
}
impl GuardrailEventHook {
pub fn as_str(&self) -> &'static str {
match self {
Self::PreCall => "pre_call",
Self::DuringCall => "during_call",
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct GuardrailError {
pub message: String,
pub kind: String,
}
impl GuardrailError {
pub fn blocked(message: impl Into<String>) -> Self {
Self {
message: message.into(),
kind: "GuardrailBlocked".to_string(),
}
}
}
impl std::fmt::Display for GuardrailError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}: {}", self.kind, self.message)
}
}
impl std::error::Error for GuardrailError {}
#[derive(Clone, Debug)]
pub struct GuardrailContext {
pub call_type: CallType,
pub selected_guardrails: Vec<String>,
pub metadata: HashMap<String, Value>,
pub user_api_key_hash: Option<String>,
pub user_api_key_user_id: Option<String>,
pub user_api_key_team_id: Option<String>,
pub trace_parent: Option<String>,
}
impl GuardrailContext {
pub fn new(call_type: CallType) -> Self {
Self {
call_type,
selected_guardrails: Vec::new(),
metadata: HashMap::new(),
user_api_key_hash: None,
user_api_key_user_id: None,
user_api_key_team_id: None,
trace_parent: None,
}
}
pub fn with_selected_guardrails(mut self, selected_guardrails: Vec<String>) -> Self {
self.selected_guardrails = selected_guardrails;
self
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct GuardrailRequest {
pub data: Value,
}
impl GuardrailRequest {
pub fn new(data: Value) -> Self {
Self { data }
}
}
#[derive(Clone, Debug, PartialEq)]
pub enum GuardrailDecision {
Allow(GuardrailRequest),
Mask(GuardrailRequest),
Block(GuardrailError),
}
impl GuardrailDecision {
pub(super) fn into_request(self) -> Result<GuardrailRequest, GuardrailError> {
match self {
Self::Allow(request) | Self::Mask(request) => Ok(request),
Self::Block(error) => Err(error),
}
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct GuardrailDispatchReport {
pub invoked: usize,
}

317
litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs
View file

@ -1,317 +0,0 @@
//! The `CustomLogger` trait — the Rust mirror of Python
//! `litellm/integrations/custom_logger.py::CustomLogger`.
//!
//! The Python-named async terminal methods are the public Rust callback shape.
use std::sync::Arc;
pub mod types;
pub use types::{
CallType, CallbackDispatchReport, CallbackTiming, CallbackValue, LogError, LogFuture,
LoggingError, ModelCallDetails,
};
pub trait CustomLogger: Send + Sync {
/// Python 1:1 name: `async_log_success_event(model_call_details, response_obj, start_time, end_time)`.
fn async_log_success_event<'a>(
&'a self,
_model_call_details: &'a ModelCallDetails,
_response_obj: &'a CallbackValue,
_timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async { Ok(()) })
}
/// Python 1:1 name: `async_log_failure_event(model_call_details, response_obj, start_time, end_time)`.
fn async_log_failure_event<'a>(
&'a self,
_model_call_details: &'a ModelCallDetails,
_response_obj: Option<&'a CallbackValue>,
_timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async { Ok(()) })
}
}
pub struct CustomLoggerRunner {
loggers: Vec<Arc<dyn CustomLogger>>,
}
impl CustomLoggerRunner {
pub fn new(loggers: Vec<Arc<dyn CustomLogger>>) -> Self {
Self { loggers }
}
pub fn is_empty(&self) -> bool {
self.loggers.is_empty()
}
pub async fn async_log_success_event(
&self,
model_call_details: &ModelCallDetails,
response_obj: &CallbackValue,
timing: CallbackTiming,
) -> CallbackDispatchReport {
if self.loggers.is_empty() {
return CallbackDispatchReport::default();
}
let mut report = CallbackDispatchReport::default();
for logger in &self.loggers {
report.invoked += 1;
if let Err(err) = logger
.async_log_success_event(model_call_details, response_obj, timing)
.await
{
report.dropped += 1;
eprintln!("litellm-ai-gateway: async_log_success_event dropped: {err}");
}
}
report
}
pub async fn async_log_failure_event(
&self,
model_call_details: &ModelCallDetails,
response_obj: Option<&CallbackValue>,
timing: CallbackTiming,
) -> CallbackDispatchReport {
if self.loggers.is_empty() {
return CallbackDispatchReport::default();
}
let mut report = CallbackDispatchReport::default();
for logger in &self.loggers {
report.invoked += 1;
if let Err(err) = logger
.async_log_failure_event(model_call_details, response_obj, timing)
.await
{
report.dropped += 1;
eprintln!("litellm-ai-gateway: async_log_failure_event dropped: {err}");
}
}
report
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload};
use serde_json::json;
use std::sync::Mutex;
#[derive(Clone, Debug, PartialEq)]
struct RecordedEvent {
hook: &'static str,
model: String,
provider: String,
call_type: String,
request_id: Option<String>,
litellm_call_id: Option<String>,
user_id: Option<String>,
response_object: Option<String>,
error_kind: Option<String>,
start_time: f64,
end_time: f64,
standard_logging_model: Option<String>,
}
#[derive(Default)]
struct RecordingCustomLogger {
events: Mutex<Vec<RecordedEvent>>,
}
impl RecordingCustomLogger {
fn events(&self) -> Vec<RecordedEvent> {
self.events.lock().unwrap().clone()
}
}
impl CustomLogger for RecordingCustomLogger {
fn async_log_success_event<'a>(
&'a self,
model_call_details: &'a ModelCallDetails,
response_obj: &'a CallbackValue,
timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push(RecordedEvent {
hook: "async_log_success_event",
model: model_call_details.model.clone(),
provider: model_call_details.custom_llm_provider.clone(),
call_type: model_call_details.call_type.to_string(),
request_id: model_call_details.request_id.clone(),
litellm_call_id: model_call_details.litellm_call_id.clone(),
user_id: model_call_details.metadata.user_api_key_user_id.clone(),
response_object: Some(response_obj.object.clone()),
error_kind: None,
start_time: timing.start_time,
end_time: timing.end_time,
standard_logging_model: model_call_details
.standard_logging_payload
.as_ref()
.map(|payload| payload.model.clone()),
});
Ok(())
})
}
fn async_log_failure_event<'a>(
&'a self,
model_call_details: &'a ModelCallDetails,
response_obj: Option<&'a CallbackValue>,
timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push(RecordedEvent {
hook: "async_log_failure_event",
model: model_call_details.model.clone(),
provider: model_call_details.custom_llm_provider.clone(),
call_type: model_call_details.call_type.to_string(),
request_id: model_call_details.request_id.clone(),
litellm_call_id: model_call_details.litellm_call_id.clone(),
user_id: model_call_details.metadata.user_api_key_user_id.clone(),
response_object: response_obj.map(|value| value.object.clone()),
error_kind: model_call_details
.failure_error
.as_ref()
.map(|error| error.kind.clone()),
start_time: timing.start_time,
end_time: timing.end_time,
standard_logging_model: model_call_details
.standard_logging_payload
.as_ref()
.map(|payload| payload.model.clone()),
});
Ok(())
})
}
}
fn payload(call_type: &str, model: &str, provider: &str) -> StandardLoggingPayload {
StandardLoggingPayload {
id: format!("req_{call_type}"),
litellm_call_id: format!("call_{call_type}"),
call_type: call_type.to_string(),
model: model.to_string(),
custom_llm_provider: provider.to_string(),
response_cost: 0.25,
prompt_tokens: 3,
completion_tokens: 4,
total_tokens: 7,
start_time: 10.0,
end_time: 11.5,
stream: false,
metadata: StandardLoggingMetadata {
user_api_key_hash: Some("hash".to_string()),
user_api_key_user_id: Some("user".to_string()),
user_api_key_team_id: Some("team".to_string()),
..Default::default()
},
messages: Some(json!([{"role": "user", "content": "read this"}])),
}
}
#[tokio::test]
async fn rust_custom_logger_reads_success_payload_for_ocr() {
let logger = Arc::new(RecordingCustomLogger::default());
let runner = CustomLoggerRunner::new(vec![logger.clone()]);
let details = ModelCallDetails::from_standard_logging_payload(payload(
"ocr",
"mistral-ocr-latest",
"mistral",
));
let response = CallbackValue::new("ocr", json!({"pages": [{"markdown": "ok"}]}));
let report = runner
.async_log_success_event(&details, &response, CallbackTiming::new(10.0, 11.5))
.await;
assert_eq!(report.invoked, 1);
assert_eq!(report.dropped, 0);
assert_eq!(
logger.events(),
vec![RecordedEvent {
hook: "async_log_success_event",
model: "mistral-ocr-latest".to_string(),
provider: "mistral".to_string(),
call_type: "ocr".to_string(),
request_id: Some("req_ocr".to_string()),
litellm_call_id: Some("call_ocr".to_string()),
user_id: Some("user".to_string()),
response_object: Some("ocr".to_string()),
error_kind: None,
start_time: 10.0,
end_time: 11.5,
standard_logging_model: Some("mistral-ocr-latest".to_string()),
}]
);
}
#[tokio::test]
async fn rust_custom_logger_reads_failure_payload_for_non_ocr_call_type() {
let logger = Arc::new(RecordingCustomLogger::default());
let runner = CustomLoggerRunner::new(vec![logger.clone()]);
let details = ModelCallDetails::from_standard_logging_payload(payload(
"acompletion",
"gpt-4.1-mini",
"openai",
))
.with_failure_error(LoggingError {
message: "provider failed".to_string(),
kind: "ProviderError".to_string(),
});
let response = CallbackValue::new("error", json!({"message": "provider failed"}));
let report = runner
.async_log_failure_event(&details, Some(&response), CallbackTiming::new(2.0, 3.0))
.await;
assert_eq!(report.invoked, 1);
assert_eq!(report.dropped, 0);
assert_eq!(
logger.events(),
vec![RecordedEvent {
hook: "async_log_failure_event",
model: "gpt-4.1-mini".to_string(),
provider: "openai".to_string(),
call_type: "acompletion".to_string(),
request_id: Some("req_acompletion".to_string()),
litellm_call_id: Some("call_acompletion".to_string()),
user_id: Some("user".to_string()),
response_object: Some("error".to_string()),
error_kind: Some("ProviderError".to_string()),
start_time: 2.0,
end_time: 3.0,
standard_logging_model: Some("gpt-4.1-mini".to_string()),
}]
);
}
#[tokio::test]
async fn no_callback_fast_path_dispatches_nothing() {
let runner = CustomLoggerRunner::new(Vec::new());
let details = ModelCallDetails::new("mistral-ocr-latest", "mistral", CallType::Ocr);
let response = CallbackValue::new("ocr", json!({}));
let report = runner
.async_log_success_event(&details, &response, CallbackTiming::new(1.0, 1.5))
.await;
assert!(runner.is_empty());
assert_eq!(report, CallbackDispatchReport::default());
}
#[test]
fn with_standard_logging_payload_keeps_top_level_fields_in_sync() {
let details = ModelCallDetails::new("old-model", "old-provider", CallType::Completion)
.with_standard_logging_payload(payload("ocr", "mistral-ocr-latest", "mistral"));
assert_eq!(details.model, "mistral-ocr-latest");
assert_eq!(details.custom_llm_provider, "mistral");
assert_eq!(details.call_type, CallType::Ocr);
assert_eq!(details.request_id, Some("req_ocr".to_string()));
assert_eq!(details.litellm_call_id, Some("call_ocr".to_string()));
}
}

194
litellm-rust/crates/ai-gateway/src/integrations/custom_logger/types.rs
View file

@ -1,194 +0,0 @@
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use serde_json::Value;
use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload};
pub type LogFuture<'a> = Pin<Box<dyn Future<Output = Result<(), LogError>> + Send + 'a>>;
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct CallbackDispatchReport {
pub invoked: usize,
pub dropped: usize,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CallType {
Ocr,
Realtime,
Completion,
Acompletion,
ChatCompletion,
Other(String),
}
impl CallType {
pub fn as_str(&self) -> &str {
match self {
Self::Ocr => "ocr",
Self::Realtime => "realtime",
Self::Completion => "completion",
Self::Acompletion => "acompletion",
Self::ChatCompletion => "chat_completion",
Self::Other(value) => value.as_str(),
}
}
}
impl From<&str> for CallType {
fn from(value: &str) -> Self {
match value {
"ocr" => Self::Ocr,
"realtime" => Self::Realtime,
"completion" => Self::Completion,
"acompletion" => Self::Acompletion,
"chat_completion" => Self::ChatCompletion,
other => Self::Other(other.to_string()),
}
}
}
impl std::fmt::Display for CallType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct CallbackTiming {
pub start_time: f64,
pub end_time: f64,
}
impl CallbackTiming {
pub fn new(start_time: f64, end_time: f64) -> Self {
Self {
start_time,
end_time,
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct CallbackValue {
pub object: String,
pub value: Value,
}
impl CallbackValue {
pub fn new(object: impl Into<String>, value: Value) -> Self {
Self {
object: object.into(),
value,
}
}
}
#[derive(Clone, Debug)]
pub struct ModelCallDetails {
pub model: String,
pub custom_llm_provider: String,
pub call_type: CallType,
pub metadata: StandardLoggingMetadata,
pub extra_metadata: HashMap<String, Value>,
pub request_id: Option<String>,
pub litellm_call_id: Option<String>,
pub response_cost: Option<f64>,
pub standard_logging_payload: Option<StandardLoggingPayload>,
pub failure_error: Option<LoggingError>,
}
impl ModelCallDetails {
pub fn new(
model: impl Into<String>,
custom_llm_provider: impl Into<String>,
call_type: CallType,
) -> Self {
Self {
model: model.into(),
custom_llm_provider: custom_llm_provider.into(),
call_type,
metadata: StandardLoggingMetadata::default(),
extra_metadata: HashMap::new(),
request_id: None,
litellm_call_id: None,
response_cost: None,
standard_logging_payload: None,
failure_error: None,
}
}
pub fn from_standard_logging_payload(payload: StandardLoggingPayload) -> Self {
let request_id = Some(payload.id.clone());
let litellm_call_id = Some(payload.litellm_call_id.clone());
let response_cost = Some(payload.response_cost);
let metadata = payload.metadata.clone();
Self {
model: payload.model.clone(),
custom_llm_provider: payload.custom_llm_provider.clone(),
call_type: CallType::from(payload.call_type.as_str()),
metadata,
extra_metadata: HashMap::new(),
request_id,
litellm_call_id,
response_cost,
standard_logging_payload: Some(payload),
failure_error: None,
}
}
pub fn with_standard_logging_payload(mut self, payload: StandardLoggingPayload) -> Self {
self.model = payload.model.clone();
self.custom_llm_provider = payload.custom_llm_provider.clone();
self.call_type = CallType::from(payload.call_type.as_str());
self.request_id = Some(payload.id.clone());
self.litellm_call_id = Some(payload.litellm_call_id.clone());
self.response_cost = Some(payload.response_cost);
self.metadata = payload.metadata.clone();
self.standard_logging_payload = Some(payload);
self
}
pub fn with_failure_error(mut self, error: LoggingError) -> Self {
self.failure_error = Some(error);
self
}
}
#[derive(Clone, Debug)]
pub struct LoggingError {
pub message: String,
pub kind: String,
}
#[derive(Clone, Debug)]
pub struct LogError {
pub message: String,
pub kind: String,
}
impl LogError {
pub fn channel_full() -> Self {
Self {
message: "logging channel is full; dropping record".to_string(),
kind: "ChannelFull".to_string(),
}
}
pub fn channel_closed() -> Self {
Self {
message: "logging channel is closed; worker has shut down".to_string(),
kind: "ChannelClosed".to_string(),
}
}
}
impl std::fmt::Display for LogError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}: {}", self.kind, self.message)
}
}
impl std::error::Error for LogError {}

197
litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/mod.rs
View file

@ -1,197 +0,0 @@
//! A `CustomLogger` that ships finished events to the LiteLLM Python proxy's
//! `/v1/rust_control_plane/logs` endpoint.
//!
//! The callback path is non-blocking: `async_log_success_event` /
//! `async_log_failure_event`
//! build a `LogRecord` and `try_send` it onto a bounded channel, returning a
//! `LogError` (never panicking, never awaiting) if the channel is full or the
//! worker has gone away. A spawned background worker drains the channel, batches
//! records into `{"records":[...]}`, and POSTs them to the proxy with a pooled
//! `reqwest::Client`.
use std::sync::Arc;
use std::time::Duration;
use reqwest::Client;
use tokio::sync::mpsc::{self, Receiver, Sender};
use tokio::time::interval;
use crate::constants::{DEFAULT_PROXY_BASE_URL, RUST_CONTROL_PLANE_LOGS_PATH};
use crate::integrations::custom_logger::{
CallbackTiming, CallbackValue, CustomLogger, LogError, LogFuture, LoggingError,
ModelCallDetails,
};
use types::{CallbackLogsRequest, EgressTunables, LogRecord};
pub mod types;
/// Ships realtime logging events to the LiteLLM Python proxy.
pub struct LiteLLMPythonProxyAPILogger {
sink: Sender<LogRecord>,
}
impl LiteLLMPythonProxyAPILogger {
/// Spawn the background worker and return a logger handle. `base` is the
/// proxy base URL (no trailing path); `master_key` is sent as a bearer token.
pub fn start(base: String, master_key: String) -> Arc<Self> {
let tunables = EgressTunables::from_env();
let (sink, receiver) = mpsc::channel::<LogRecord>(tunables.channel_capacity);
let url = format!(
"{}{}",
base.trim_end_matches('/'),
RUST_CONTROL_PLANE_LOGS_PATH
);
let client = Client::new();
tokio::spawn(worker_loop(
receiver,
client,
url,
master_key,
tunables.max_batch_size,
tunables.flush_interval,
));
Arc::new(Self { sink })
}
/// Build a logger from the environment: `LITELLM_PROXY_BASE_URL` (default
/// `http://localhost:4000`) and `LITELLM_MASTER_KEY`.
///
/// `LITELLM_PROXY_BASE_URL` is treated as the full base and the route is
/// appended verbatim, so if the proxy runs under a `SERVER_ROOT_PATH`
/// (e.g. served at `https://host/litellm`), include it in the base
/// (`LITELLM_PROXY_BASE_URL=https://host/litellm`) and the POST lands at
/// `https://host/litellm/v1/rust_control_plane/logs`.
pub fn from_env() -> Arc<Self> {
let base = std::env::var("LITELLM_PROXY_BASE_URL")
.ok()
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| DEFAULT_PROXY_BASE_URL.to_string());
let key = std::env::var("LITELLM_MASTER_KEY").unwrap_or_default();
Self::start(base, key)
}
fn enqueue(&self, record: LogRecord) -> Result<(), LogError> {
self.sink.try_send(record).map_err(|err| match err {
mpsc::error::TrySendError::Full(_) => LogError::channel_full(),
mpsc::error::TrySendError::Closed(_) => LogError::channel_closed(),
})
}
}
impl CustomLogger for LiteLLMPythonProxyAPILogger {
fn async_log_success_event<'a>(
&'a self,
model_call_details: &'a ModelCallDetails,
_response_obj: &'a CallbackValue,
_timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async move {
if let Some(payload) = &model_call_details.standard_logging_payload {
self.enqueue(LogRecord {
status: "success".to_string(),
payload: payload.clone(),
error: None,
})?;
}
Ok(())
})
}
fn async_log_failure_event<'a>(
&'a self,
model_call_details: &'a ModelCallDetails,
_response_obj: Option<&'a CallbackValue>,
_timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async move {
if let Some(payload) = &model_call_details.standard_logging_payload {
let fallback_error;
let error = match &model_call_details.failure_error {
Some(error) => error,
None => {
fallback_error = LoggingError {
message: "callback failure event".to_string(),
kind: "CallbackFailure".to_string(),
};
&fallback_error
}
};
self.enqueue(LogRecord {
status: "failure".to_string(),
payload: payload.clone(),
error: Some(format!("{}: {}", error.kind, error.message)),
})?;
}
Ok(())
})
}
}
/// Drain the channel, batching records and POSTing them to the proxy. Exits when
/// the channel is closed (all senders dropped) and drained.
async fn worker_loop(
mut receiver: Receiver<LogRecord>,
client: Client,
url: String,
master_key: String,
max_batch_size: usize,
flush_interval: Duration,
) {
let mut ticker = interval(flush_interval);
let mut batch: Vec<LogRecord> = Vec::with_capacity(max_batch_size);
loop {
tokio::select! {
maybe_record = receiver.recv() => {
match maybe_record {
Some(record) => {
batch.push(record);
if batch.len() >= max_batch_size {
flush(&client, &url, &master_key, &mut batch).await;
}
}
None => {
// Channel closed: flush remaining and exit.
flush(&client, &url, &master_key, &mut batch).await;
break;
}
}
}
_ = ticker.tick() => {
flush(&client, &url, &master_key, &mut batch).await;
}
}
}
}
/// POST the current batch (if any), clearing it. Errors are logged, not fatal.
async fn flush(client: &Client, url: &str, master_key: &str, batch: &mut Vec<LogRecord>) {
if batch.is_empty() {
return;
}
let records = std::mem::take(batch)
.into_iter()
.map(LogRecord::into_callback_record)
.collect();
let body = CallbackLogsRequest { records };
let response = client
.post(url)
.bearer_auth(master_key)
.json(&body)
.send()
.await;
match response {
Ok(resp) if resp.status().is_success() => {}
Ok(resp) => {
eprintln!(
"litellm-ai-gateway: callback logs POST returned {} to {url}",
resp.status()
);
}
Err(err) => {
eprintln!("litellm-ai-gateway: callback logs POST failed to {url}: {err}");
}
}
}

72
litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/types.rs
View file

@ -1,72 +0,0 @@
use std::time::Duration;
use serde::Serialize;
use crate::constants::{
DEFAULT_CHANNEL_CAPACITY, DEFAULT_FLUSH_INTERVAL_MS, DEFAULT_MAX_BATCH_SIZE,
};
use crate::integrations::types::StandardLoggingPayload;
#[derive(Serialize)]
pub struct CallbackLogsRequest {
pub records: Vec<CallbackLogRecord>,
}
#[derive(Serialize)]
pub struct CallbackLogRecord {
pub status: String,
pub standard_logging_payload: StandardLoggingPayload,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}
#[derive(Clone, Debug)]
pub struct LogRecord {
pub status: String,
pub payload: StandardLoggingPayload,
pub error: Option<String>,
}
impl LogRecord {
pub fn into_callback_record(self) -> CallbackLogRecord {
CallbackLogRecord {
status: self.status,
standard_logging_payload: self.payload,
error: self.error,
}
}
}
pub(super) struct EgressTunables {
pub channel_capacity: usize,
pub max_batch_size: usize,
pub flush_interval: Duration,
}
impl EgressTunables {
pub fn from_env() -> Self {
Self {
channel_capacity: env_positive(
"LITELLM_LOG_CHANNEL_CAPACITY",
DEFAULT_CHANNEL_CAPACITY,
),
max_batch_size: env_positive("LITELLM_LOG_BATCH_SIZE", DEFAULT_MAX_BATCH_SIZE),
flush_interval: Duration::from_millis(env_positive(
"LITELLM_LOG_FLUSH_INTERVAL_MS",
DEFAULT_FLUSH_INTERVAL_MS,
)),
}
}
}
fn env_positive<T>(name: &str, default: T) -> T
where
T: std::str::FromStr + PartialOrd + From<u8>,
{
let zero = T::from(0u8);
std::env::var(name)
.ok()
.and_then(|value| value.trim().parse::<T>().ok())
.filter(|n| *n > zero)
.unwrap_or(default)
}

12
litellm-rust/crates/ai-gateway/src/integrations/mod.rs
View file

@ -1,12 +0,0 @@
//! Pure-Rust logging integrations. Names map 1:1 to Python
//! `litellm/integrations/`:
//! - [`custom_guardrail::CustomGuardrail`] — the guardrail callback trait
//! - [`custom_logger::CustomLogger`] — the callback trait
//! - [`litellm_python_proxy_api::LiteLLMPythonProxyAPILogger`] — ships events
//! to the Python proxy's `/v1/rust_control_plane/logs` endpoint
//! - [`types`] — the typed `StandardLoggingPayload` wire contract
pub mod custom_guardrail;
pub mod custom_logger;
pub mod litellm_python_proxy_api;
pub mod types;

83
litellm-rust/crates/ai-gateway/src/integrations/types.rs
View file

@ -1,83 +0,0 @@
//! Typed payloads for the LiteLLM `/v1/callbacks/logs` realtime-logging contract.
//!
//! Field names below are the EXACT JSON keys the Python replay path + spend-logs
//! builder read. Note the deliberate mix:
//! - `startTime` / `endTime` are camelCase (epoch f64 seconds)
//! - `response_cost` / `prompt_tokens` / etc. are snake_case
//!
//! Mirrors Python `litellm/integrations/` + the proxy `CallbackLogsRequest`
//! contract 1:1.
use serde::Serialize;
use serde_json::Value;
use std::collections::HashMap;
/// Cumulative token usage for a realtime session.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Usage {
pub prompt_tokens: u64,
pub completion_tokens: u64,
pub total_tokens: u64,
}
/// Cost-attribution metadata threaded from the authenticated request.
#[derive(Clone, Debug, Default)]
pub struct RequestMetadata {
pub user_api_key_hash: Option<String>,
pub user_api_key_user_id: Option<String>,
pub user_api_key_team_id: Option<String>,
}
/// The self-describing payload. Field names are the EXACT JSON keys the Python
/// replay path + spend-logs builder read.
#[derive(Clone, Debug, Serialize)]
pub struct StandardLoggingPayload {
pub id: String,
pub litellm_call_id: String,
/// e.g. "realtime", "acompletion". Falls back to "acompletion" if absent.
pub call_type: String,
pub model: String,
pub custom_llm_provider: String,
/// Spend ($) written to LiteLLM_SpendLogs.spend.
pub response_cost: f64,
pub prompt_tokens: u64,
pub completion_tokens: u64,
pub total_tokens: u64,
/// EPOCH SECONDS as float — camelCase keys, NOT snake_case.
#[serde(rename = "startTime")]
pub start_time: f64,
#[serde(rename = "endTime")]
pub end_time: f64,
pub stream: bool,
pub metadata: StandardLoggingMetadata,
/// Optional; stored as request input on the spend log row.
#[serde(skip_serializing_if = "Option::is_none")]
pub messages: Option<Value>,
}
/// Cost-attribution keys. The replayer maps these into litellm_params.metadata,
/// which the spend-logs builder reads to set user / team_id / organization_id.
#[derive(Clone, Debug, Serialize, Default)]
pub struct StandardLoggingMetadata {
pub user_api_key_hash: Option<String>, // -> SpendLogs.api_key
pub user_api_key_user_id: Option<String>, // -> SpendLogs.user
pub user_api_key_team_id: Option<String>, // -> SpendLogs.team_id
// Optional but read by the builder; include when known:
#[serde(skip_serializing_if = "Option::is_none")]
pub user_api_key_alias: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub user_api_key_org_id: Option<String>, // -> SpendLogs.organization_id
#[serde(skip_serializing_if = "Option::is_none")]
pub user_api_key_end_user_id: Option<String>, // -> SpendLogs.end_user
#[serde(skip_serializing_if = "Option::is_none")]
pub spend_logs_metadata: Option<HashMap<String, Value>>,
}

1
litellm-rust/crates/ai-gateway/src/io/audio_transcription.rs
View file

@ -1 +0,0 @@
pub use crate::audio_transcription::{AudioTranscriptionRequest, audio_transcription};

6
litellm-rust/crates/ai-gateway/src/io/mod.rs
View file

@ -1,6 +0,0 @@
pub mod audio_transcription;
pub mod ocr;
pub mod realtime;
pub mod realtime_pool;
pub mod responses_ws;
pub(crate) mod tls;

1
litellm-rust/crates/ai-gateway/src/io/ocr.rs
View file

@ -1 +0,0 @@
pub use crate::ocr::{OcrRequest, ocr};

418
litellm-rust/crates/ai-gateway/src/io/realtime.rs
View file

@ -1,418 +0,0 @@
//! End-to-end OpenAI realtime invocation.
//!
//! The host-facing entry point opens the WebSocket to OpenAI, then splices a
//! client realtime stream to the upstream, driving typed events through the pure
//! `OPENAI_REALTIME_CONFIG` transforms.
//! Network, auth header, key resolution, and wire (de)serialization live here so
//! the `transformation` module stays pure and typed.
//!
//! The dial and splice steps are factored out ([`dial_upstream`], [`splice`]) so
//! the connection pool ([`crate::io::realtime_pool`]) can pre-establish an upstream,
//! buffer its `session.created`, and later hand the live socket to the same
//! splice loop a fresh dial uses.
use std::time::Duration;
use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{Sink, SinkExt, Stream, StreamExt};
use litellm_core::AuthError;
use litellm_core::auth::error::MissingCredential;
use litellm_core::error::Error;
use litellm_core::realtime::transformation::RealtimeProviderConfig;
use litellm_core::realtime::types::RealtimeEvent;
use tokio::net::TcpStream;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG;
use crate::io::tls::connect_upstream;
/// Environment variable holding the OpenAI API key (last-resort fallback).
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
/// Default **idle** timeout: if neither side sends a frame for this long, the
/// session is reaped. It resets on any activity, so it does not cap a healthy
/// (continuously streaming) session — it only frees a stalled one (e.g. a
/// half-open upstream that keeps the socket open but stops sending).
const DEFAULT_IDLE_TIMEOUT_SECS: u64 = 300;
/// The concrete upstream WebSocket type (TLS or plain). Shared by the dial path
/// and the pool so warm sockets and fresh sockets are the exact same type.
pub type UpstreamWs = WebSocketStream<MaybeTlsStream<TcpStream>>;
pub(crate) type UpstreamTx = SplitSink<UpstreamWs, Message>;
pub(crate) type UpstreamRx = SplitStream<UpstreamWs>;
/// Resolve the OpenAI API key from the explicit param or the environment.
///
/// Blank/whitespace values are treated as absent (guard at resolution time).
pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
api_key
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| {
std::env::var(OPENAI_API_KEY_ENV)
.ok()
.filter(|key| !key.trim().is_empty())
})
.ok_or_else(|| Error::from(AuthError::from(MissingCredential::OpenAiRealtimeApiKey)))
}
/// Open the upstream WebSocket to OpenAI for `(model, api_key, api_base)`.
///
/// This is the dial half of [`realtime`], factored out so the pool can
/// pre-establish sockets ahead of any client. `api_key` here is already resolved
/// (non-blank) — the pool resolves it once when it is created.
pub(crate) async fn dial_upstream(
model: &str,
api_key: &str,
api_base: Option<&str>,
) -> Result<UpstreamWs, Error> {
let url = OPENAI_REALTIME_CONFIG.complete_url(api_base, model);
let mut request = url
.as_str()
.into_client_request()
.map_err(|err| Error::Network(err.to_string()))?;
// GA realtime: only Authorization. The legacy OpenAI-Beta header triggers
// beta_api_shape_disabled, so we do not send it.
request.headers_mut().insert(
AUTHORIZATION,
HeaderValue::from_str(&format!("Bearer {api_key}"))
.map_err(|err| Error::Auth(err.to_string()))?,
);
let (upstream, _response) = connect_upstream(request)
.await
.map_err(|err| Error::Network(err.to_string()))?;
Ok(upstream)
}
/// Read the next text frame from the upstream and decode it as a typed event.
///
/// Used by the pool to pre-read OpenAI's unprompted `session.created`. Returns an
/// error on a non-text frame, a closed socket, or undecodable JSON so the pool can
/// discard a misbehaving socket rather than warm it.
pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> Result<RealtimeEvent, Error> {
loop {
let message = upstream_rx
.next()
.await
.ok_or_else(|| Error::Network("upstream closed before first event".to_string()))?
.map_err(|err| Error::Network(err.to_string()))?;
match message {
Message::Text(text) => {
return serde_json::from_str(&text)
.map_err(|err| Error::InvalidResponse(err.to_string()));
}
// Ignore protocol frames (ping/pong) while waiting for the first event.
Message::Ping(_) | Message::Pong(_) => continue,
Message::Close(_) => {
return Err(Error::Network(
"upstream closed before first event".to_string(),
));
}
_ => continue,
}
}
}
/// Splice an already-connected upstream to the client streams.
///
/// `prelude` is relayed to the client first (the pool passes the buffered
/// `session.created` here; the fresh-dial path passes `None` and lets the upstream
/// deliver it). Then a single select loop forwards both directions through the
/// transforms until either side closes or the idle timeout fires.
/// `observe` is invoked on **upstream→client** events only (the trusted side that
/// carries `session.created` and `response.done` usage) — never on client events,
/// so a client cannot fabricate usage into its own logs.
#[allow(clippy::too_many_arguments)]
pub(crate) async fn splice<In, Out>(
model: &str,
mut upstream_tx: UpstreamTx,
mut upstream_rx: UpstreamRx,
prelude: Option<RealtimeEvent>,
idle_timeout: Option<Duration>,
mut observe: impl FnMut(&RealtimeEvent) + Send,
mut client_in: In,
mut client_out: Out,
) -> Result<(), Error>
where
In: Stream<Item = RealtimeEvent> + Unpin + Send,
Out: Sink<RealtimeEvent> + Unpin + Send,
<Out as Sink<RealtimeEvent>>::Error: std::fmt::Display,
{
let config = &OPENAI_REALTIME_CONFIG;
// Relay a buffered backend event (warm handoff's session.created) first, so a
// warm session looks identical to a fresh one from the client's view.
if let Some(event) = prelude {
for outbound in config.transform_realtime_response(&event, model)?.events {
client_out
.send(outbound)
.await
.map_err(|err| Error::Network(err.to_string()))?;
}
}
let idle = idle_timeout.unwrap_or(Duration::from_secs(DEFAULT_IDLE_TIMEOUT_SECS));
// One loop forwarding both directions. The `sleep(idle)` arm is rebuilt every
// iteration, so any frame (either way) resets it — it fires only when the
// session has been fully idle for `idle`, reaping a stalled connection
// (task + upstream TCP socket) instead of leaking it.
loop {
tokio::select! {
// client -> upstream
client_event = client_in.next() => {
let Some(event) = client_event else { break }; // client disconnected
// NOTE: do NOT observe client events. session.created / response.done
// (carrying usage) are server→client events; observing the client arm
// would let an authenticated client POST a fabricated response.done and
// inflate its own spend log. Logging observes upstream events only.
for outbound in config.transform_realtime_request(&event, model)?.events {
let payload = serde_json::to_string(&outbound)
.map_err(|err| Error::InvalidResponse(err.to_string()))?;
upstream_tx
.send(Message::Text(payload))
.await
.map_err(|err| Error::Network(err.to_string()))?;
}
}
// upstream -> client
upstream_message = upstream_rx.next() => {
let Some(message) = upstream_message else { break }; // upstream closed
match message.map_err(|err| Error::Network(err.to_string()))? {
Message::Text(text) => {
let event: RealtimeEvent = serde_json::from_str(&text)
.map_err(|err| Error::InvalidResponse(err.to_string()))?;
observe(&event);
for outbound in config.transform_realtime_response(&event, model)?.events {
client_out
.send(outbound)
.await
.map_err(|err| Error::Network(err.to_string()))?;
}
}
Message::Close(_) => break,
_ => {}
}
}
// idle timeout: no activity from either side within `idle`
_ = tokio::time::sleep(idle) => break,
}
}
Ok(())
}
/// Splice a client realtime stream to OpenAI: forward client events upstream
/// (via `transform_realtime_request`) and backend events downstream (via
/// `transform_realtime_response`). Returns when either side closes.
///
/// Generic over the client transport (typed events) so this crate stays
/// framework-agnostic; the gateway adapts its axum socket to these. This is the
/// fresh-dial path: dial, then splice. The pool's warm-handoff path skips the dial
/// and calls [`splice`] directly with a buffered `session.created`.
#[allow(clippy::too_many_arguments)]
pub async fn realtime<In, Out>(
model: &str,
api_key: Option<&str>,
api_base: Option<&str>,
idle_timeout: Option<Duration>,
observe: impl FnMut(&RealtimeEvent) + Send,
client_in: In,
client_out: Out,
) -> Result<(), Error>
where
In: Stream<Item = RealtimeEvent> + Unpin + Send,
Out: Sink<RealtimeEvent> + Unpin + Send,
<Out as Sink<RealtimeEvent>>::Error: std::fmt::Display,
{
let api_key = resolve_api_key(api_key)?;
let upstream = dial_upstream(model, &api_key, api_base).await?;
let (upstream_tx, upstream_rx) = upstream.split();
splice(
model,
upstream_tx,
upstream_rx,
None,
idle_timeout,
observe,
client_in,
client_out,
)
.await
}
/// Splice a pre-warmed upstream (taken from [`crate::io::realtime_pool`]) to the
/// client. Relays the buffered `session.created` first, then splices exactly like
/// the fresh-dial path — so a warm session is indistinguishable from a fresh one.
#[allow(clippy::too_many_arguments)]
pub async fn realtime_warm<In, Out>(
model: &str,
handoff: crate::io::realtime_pool::WarmHandoff,
idle_timeout: Option<Duration>,
observe: impl FnMut(&RealtimeEvent) + Send,
client_in: In,
client_out: Out,
) -> Result<(), Error>
where
In: Stream<Item = RealtimeEvent> + Unpin + Send,
Out: Sink<RealtimeEvent> + Unpin + Send,
<Out as Sink<RealtimeEvent>>::Error: std::fmt::Display,
{
splice(
model,
handoff.tx,
handoff.rx,
Some(handoff.session_created),
idle_timeout,
observe,
client_in,
client_out,
)
.await
}
#[cfg(test)]
mod tests {
use super::*;
fn event(raw: &str) -> RealtimeEvent {
serde_json::from_str(raw).expect("valid event json")
}
/// The realtime dial has to reach a `wss://` upstream without a process-wide
/// crypto provider installed, which is what dialing through `io::tls` buys.
#[tokio::test]
async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind a loopback port");
let port = listener
.local_addr()
.expect("read the bound address")
.port();
tokio::spawn(async move {
while let Ok((stream, _peer)) = listener.accept().await {
drop(stream);
}
});
let result = dial_upstream(
"gpt-realtime",
"sk-test",
Some(&format!("wss://127.0.0.1:{port}")),
)
.await;
assert!(matches!(result, Err(Error::Network(_))));
}
#[test]
fn resolve_api_key_prefers_param_then_blank_falls_through() {
assert_eq!(resolve_api_key(Some("sk-test")).unwrap(), "sk-test");
// A blank param with no env set should error.
if std::env::var(OPENAI_API_KEY_ENV).is_err() {
assert!(resolve_api_key(Some(" ")).is_err());
}
}
/// Live end-to-end check against OpenAI. Ignored by default (CI never runs
/// it); run explicitly with `OPENAI_API_KEY` set:
/// `cargo test -p litellm-ai-gateway --features server realtime_invokes_openai -- --ignored --nocapture`
#[tokio::test]
#[ignore = "hits the live OpenAI realtime API; needs OPENAI_API_KEY"]
async fn realtime_invokes_openai_and_responds() {
use futures_channel::mpsc;
let key =
std::env::var(OPENAI_API_KEY_ENV).expect("set OPENAI_API_KEY to run this ignored test");
// client -> provider (we hold `client_tx` to push events upstream)
let (mut client_tx, client_in) = mpsc::unbounded::<RealtimeEvent>();
// provider -> client (we hold `backend_rx` to read backend events)
let (client_out, mut backend_rx) = mpsc::unbounded::<RealtimeEvent>();
// Clone the key so the spawned task owns its `String` (no borrow across await).
let key_owned = key.clone();
let call = tokio::spawn(async move {
realtime(
"gpt-realtime",
Some(&key_owned),
None,
None,
|_| {},
client_in,
client_out,
)
.await
});
// 1. First backend event should be session.created.
let first = tokio::time::timeout(Duration::from_secs(30), backend_rx.next())
.await
.expect("timed out waiting for session.created")
.expect("backend stream closed before session.created");
assert_eq!(
first.event_type, "session.created",
"expected session.created, got: {}",
first.event_type
);
// 2. Ask for a short audio response.
client_tx
.send(event(
r#"{"type":"conversation.item.create","item":{"type":"message","role":"user","content":[{"type":"input_text","text":"Say hi."}]}}"#,
))
.await
.expect("send conversation.item.create");
client_tx
.send(event(r#"{"type":"response.create"}"#))
.await
.expect("send response.create");
// 3. Read backend events; require a non-empty audio delta, then response.done.
let mut saw_audio_delta = false;
let mut saw_done = false;
for _ in 0..500 {
let next = tokio::time::timeout(Duration::from_secs(30), backend_rx.next()).await;
let event = match next {
Ok(Some(event)) => event,
Ok(None) => break,
Err(_) => panic!("timed out waiting for backend events"),
};
match event.event_type.as_str() {
"response.output_audio.delta" => {
let delta = event
.data
.get("delta")
.and_then(|value| value.as_str())
.unwrap_or("");
if !delta.is_empty() {
saw_audio_delta = true;
}
}
"response.done" => {
saw_done = true;
break;
}
_ => {}
}
}
assert!(
saw_audio_delta,
"expected a response.output_audio.delta with non-empty delta"
);
assert!(saw_done, "expected a response.done event");
// Drop the client sender so the provider's to_upstream side finishes.
drop(client_tx);
let _ = call.await;
}
}

712
litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs
View file

@ -1,712 +0,0 @@
//! Pre-warmed upstream realtime connection pool.
//!
//! The gateway's realtime overhead lives entirely in session establishment: on
//! every client connect it dials a fresh upstream WS to OpenAI and waits for
//! `session.created` before it can serve. This pool keeps a small set of upstream
//! sockets **already connected and already past `session.created`** so a connect
//! can be served from a warm socket and the handshake is off the critical path.
//!
//! Layering: this lives in the gateway's `io` module next to the dial/splice it
//! reuses. The gateway holds an `Arc<RealtimePool>` in its state and asks for a
//! warm socket per connect; on a miss it fresh-dials exactly as before. The pool
//! is a latency optimization, never a correctness dependency — see the gateway's
//! `src/routes/realtime/README.md`.
//!
//! ## Caveats (enforced here)
//! - One warm socket serves exactly one session (realtime isn't multiplexed), so
//! the pool is sized to the connect *rate*, not concurrent connections.
//! - `session.created` is pre-read once and buffered; nothing else is read from a
//! warm socket before handoff, so a warm session starts at OpenAI defaults just
//! like a fresh one (`session.update` semantics unchanged).
//! - Warm sockets are short-lived (`max_idle`) and liveness-checked at handoff to
//! bound idle billing / dodge OpenAI's idle timeout.
//! - On miss or dead socket the caller fresh-dials; the pool never blocks or fails
//! a connect because it is empty.
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use futures_util::StreamExt;
use litellm_core::Error;
use litellm_core::realtime::types::RealtimeEvent;
use crate::io::realtime::{
UpstreamRx, UpstreamTx, UpstreamWs, dial_upstream, read_event, resolve_api_key,
};
/// Default target warm sockets per key when pooling is enabled.
pub const DEFAULT_POOL_SIZE: usize = 4;
/// Default max time a warm socket may sit before it is closed and replaced.
pub const DEFAULT_MAX_IDLE: Duration = Duration::from_secs(30);
/// Env var: target warm sockets per key. `0` disables pooling (fresh-dial only).
pub const POOL_SIZE_ENV: &str = "REALTIME_POOL_SIZE";
/// Env var: max warm-socket idle lifetime, in seconds.
pub const MAX_IDLE_ENV: &str = "REALTIME_POOL_MAX_IDLE_SECS";
/// How often the background replenisher wakes to top up and reap stale sockets.
const REPLENISH_TICK: Duration = Duration::from_millis(250);
/// Backoff floor after a key's warm-up dials all fail. The first failed pass
/// waits this long before retrying that key.
const BACKOFF_BASE: Duration = Duration::from_millis(500);
/// Backoff ceiling. A key that keeps failing (invalid credentials, an
/// unreachable upstream) is retried at most once per this interval — instead of
/// firing `needed` concurrent TLS dials every 250 ms tick, which would hammer
/// the upstream and risk rate-limit exhaustion that degrades valid cold-path
/// traffic. Backoff resets the moment a dial for the key succeeds.
const BACKOFF_MAX: Duration = Duration::from_secs(30);
/// Identifies an upstream connection: the tuple that fully determines the dial.
/// `api_key` is included so a warm socket is only ever reused for the same key
/// (no cross-tenant reuse).
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct UpstreamKey {
pub model: String,
pub api_key: String,
pub api_base: Option<String>,
}
impl std::fmt::Debug for UpstreamKey {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("UpstreamKey")
.field("model", &self.model)
.field("api_key", &"[REDACTED]")
.field("api_base", &self.api_base)
.finish()
}
}
/// A warm upstream: split halves + the buffered `session.created` + when it was
/// warmed (for `max_idle` expiry).
struct WarmConnection {
tx: UpstreamTx,
rx: UpstreamRx,
session_created: RealtimeEvent,
warmed_at: Instant,
}
/// A live upstream taken from the pool, ready to splice. The caller relays
/// `session_created` to the client first, then splices `(tx, rx)` as usual.
pub struct WarmHandoff {
pub tx: UpstreamTx,
pub rx: UpstreamRx,
pub session_created: RealtimeEvent,
}
/// Pool configuration, resolved once at startup from the environment.
#[derive(Clone, Copy, Debug)]
pub struct PoolConfig {
/// Target warm sockets per key. `0` disables pooling.
pub target_size: usize,
/// Max time a warm socket may sit before it is closed and replaced.
pub max_idle: Duration,
}
impl Default for PoolConfig {
fn default() -> Self {
Self {
target_size: DEFAULT_POOL_SIZE,
max_idle: DEFAULT_MAX_IDLE,
}
}
}
impl PoolConfig {
/// Read config from the environment, falling back to defaults. An invalid
/// value warns and uses the default rather than failing startup.
pub fn from_env() -> Self {
let target_size = match std::env::var(POOL_SIZE_ENV) {
Ok(raw) => raw.trim().parse().unwrap_or_else(|_| {
eprintln!("warning: {POOL_SIZE_ENV}={raw:?} is not a valid size; using {DEFAULT_POOL_SIZE}");
DEFAULT_POOL_SIZE
}),
Err(_) => DEFAULT_POOL_SIZE,
};
let max_idle = match std::env::var(MAX_IDLE_ENV) {
Ok(raw) => raw
.trim()
.parse()
.map(Duration::from_secs)
.unwrap_or_else(|_| {
eprintln!(
"warning: {MAX_IDLE_ENV}={raw:?} is not a valid number of seconds; using {}s",
DEFAULT_MAX_IDLE.as_secs()
);
DEFAULT_MAX_IDLE
}),
Err(_) => DEFAULT_MAX_IDLE,
};
Self {
target_size,
max_idle,
}
}
/// Whether pooling is on (`target_size > 0`).
pub fn enabled(&self) -> bool {
self.target_size > 0
}
}
/// Per-key warm sockets, behind a single `Mutex`. Realtime warm sockets are few
/// (the pool is small), so a plain mutex over a `VecDeque`-ish `Vec` is simpler
/// and faster than sharding; contention is negligible at this scale.
type Warm = HashMap<UpstreamKey, Vec<WarmConnection>>;
/// Per-key replenish backoff. Absent (or `consecutive_failures == 0`) means the
/// key is healthy and replenished every tick. After a pass whose dials all fail,
/// `retry_after` is pushed out with exponential backoff so a broken key (invalid
/// credentials, unreachable upstream) is not re-dialed on every 250 ms tick.
#[derive(Default)]
struct Backoff {
/// Don't attempt warm-up dials for this key until this instant. `None` =
/// eligible now.
retry_after: Option<Instant>,
consecutive_failures: u32,
}
type Backoffs = HashMap<UpstreamKey, Backoff>;
/// Pre-warmed upstream realtime connection pool.
///
/// Cheap to clone-via-`Arc`. The background replenisher is spawned by
/// [`RealtimePool::spawn`]; a pool built with [`RealtimePool::disabled`] never
/// warms anything and every `take` misses (callers fresh-dial).
pub struct RealtimePool {
config: PoolConfig,
warm: Mutex<Warm>,
/// Per-key replenish backoff so a broken key doesn't trigger unbounded
/// concurrent dials every tick. Separate lock from `warm` so the request
/// hot path (`take`) never contends on it.
backoff: Mutex<Backoffs>,
}
impl RealtimePool {
/// A disabled pool: no background task, every `take` returns `None`.
pub fn disabled() -> Arc<Self> {
Arc::new(Self {
config: PoolConfig {
target_size: 0,
..PoolConfig::default()
},
warm: Mutex::new(HashMap::new()),
backoff: Mutex::new(HashMap::new()),
})
}
/// Build a pool from config **without** the background replenisher. The pool
/// only warms when [`RealtimePool::warm_now`] is called. Used by deterministic
/// unit tests; production uses [`RealtimePool::spawn`].
#[cfg(test)]
fn new_unspawned(config: PoolConfig) -> Arc<Self> {
Arc::new(Self {
config,
warm: Mutex::new(HashMap::new()),
backoff: Mutex::new(HashMap::new()),
})
}
/// Build a pool from config and, if enabled, spawn the background replenisher.
/// Returns the shared handle the gateway stores in its state.
pub fn spawn(config: PoolConfig) -> Arc<Self> {
let pool = Arc::new(Self {
config,
warm: Mutex::new(HashMap::new()),
backoff: Mutex::new(HashMap::new()),
});
if config.enabled() {
let weak = Arc::downgrade(&pool);
tokio::spawn(async move {
let mut tick = tokio::time::interval(REPLENISH_TICK);
tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
tick.tick().await;
// Stop once the gateway has dropped its handle.
let Some(pool) = weak.upgrade() else { break };
pool.replenish_all().await;
}
});
}
pool
}
/// Resolved config (test/inspection).
pub fn config(&self) -> PoolConfig {
self.config
}
/// Register a key so the replenisher starts warming it. Idempotent. The
/// gateway calls this once per known deployment at startup; the pool only
/// warms keys it has seen, so it never dials a model nobody asked for.
pub fn register(&self, key: UpstreamKey) {
if !self.config.enabled() {
return;
}
self.warm.lock().unwrap().entry(key).or_default();
}
/// Take a warm, live socket for `key`, or `None` on miss / dead socket.
///
/// Pops the freshest non-expired socket and liveness-checks it; a socket that
/// is too old or already dead is dropped (closing it) and the next candidate
/// tried. Never blocks: if nothing warm is live, returns `None` so the caller
/// fresh-dials.
pub fn take(&self, key: &UpstreamKey) -> Option<WarmHandoff> {
if !self.config.enabled() {
return None;
}
loop {
let mut candidate = {
let mut warm = self.warm.lock().unwrap();
let bucket = warm.get_mut(key)?;
bucket.pop()?
};
// Discard sockets past their warm lifetime (idle-billing guard).
if candidate.warmed_at.elapsed() > self.config.max_idle {
continue; // drops `candidate`, closing the socket
}
// Liveness: a non-blocking check that the socket hasn't already
// delivered a Close/Err. A warm socket should be silent after
// session.created, so anything pending means it is unhealthy.
if is_dead(&mut candidate.rx) {
continue;
}
return Some(WarmHandoff {
tx: candidate.tx,
rx: candidate.rx,
session_created: candidate.session_created,
});
}
}
/// One replenish pass over every registered key: reap stale sockets, then
/// dial up to `target_size`. Dials run concurrently; failures are swallowed
/// (a key that can't be warmed just keeps fresh-dialing on the request path)
/// and put the key into exponential backoff so a broken key isn't re-dialed
/// on every tick.
async fn replenish_all(&self) {
let keys: Vec<UpstreamKey> = { self.warm.lock().unwrap().keys().cloned().collect() };
for key in keys {
self.reap_stale(&key);
// Skip keys still in backoff from a prior all-failed pass — this is
// what bounds dials against an invalid/unreachable key to once per
// `BACKOFF_MAX` instead of `needed` dials every 250 ms tick.
if self.in_backoff(&key) {
continue;
}
let needed = {
let warm = self.warm.lock().unwrap();
let have = warm.get(&key).map(Vec::len).unwrap_or(0);
self.config.target_size.saturating_sub(have)
};
if needed == 0 {
continue;
}
// Dial the missing sockets CONCURRENTLY. A sequential loop here makes
// a full refill cost `needed × handshake` (~needed × 350 ms), which
// can't keep up with a high connect rate — the pool drains faster
// than it refills and most connects miss. Firing the dials together
// refills in ~one handshake window, keeping warm supply ≈ peak
// concurrent connects so the sub-ms warm handoff becomes the median,
// not the lucky-hit tail.
let dials = (0..needed).map(|_| warm_one(&key));
let results = futures_util::future::join_all(dials).await;
let mut any_ok = false;
// `.flatten()` keeps only the successful dials; a key that can't be
// warmed just keeps fresh-dialing on the request path.
for conn in results.into_iter().flatten() {
any_ok = true;
self.warm
.lock()
.unwrap()
.entry(key.clone())
.or_default()
.push(conn);
}
// Reset backoff on any success; otherwise grow it. We only ever enter
// backoff when a pass that *attempted* dials produced none — a `needed
// == 0` pass is handled by the `continue` above and never touches it.
self.record_replenish_outcome(&key, any_ok);
}
}
/// Whether `key` is currently in a backoff window (a prior pass failed and
/// the retry time hasn't arrived). Eligible keys are pruned from the backoff
/// map so it doesn't grow unbounded for healthy keys.
fn in_backoff(&self, key: &UpstreamKey) -> bool {
let mut backoff = self.backoff.lock().unwrap();
match backoff.get(key).and_then(|b| b.retry_after) {
Some(retry_after) if Instant::now() < retry_after => true,
Some(_) => {
// Window elapsed — allow the attempt. Keep the failure count so a
// still-broken key backs off further, but clear the gate so this
// tick proceeds.
if let Some(b) = backoff.get_mut(key) {
b.retry_after = None;
}
false
}
None => false,
}
}
/// Update a key's backoff after a replenish attempt. Success clears it;
/// failure grows the retry delay exponentially up to `BACKOFF_MAX`.
fn record_replenish_outcome(&self, key: &UpstreamKey, any_ok: bool) {
let mut backoff = self.backoff.lock().unwrap();
if any_ok {
backoff.remove(key);
return;
}
let entry = backoff.entry(key.clone()).or_default();
entry.consecutive_failures = entry.consecutive_failures.saturating_add(1);
// Exponential: BASE * 2^(failures-1), saturating at MAX. `min` of the
// shift exponent keeps the doubling from overflowing.
let shift = (entry.consecutive_failures - 1).min(16);
let delay = BACKOFF_BASE.saturating_mul(1u32 << shift).min(BACKOFF_MAX);
entry.retry_after = Some(Instant::now() + delay);
}
/// Drop sockets past `max_idle` or already dead for a key.
fn reap_stale(&self, key: &UpstreamKey) {
let mut warm = self.warm.lock().unwrap();
if let Some(bucket) = warm.get_mut(key) {
bucket.retain_mut(|conn| {
conn.warmed_at.elapsed() <= self.config.max_idle && !is_dead(&mut conn.rx)
});
}
}
/// Test/inspection: number of warm sockets currently held for `key`.
#[cfg(test)]
pub fn warm_len(&self, key: &UpstreamKey) -> usize {
self.warm
.lock()
.unwrap()
.get(key)
.map(Vec::len)
.unwrap_or(0)
}
/// Test/inspection: consecutive replenish failures recorded for `key` (0 if
/// the key is healthy / has no backoff entry).
#[cfg(test)]
pub fn backoff_failures(&self, key: &UpstreamKey) -> u32 {
self.backoff
.lock()
.unwrap()
.get(key)
.map(|b| b.consecutive_failures)
.unwrap_or(0)
}
/// Test helper: synchronously warm `target_size` sockets for `key` (no
/// background task). Lets tests assert handoff behavior deterministically.
#[cfg(test)]
pub async fn warm_now(&self, key: &UpstreamKey) {
let needed = {
let warm = self.warm.lock().unwrap();
let have = warm.get(key).map(Vec::len).unwrap_or(0);
self.config.target_size.saturating_sub(have)
};
for _ in 0..needed {
if let Ok(conn) = warm_one(key).await {
self.warm
.lock()
.unwrap()
.entry(key.clone())
.or_default()
.push(conn);
}
}
}
/// Test helper: insert an already-built warm connection (used to inject a
/// dead socket and assert it is discarded at handoff).
#[cfg(test)]
fn insert_warm(&self, key: UpstreamKey, conn: WarmConnection) {
self.warm.lock().unwrap().entry(key).or_default().push(conn);
}
}
/// Dial one upstream and pre-read its `session.created` into a [`WarmConnection`].
///
/// `key.api_key` is already resolved (non-blank). The first frame OpenAI sends
/// unprompted is `session.created`; we buffer exactly that and read nothing more.
async fn warm_one(key: &UpstreamKey) -> Result<WarmConnection, Error> {
let upstream: UpstreamWs =
dial_upstream(&key.model, &key.api_key, key.api_base.as_deref()).await?;
let (tx, mut rx) = upstream.split();
let session_created = read_event(&mut rx).await?;
Ok(WarmConnection {
tx,
rx,
session_created,
warmed_at: Instant::now(),
})
}
/// Resolve a deployment's API key into the pool key, returning `None` when no key
/// can be resolved (those deployments simply aren't pooled — the request path
/// still fresh-dials and surfaces the auth error there).
pub fn upstream_key(
model: &str,
api_key: Option<&str>,
api_base: Option<&str>,
) -> Option<UpstreamKey> {
let api_key = resolve_api_key(api_key).ok()?;
Some(UpstreamKey {
model: model.to_string(),
api_key,
api_base: api_base.map(str::to_string),
})
}
/// Non-blocking liveness check: poll the upstream once. A warm socket is silent
/// after `session.created`, so a pending `Close`/`Err`/`None` means it is dead.
/// A pending data frame (shouldn't happen pre-handoff) is also treated as
/// unhealthy — we'd rather discard and fresh-dial than hand over a socket in an
/// unexpected state. `Pending` (the healthy case) returns `false`.
fn is_dead(rx: &mut UpstreamRx) -> bool {
use futures_util::Stream;
use futures_util::task::noop_waker_ref;
use std::pin::Pin;
use std::task::{Context, Poll};
let mut cx = Context::from_waker(noop_waker_ref());
match Pin::new(rx).poll_next(&mut cx) {
Poll::Pending => false,
Poll::Ready(None) => true,
Poll::Ready(Some(Err(_))) => true,
// Any frame arriving before handoff is unexpected for a silent warm
// socket; treat it as unhealthy.
Poll::Ready(Some(Ok(_))) => true,
}
}
#[cfg(test)]
mod tests {
use super::*;
use futures_util::SinkExt;
use std::net::SocketAddr;
use tokio::net::TcpListener;
use tokio_tungstenite::tungstenite::Message;
/// An in-process fake OpenAI realtime WS server. On connect it sends
/// `session.created`; on `response.create` it sends `response.created` +
/// `response.output_audio.delta` + `response.done`. Returns its `ws://` base.
async fn spawn_fake_openai() -> String {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr: SocketAddr = listener.local_addr().unwrap();
tokio::spawn(async move {
while let Ok((stream, _)) = listener.accept().await {
tokio::spawn(handle_fake_conn(stream));
}
});
format!("ws://{addr}")
}
async fn handle_fake_conn(stream: tokio::net::TcpStream) {
let mut ws = match tokio_tungstenite::accept_async(stream).await {
Ok(ws) => ws,
Err(_) => return,
};
// Unprompted session.created, exactly like OpenAI.
let _ = ws
.send(Message::Text(
r#"{"type":"session.created","session":{"id":"sess_fake"}}"#.to_string(),
))
.await;
while let Some(Ok(msg)) = ws.next().await {
if let Message::Text(text) = msg
&& text.contains("response.create")
{
for frame in [
r#"{"type":"response.created"}"#,
r#"{"type":"response.output_audio.delta","delta":"AAAA"}"#,
r#"{"type":"response.done"}"#,
] {
let _ = ws.send(Message::Text(frame.to_string())).await;
}
}
}
}
fn test_config() -> PoolConfig {
PoolConfig {
target_size: 2,
max_idle: Duration::from_secs(30),
}
}
fn key_for(base: &str) -> UpstreamKey {
UpstreamKey {
model: "gpt-realtime".to_string(),
api_key: "sk-test".to_string(),
api_base: Some(base.to_string()),
}
}
#[tokio::test]
async fn warm_handoff_relays_buffered_session_created() {
let base = spawn_fake_openai().await;
let pool = RealtimePool::new_unspawned(test_config());
let key = key_for(&base);
pool.register(key.clone());
pool.warm_now(&key).await;
assert_eq!(pool.warm_len(&key), 2);
let handoff = pool.take(&key).expect("a warm socket should be available");
assert_eq!(handoff.session_created.event_type, "session.created");
assert_eq!(
handoff
.session_created
.data
.get("session")
.and_then(|s| s.get("id"))
.and_then(|v| v.as_str()),
Some("sess_fake")
);
// Taking one leaves one.
assert_eq!(pool.warm_len(&key), 1);
}
#[tokio::test]
async fn pool_miss_returns_none_for_fresh_dial_fallback() {
let base = spawn_fake_openai().await;
let pool = RealtimePool::new_unspawned(test_config());
let key = key_for(&base);
// Registered but never warmed → empty bucket → miss.
pool.register(key.clone());
assert!(pool.take(&key).is_none());
// Unknown key → miss.
let other = key_for("ws://127.0.0.1:1");
assert!(pool.take(&other).is_none());
}
#[tokio::test]
async fn disabled_pool_never_hands_off() {
let pool = RealtimePool::disabled();
let key = key_for("ws://127.0.0.1:1");
pool.register(key.clone());
assert_eq!(pool.warm_len(&key), 0);
assert!(pool.take(&key).is_none());
}
#[tokio::test]
async fn dead_warm_socket_is_discarded() {
let base = spawn_fake_openai().await;
let pool = RealtimePool::new_unspawned(test_config());
let key = key_for(&base);
pool.register(key.clone());
// Build one real warm connection, then kill the upstream by dropping the
// server side: easiest is to dial, read session.created, then close our
// own rx's peer. Instead we forge "dead" via an already-closed socket:
// dial a connection and immediately send a Close from the client side so
// the server closes back, then warm it. Simpler: warm normally, then
// mark it stale by backdating warmed_at past max_idle and confirm it's
// dropped — that exercises the same discard path.
let mut conn = warm_one(&key).await.expect("warm one");
conn.warmed_at = Instant::now() - Duration::from_secs(3600); // past max_idle
pool.insert_warm(key.clone(), conn);
assert_eq!(pool.warm_len(&key), 1);
// take() must discard the stale socket and report a miss.
assert!(pool.take(&key).is_none());
assert_eq!(pool.warm_len(&key), 0);
}
#[tokio::test]
async fn background_replenisher_tops_up_registered_key() {
let base = spawn_fake_openai().await;
let pool = RealtimePool::spawn(test_config());
let key = key_for(&base);
pool.register(key.clone());
// Wait (bounded) for the background task to reach the target size.
let mut warmed = 0;
for _ in 0..40 {
tokio::time::sleep(Duration::from_millis(50)).await;
warmed = pool.warm_len(&key);
if warmed >= test_config().target_size {
break;
}
}
assert_eq!(
warmed,
test_config().target_size,
"background replenisher should warm up to target_size"
);
let handoff = pool.take(&key).expect("a warm socket should be available");
assert_eq!(handoff.session_created.event_type, "session.created");
}
#[tokio::test]
async fn closed_upstream_socket_is_detected_dead() {
// A genuinely dead socket: dial the fake, read session.created, then drop
// the server by closing from our side and waiting for the close to land.
let base = spawn_fake_openai().await;
let pool = RealtimePool::new_unspawned(test_config());
let key = key_for(&base);
pool.register(key.clone());
let mut conn = warm_one(&key).await.expect("warm one");
// Close the upstream from the client side; the server echoes a close.
let _ = conn.tx.send(Message::Close(None)).await;
// Give the close a moment to arrive on rx.
tokio::time::sleep(Duration::from_millis(50)).await;
pool.insert_warm(key.clone(), conn);
// Liveness check at take() should detect the close and discard it.
assert!(pool.take(&key).is_none());
assert_eq!(pool.warm_len(&key), 0);
}
#[tokio::test]
async fn broken_key_backs_off_instead_of_dialing_every_tick() {
// A key whose upstream is unreachable: every warm-up dial fails.
let pool = RealtimePool::new_unspawned(test_config());
let key = key_for("ws://127.0.0.1:1"); // nothing listens here
pool.register(key.clone());
// First pass attempts dials, they all fail → key enters backoff, no warm
// sockets, one recorded failure.
pool.replenish_all().await;
assert_eq!(pool.warm_len(&key), 0);
assert_eq!(pool.backoff_failures(&key), 1);
assert!(
pool.in_backoff(&key),
"a key whose dials all failed must be in backoff"
);
// An immediate next pass must be SKIPPED (still in the backoff window), so
// it does NOT fire another round of dials — the failure count is unchanged.
pool.replenish_all().await;
assert_eq!(
pool.backoff_failures(&key),
1,
"replenish during the backoff window must not re-dial the broken key"
);
}
#[tokio::test]
async fn healthy_key_never_enters_backoff_and_clears_after_recovery() {
let base = spawn_fake_openai().await;
let pool = RealtimePool::new_unspawned(test_config());
let key = key_for(&base);
pool.register(key.clone());
// A reachable upstream: the pass succeeds, so the key is never backed off.
pool.replenish_all().await;
assert_eq!(pool.warm_len(&key), test_config().target_size);
assert_eq!(pool.backoff_failures(&key), 0);
assert!(!pool.in_backoff(&key));
}
}

485
litellm-rust/crates/ai-gateway/src/io/responses_ws.rs
View file

@ -1,485 +0,0 @@
use std::time::Duration;
use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{Sink, SinkExt, Stream, StreamExt};
use litellm_core::AuthError;
use litellm_core::Error;
use litellm_core::auth::error::MissingCredential;
use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG;
use litellm_core::responses::types::ResponsesWsEvent;
use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
use litellm_core::responses::websocket::{ResponsesUpstreamWs, connect_upstream};
use crate::constants::{
DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS,
};
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
type UpstreamTx = SplitSink<ResponsesUpstreamWs, Message>;
type UpstreamRx = SplitStream<ResponsesUpstreamWs>;
pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
api_key
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
.or_else(|| {
std::env::var(OPENAI_API_KEY_ENV)
.ok()
.filter(|value| !value.trim().is_empty())
})
.ok_or_else(|| Error::from(AuthError::from(MissingCredential::OpenAiResponsesApiKey)))
}
async fn dial_upstream(
model: &str,
api_key: &str,
api_base: Option<&str>,
) -> Result<ResponsesUpstreamWs, Error> {
let url = OPENAI_RESPONSES_WS_CONFIG.complete_websocket_url(api_base, model);
let mut request = url
.as_str()
.into_client_request()
.map_err(|error| Error::Network(error.to_string()))?;
request.headers_mut().insert(
AUTHORIZATION,
HeaderValue::from_str(&format!("Bearer {api_key}"))
.map_err(|error| Error::Auth(error.to_string()))?,
);
let result = tokio::time::timeout(
Duration::from_secs(DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS),
connect_upstream(request),
)
.await
.map_err(|_| Error::Network("Responses WebSocket connection timed out".to_string()))?;
result
.map(|(socket, _)| socket)
.map_err(|error| match *error {
tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
status: response.status().as_u16(),
body: String::new(),
},
other => Error::Network(other.to_string()),
})
}
pub struct ResponsesWebSocketStreaming;
impl ResponsesWebSocketStreaming {
pub async fn bidirectional_forward<In, Out>(
model: &str,
upstream_tx: UpstreamTx,
upstream_rx: UpstreamRx,
idle_timeout: Option<Duration>,
observe: impl FnMut(&ResponsesWsEvent) + Send,
client_in: In,
client_out: Out,
) -> Result<(), Error>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
Out::Error: std::fmt::Display,
{
splice(
model,
upstream_tx,
upstream_rx,
idle_timeout,
observe,
client_in,
client_out,
)
.await
}
}
pub(crate) async fn splice<In, Out>(
model: &str,
mut upstream_tx: UpstreamTx,
mut upstream_rx: UpstreamRx,
idle_timeout: Option<Duration>,
mut observe: impl FnMut(&ResponsesWsEvent) + Send,
mut client_in: In,
mut client_out: Out,
) -> Result<(), Error>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
Out::Error: std::fmt::Display,
{
let idle =
idle_timeout.unwrap_or_else(|| Duration::from_secs(DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS));
loop {
tokio::select! {
event = client_in.next() => {
let Some(event) = event else { break };
for outbound in OPENAI_RESPONSES_WS_CONFIG
.transform_ws_request(&event, model)?
.events
{
let payload = serde_json::to_string(&outbound)
.map_err(|error| Error::InvalidResponse(error.to_string()))?;
upstream_tx.send(Message::Text(payload))
.await
.map_err(|error| Error::Network(error.to_string()))?;
}
}
message = upstream_rx.next() => {
let Some(message) = message else { break };
match message.map_err(|error| Error::Network(error.to_string()))? {
Message::Text(text) => {
let event = serde_json::from_str::<ResponsesWsEvent>(&text)
.map_err(|error| Error::InvalidResponse(error.to_string()))?;
observe(&event);
for outbound in OPENAI_RESPONSES_WS_CONFIG
.transform_ws_response(&event, model)?
.events
{
client_out.send(outbound)
.await
.map_err(|error| Error::Network(error.to_string()))?;
}
}
Message::Close(_) => break,
_ => {}
}
}
_ = tokio::time::sleep(idle) => break,
}
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
pub async fn async_responses_websocket<In, Out>(
model: &str,
api_key: Option<&str>,
api_base: Option<&str>,
first_frame: Option<ResponsesWsEvent>,
idle_timeout: Option<Duration>,
mut observe: impl FnMut(&ResponsesWsEvent) + Send,
client_in: In,
client_out: Out,
) -> Result<(), Error>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
Out::Error: std::fmt::Display,
{
let key = resolve_api_key(api_key)?;
let upstream = dial_upstream(model, &key, api_base).await?;
let (mut upstream_tx, upstream_rx) = upstream.split();
if let Some(first_frame) = first_frame {
for outbound in OPENAI_RESPONSES_WS_CONFIG
.transform_ws_request(&first_frame, model)?
.events
{
let payload = serde_json::to_string(&outbound)
.map_err(|error| Error::InvalidResponse(error.to_string()))?;
upstream_tx
.send(Message::Text(payload))
.await
.map_err(|error| Error::Network(error.to_string()))?;
}
}
ResponsesWebSocketStreaming::bidirectional_forward(
model,
upstream_tx,
upstream_rx,
idle_timeout,
&mut observe,
client_in,
client_out,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub async fn responses_ws<In, Out>(
model: &str,
api_key: Option<&str>,
api_base: Option<&str>,
first_frame: Option<ResponsesWsEvent>,
idle_timeout: Option<Duration>,
observe: impl FnMut(&ResponsesWsEvent) + Send,
client_in: In,
client_out: Out,
) -> Result<(), Error>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
Out::Error: std::fmt::Display,
{
async_responses_websocket(
model,
api_key,
api_base,
first_frame,
idle_timeout,
observe,
client_in,
client_out,
)
.await
}
#[cfg(test)]
mod tests {
use super::*;
use futures_channel::mpsc;
use futures_util::{SinkExt, StreamExt};
use litellm_core::responses::types::ResponsesWsEventType;
use serde_json::json;
use tokio::io::AsyncWriteExt;
use tokio::net::TcpListener;
use tokio_tungstenite::accept_async;
/// The Responses dial has to reach a `wss://` upstream without a process-wide
/// crypto provider installed, which is what dialing through `io::tls` buys.
#[tokio::test]
async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind a loopback port");
let port = listener
.local_addr()
.expect("read the bound address")
.port();
tokio::spawn(async move {
while let Ok((stream, _peer)) = listener.accept().await {
drop(stream);
}
});
let result =
dial_upstream("gpt-5", "sk-test", Some(&format!("wss://127.0.0.1:{port}"))).await;
assert!(matches!(result, Err(Error::Network(_))));
}
async fn websocket_base() -> (String, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let address = listener.local_addr().expect("local address");
let task = tokio::spawn(async move {
let (stream, _) = listener.accept().await.expect("accept");
let mut socket = accept_async(stream).await.expect("websocket handshake");
while let Some(Ok(Message::Text(text))) = socket.next().await {
let request: serde_json::Value = serde_json::from_str(&text).expect("request json");
let model = request
.get("model")
.and_then(serde_json::Value::as_str)
.or_else(|| {
request
.get("response")
.and_then(serde_json::Value::as_object)
.and_then(|response| {
response.get("model").and_then(serde_json::Value::as_str)
})
})
.expect("enforced model");
socket
.send(Message::Text(
json!({
"type": "response.created",
"response": {
"id": format!("resp-{model}"),
"model": model,
"extra": "preserved"
}
})
.to_string(),
))
.await
.expect("created event");
socket
.send(Message::Text(
json!({
"type": "response.completed",
"response": {
"id": format!("resp-{model}"),
"model": model,
"usage": {
"input_tokens": 1,
"output_tokens": 2,
"total_tokens": 3
}
}
})
.to_string(),
))
.await
.expect("completed event");
}
});
(format!("http://{address}"), task)
}
fn event(value: serde_json::Value) -> ResponsesWsEvent {
serde_json::from_value(value).expect("event")
}
#[test]
fn explicit_nonblank_key_wins() {
assert_eq!(
resolve_api_key(Some(" explicit ")).expect("key"),
"explicit"
);
}
#[test]
fn blank_key_is_not_accepted_without_environment_key() {
if std::env::var(OPENAI_API_KEY_ENV).is_err() {
assert!(resolve_api_key(Some(" ")).is_err());
}
}
#[tokio::test]
async fn forwards_events_sequentially_and_enforces_model() {
let (api_base, server) = websocket_base().await;
let (client_tx, client_rx) = mpsc::unbounded();
let (output_tx, mut output_rx) = mpsc::unbounded();
let (observed_tx, observed_rx) = mpsc::unbounded();
client_tx
.unbounded_send(event(json!({
"type": "response.create",
"model": "wrong"
})))
.expect("first request");
client_tx
.unbounded_send(event(json!({
"type": "response.create",
"response": {"model": "also-wrong"}
})))
.expect("second request");
let task = tokio::spawn(async move {
responses_ws(
"authorized-model",
Some("test-key"),
Some(&api_base),
None,
Some(Duration::from_secs(1)),
move |event| {
observed_tx
.unbounded_send(event.clone())
.expect("observe event");
},
client_rx,
output_tx,
)
.await
});
let first = output_rx.next().await.expect("first output");
let second = output_rx.next().await.expect("second output");
let third = output_rx.next().await.expect("third output");
let fourth = output_rx.next().await.expect("fourth output");
drop(client_tx);
task.await.expect("splice task").expect("successful splice");
server.await.expect("server task");
assert_eq!(first.event_type, ResponsesWsEventType::ResponseCreated);
assert_eq!(first.model(), Some("authorized-model"));
assert_eq!(first.data["response"]["extra"], "preserved");
assert_eq!(second.event_type, ResponsesWsEventType::ResponseCompleted);
assert_eq!(third.event_type, ResponsesWsEventType::ResponseCreated);
assert_eq!(fourth.event_type, ResponsesWsEventType::ResponseCompleted);
let observed: Vec<_> = observed_rx.collect().await;
assert_eq!(observed.len(), 4);
assert!(
observed
.iter()
.all(|event| event.event_type != ResponsesWsEventType::ResponseCreate)
);
}
#[tokio::test]
async fn idle_timeout_ends_without_upstream_events() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let address = listener.local_addr().expect("address");
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.expect("accept");
let _socket = accept_async(stream).await.expect("handshake");
tokio::time::sleep(Duration::from_secs(1)).await;
});
let (_client_tx, client_rx) = mpsc::unbounded::<ResponsesWsEvent>();
let (output_tx, mut output_rx) = mpsc::unbounded();
let result = responses_ws(
"model",
Some("key"),
Some(&format!("http://{address}")),
None,
Some(Duration::from_millis(20)),
|_| {},
client_rx,
output_tx,
)
.await;
assert!(result.is_ok());
assert!(output_rx.next().await.is_none());
server.abort();
}
#[tokio::test]
async fn dial_http_status_is_preserved() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let address = listener.local_addr().expect("address");
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.expect("accept");
stream
.write_all(b"HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\n\r\n")
.await
.expect("response");
});
let (_client_tx, client_rx) = mpsc::unbounded::<ResponsesWsEvent>();
let (output_tx, _output_rx) = mpsc::unbounded();
let error = responses_ws(
"model",
Some("key"),
Some(&format!("http://{address}")),
None,
Some(Duration::from_millis(20)),
|_| {},
client_rx,
output_tx,
)
.await
.expect_err("status error");
assert!(matches!(error, Error::Http { status: 401, .. }));
server.await.expect("server task");
}
#[tokio::test]
async fn dial_http_500_status_is_preserved() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let address = listener.local_addr().expect("address");
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.expect("accept");
stream
.write_all(b"HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\n\r\n")
.await
.expect("response");
});
let (_client_tx, client_rx) = mpsc::unbounded::<ResponsesWsEvent>();
let (output_tx, _output_rx) = mpsc::unbounded();
let error = responses_ws(
"model",
Some("key"),
Some(&format!("http://{address}")),
None,
Some(Duration::from_millis(20)),
|_| {},
client_rx,
output_tx,
)
.await
.expect_err("status error");
assert!(matches!(error, Error::Http { status: 500, .. }));
server.await.expect("server task");
}
}

80
litellm-rust/crates/ai-gateway/src/io/tls.rs
View file

@ -1,80 +0,0 @@
//! Outbound WebSocket dials over a TLS config this crate builds once and owns.
//!
//! `reqwest/rustls-tls` enables `rustls/ring` and `litellm-core`'s `bedrock-auth`
//! enables `rustls/aws-lc-rs`, so the bare `ClientConfig::builder()` that
//! `tokio-tungstenite` uses when handed no connector panics rather than guess
//! between them. Naming ring on a connector of our own settles that for these
//! dials without touching the process-wide default, and building the config
//! once keeps the platform trust store, which `tokio-tungstenite` would
//! otherwise re-read on every dial, off the dial path.
use std::io;
use std::sync::{Arc, OnceLock};
use rustls::{ClientConfig, RootCertStore};
use tokio::net::TcpStream;
use tokio_tungstenite::tungstenite::Error;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::error::TlsError;
use tokio_tungstenite::tungstenite::handshake::client::Response;
use tokio_tungstenite::{
Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config,
};
static TLS_CONFIG: OnceLock<Arc<ClientConfig>> = OnceLock::new();
fn build_config() -> Result<ClientConfig, Box<Error>> {
let native = rustls_native_certs::load_native_certs();
let roots = {
let mut store = RootCertStore::empty();
let (added, _ignored) = store.add_parsable_certificates(native.certs);
if added == 0 {
return Err(Box::new(Error::Io(io::Error::other(format!(
"no usable native root certificates: {:?}",
native.errors
)))));
}
store
};
ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider()))
.with_safe_default_protocol_versions()
.map(|builder| builder.with_root_certificates(roots).with_no_client_auth())
.map_err(|error| Box::new(Error::Tls(TlsError::Rustls(error))))
}
fn tls_config() -> Result<Arc<ClientConfig>, Box<Error>> {
if let Some(config) = TLS_CONFIG.get() {
return Ok(Arc::clone(config));
}
let built = Arc::new(build_config()?);
Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built)))
}
pub(crate) async fn connect_upstream<R>(
request: R,
) -> Result<(WebSocketStream<MaybeTlsStream<TcpStream>>, Response), Box<Error>>
where
R: IntoClientRequest + Unpin,
{
let request = request.into_client_request().map_err(Box::new)?;
let connector = match request.uri().scheme_str() {
Some("wss") => Some(Connector::Rustls(tls_config()?)),
_ => None,
};
connect_async_tls_with_config(request, None, false, connector)
.await
.map_err(Box::new)
}
#[cfg(test)]
mod tests {
use super::build_config;
#[test]
fn builds_a_usable_config_with_both_provider_features_enabled() {
let config = build_config().expect("a client config");
assert!(!config.crypto_provider().cipher_suites.is_empty());
}
}

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

@ -1,32 +0,0 @@
//! LiteLLM AI Gateway library.
//!
//! Two layers, split by feature so the Python `cdylib` can depend on the I/O
//! without pulling in the HTTP server:
//!
//! - Call-type modules such as [`ocr`]: provider transforms, lifecycle hooks,
//! and provider I/O. Always available — no feature required. These predate the
//! rule that a route's entrypoint and handler live in `litellm-core` (see
//! `litellm_core::messages`) and move there as they are touched.
//! - [`io`]: compatibility exports and realtime WebSocket splice helpers.
//! - The server modules ([`auth`], [`routes`], [`state`]) and anything pulling
//! `axum` are gated behind the `server` feature, which the `litellm-ai-gateway`
//! binary turns on.
pub mod audio_transcription;
mod client;
pub mod io;
pub mod ocr;
#[cfg(feature = "server")]
pub mod auth;
#[cfg(feature = "server")]
pub mod routes;
#[cfg(feature = "server")]
pub mod state;
#[cfg(feature = "trace-parity")]
pub mod trace_parity;
mod constants;
pub mod integrations;
#[cfg(feature = "server")]
mod realtime;

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

@ -1,162 +0,0 @@
//! LiteLLM AI Gateway — a minimal Axum server fronting the Rust router.
//!
//! Flow: client → `POST /v1/realtime` → `router.realtime()` selects a deployment
//! (simple-shuffle) → `io::realtime::realtime()` invokes OpenAI. The
//! server owns transport + config; routing lives in the `router` crate.
//!
//! The binary requires the `server` feature (declared in `Cargo.toml` via
//! `required-features`), so cargo skips it unless that feature is on. Everything
//! the binary needs lives in the library (`litellm_ai_gateway`); `main` just
//! wires startup.
use std::sync::Arc;
use litellm_ai_gateway::io::realtime_pool::{PoolConfig, RealtimePool, upstream_key};
use litellm_ai_gateway::routes;
use litellm_ai_gateway::state::AppState;
#[cfg(feature = "python-config")]
use litellm_config::load_model_list;
use litellm_core::router::{Deployment, LiteLLMParams, Router};
use litellm_ai_gateway::integrations::custom_logger::CustomLogger;
use litellm_ai_gateway::integrations::litellm_python_proxy_api::LiteLLMPythonProxyAPILogger;
/// Bind to localhost by default so the gateway is not a public, unauthenticated
/// provider proxy out of the box. Override with `HOST` (e.g. `0.0.0.0`).
const DEFAULT_HOST: &str = "127.0.0.1";
const DEFAULT_PORT: u16 = 4001;
#[tokio::main]
async fn main() {
// Trim before storing so it matches the trimmed bearer token in `auth`
// (avoids a silent auth failure when the env var has surrounding whitespace).
let master_key: Option<Arc<str>> = std::env::var("LITELLM_MASTER_KEY")
.ok()
.map(|key| key.trim().to_string())
.filter(|key| !key.is_empty())
.map(Arc::from);
if master_key.is_none() {
eprintln!(
"warning: LITELLM_MASTER_KEY is not set; /v1/realtime will reject all requests (fail closed)"
);
}
// Spawn the realtime-logging worker (drains a channel → POSTs batches to the
// Python proxy's /v1/callbacks/logs). Built here so the spawn lands on the
// tokio runtime. `from_env` reads LITELLM_PROXY_BASE_URL + LITELLM_MASTER_KEY.
let proxy_logger = LiteLLMPythonProxyAPILogger::from_env();
let loggers: Vec<Arc<dyn CustomLogger>> = vec![proxy_logger];
let router = Arc::new(build_router());
// Build the pre-warmed realtime pool and register each deployment's upstream
// so the background replenisher starts warming it. `REALTIME_POOL_SIZE=0`
// yields a disabled pool → every connect fresh-dials (original behavior).
let pool_config = PoolConfig::from_env();
let realtime_pool = RealtimePool::spawn(pool_config);
if pool_config.enabled() {
register_deployments(&router, &realtime_pool);
eprintln!(
"realtime connection pool enabled: target {} warm sockets/key, max idle {}s",
pool_config.target_size,
pool_config.max_idle.as_secs()
);
} else {
eprintln!(
"realtime connection pool disabled (REALTIME_POOL_SIZE=0); fresh-dialing each connect"
);
}
let state = AppState {
router,
master_key,
loggers: Arc::new(loggers),
realtime_pool,
};
let host = std::env::var("HOST").unwrap_or_else(|_| DEFAULT_HOST.to_string());
let port = resolve_port();
let listener = tokio::net::TcpListener::bind((host.as_str(), port))
.await
.expect("failed to bind listener");
eprintln!("litellm-ai-gateway listening on {host}:{port}");
axum::serve(listener, routes::app(state))
.await
.expect("server error");
}
/// Register every deployment's upstream key with the pool so the replenisher
/// pre-warms it. Mirrors `service::run`'s key derivation (strip `openai/`, resolve
/// api_key); deployments whose key can't be resolved are skipped (they fresh-dial
/// and surface the auth error on the request path, as before).
fn register_deployments(router: &Router, pool: &RealtimePool) {
for deployment in router.deployments() {
let params = &deployment.litellm_params;
let provider_model = params
.model
.strip_prefix("openai/")
.unwrap_or(&params.model);
if let Some(key) = upstream_key(
provider_model,
params.api_key.as_deref(),
params.api_base.as_deref(),
) {
pool.register(key);
}
}
}
/// Resolve `PORT`, warning (rather than silently defaulting) on an invalid value.
fn resolve_port() -> u16 {
match std::env::var("PORT") {
Ok(raw) => raw.parse().unwrap_or_else(|_| {
eprintln!("warning: PORT={raw:?} is not a valid port; using {DEFAULT_PORT}");
DEFAULT_PORT
}),
Err(_) => DEFAULT_PORT,
}
}
/// Build the router. With the `python-config` feature and `LITELLM_CONFIG_PATH`
/// set, load the resolved `model_list` from the proxy config via the embedded
/// Python reader (load time only). Otherwise fall back to the env stand-in.
fn build_router() -> Router {
#[cfg(feature = "python-config")]
if let Ok(config_path) = std::env::var("LITELLM_CONFIG_PATH") {
match load_model_list(std::path::Path::new(&config_path)) {
Ok(deployments) => {
eprintln!("loaded model_list from {config_path} via python config reader");
return Router::new(deployments);
}
Err(err) => {
eprintln!("config load failed ({err}); falling back to env deployment");
}
}
}
build_router_from_env()
}
/// Build a minimal single-deployment `model_list` from the environment.
///
/// A real deployment loads `model_list` from config; this is the minimal stand-in
/// so the gateway has one OpenAI deployment to route to.
fn build_router_from_env() -> Router {
let model =
std::env::var("OPENAI_REALTIME_MODEL").unwrap_or_else(|_| "gpt-realtime".to_string());
let api_key = std::env::var("OPENAI_API_KEY").ok();
if api_key.is_none() {
eprintln!(
"warning: OPENAI_API_KEY is not set; realtime requests will fail with auth errors"
);
}
let deployment = Deployment {
model_name: model.clone(),
litellm_params: LiteLLMParams {
model,
api_key,
api_base: None,
},
};
Router::new(vec![deployment])
}

127
litellm-rust/crates/ai-gateway/src/ocr/mod.rs
View file

@ -1,127 +0,0 @@
use litellm_core::Error;
use litellm_core::ocr::{
OcrClient,
wire::{OcrWireRequest, decode_request},
};
use serde_json::Value;
mod types;
pub use types::OcrRequest;
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub async fn ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
core_ocr(request).await
}
async fn core_ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
validate_host_hooks(&request)?;
let client = OcrClient::new(crate::client::http_client().clone())?;
let core_request = decode_request(OcrWireRequest {
model: request.model.to_string(),
document: request.document,
api_key: request.api_key.map(str::to_string),
api_base: request.api_base.map(str::to_string),
custom_llm_provider: request.custom_llm_provider.map(str::to_string),
extra_headers: request.extra_headers,
optional_params: request.optional_params,
input_sources: Default::default(),
timeout_seconds: request.timeout.map(|timeout| timeout.as_secs_f64()),
})?;
client
.perform(core_request)
.await
.map(|response| response.into_json())
}
fn validate_host_hooks(request: &OcrRequest<'_>) -> Result<(), Error> {
if !request.guardrails.is_empty() {
return Err(Error::Unsupported(
"OCR host guardrails are not wired to the core path",
));
}
if !request.callbacks.is_empty() {
return Err(Error::Unsupported(
"OCR host callbacks are not wired to the core path",
));
}
Ok(())
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use litellm_core::ocr::wire::is_supported_request;
use serde_json::{Map, json};
use super::{OcrRequest, validate_host_hooks};
use crate::integrations::custom_guardrail::{CustomGuardrail, GuardrailEventHook};
use crate::integrations::custom_logger::CustomLogger;
struct TestGuardrail;
impl CustomGuardrail for TestGuardrail {
fn guardrail_name(&self) -> &str {
"test"
}
fn supported_event_hooks(&self) -> &[GuardrailEventHook] {
&[]
}
}
struct TestLogger;
impl CustomLogger for TestLogger {}
fn request() -> OcrRequest<'static> {
OcrRequest {
model: "model",
document: json!({"type":"image_url","image_url":"data:image/png;base64,YQ=="}),
api_key: None,
api_base: None,
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
timeout: None,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: Default::default(),
litellm_call_id: None,
}
}
#[test]
fn core_activation_includes_migrated_providers() {
assert!(is_supported_request("model", Some("mistral")));
assert!(is_supported_request("pixtral-12b", Some("azure_ai")));
assert!(is_supported_request(
"doc-intelligence/prebuilt-layout",
Some("azure_ai")
));
assert!(is_supported_request("parse-v3", Some("reducto")));
assert!(is_supported_request("mistral-ocr", Some("vertex_ai")));
assert!(is_supported_request("deepseek-ocr", Some("vertex_ai")));
}
#[test]
fn core_path_rejects_unwired_guardrails() {
let request = OcrRequest {
guardrails: vec![Arc::new(TestGuardrail)],
..request()
};
let error = validate_host_hooks(&request).unwrap_err();
assert!(error.to_string().contains("guardrails are not wired"));
}
#[test]
fn core_path_rejects_unwired_callbacks() {
let request = OcrRequest {
callbacks: vec![Arc::new(TestLogger)],
..request()
};
let error = validate_host_hooks(&request).unwrap_err();
assert!(error.to_string().contains("callbacks are not wired"));
}
}

23
litellm-rust/crates/ai-gateway/src/ocr/types.rs
View file

@ -1,23 +0,0 @@
use std::sync::Arc;
use std::time::Duration;
use serde_json::{Map, Value};
use crate::integrations::custom_guardrail::CustomGuardrail;
use crate::integrations::custom_logger::CustomLogger;
use crate::integrations::types::RequestMetadata;
pub struct OcrRequest<'a> {
pub model: &'a str,
pub document: Value,
pub api_key: Option<&'a str>,
pub api_base: Option<&'a str>,
pub custom_llm_provider: Option<&'a str>,
pub extra_headers: Option<Map<String, Value>>,
pub optional_params: Map<String, Value>,
pub timeout: Option<Duration>,
pub callbacks: Vec<Arc<dyn CustomLogger>>,
pub guardrails: Vec<Arc<dyn CustomGuardrail>>,
pub request_metadata: RequestMetadata,
pub litellm_call_id: Option<&'a str>,
}

4
litellm-rust/crates/ai-gateway/src/realtime/mod.rs
View file

@ -1,4 +0,0 @@
//! Realtime logging collector. Observes the realtime event stream and emits a
//! `StandardLoggingPayload` to the registered callbacks on session close.
pub mod streaming;

414
litellm-rust/crates/ai-gateway/src/realtime/streaming.rs
View file

@ -1,414 +0,0 @@
//! `RealTimeStreaming` — the realtime logging collector.
//!
//! Mirrors Python `litellm.realtime_api.main.RealTimeStreaming`: it observes the
//! event stream in O(1) (never buffering frames), accumulating just the fields
//! the spend log needs (model, id, cumulative usage), then on session close
//! builds a `StandardLoggingPayload` and fans it out to every registered
//! `CustomLogger`.
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use litellm_core::realtime::types::RealtimeEvent;
use serde_json::Value;
use crate::constants::DEFAULT_PROVIDER;
use crate::integrations::custom_logger::{
CallbackTiming, CallbackValue, CustomLogger, CustomLoggerRunner, LoggingError, ModelCallDetails,
};
use crate::integrations::types::{
RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload, Usage,
};
/// Current wall-clock time as epoch seconds (float), matching the Python
/// `startTime`/`endTime` contract.
fn epoch_seconds() -> f64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs_f64())
.unwrap_or(0.0)
}
/// Status of a finished realtime session, mapped to the callback record status.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SessionStatus {
Success,
Failure,
}
/// Accumulates realtime session state and emits a logging payload on close.
pub struct RealTimeStreaming {
callbacks: Vec<Arc<dyn CustomLogger>>,
/// REQUEST-ID RULE: the SpendLogs `request_id` == the OpenAI realtime session
/// id (`sess_…`), captured from `session.created`. Both `id` and
/// `litellm_call_id` are set to that value so the Python writer logs the same
/// id regardless of which field it reads. The gateway-generated `rt-…` id
/// (the constructor seed) is only a fallback for sessions that fail before
/// `session.created` arrives.
litellm_call_id: String,
/// See the request-id rule above — mirrors `litellm_call_id`.
id: String,
model: String,
custom_llm_provider: String,
usage: Usage,
response_cost: f64,
start_time: f64,
end_time: f64,
metadata: RequestMetadata,
/// Count of logging callbacks that failed to enqueue (non-fatal).
dropped: u64,
}
impl RealTimeStreaming {
/// Create a collector for one session. `litellm_call_id` is the gateway's
/// per-connection id; `model` is the requested model (a sane default until
/// `session.created` reports the upstream model).
pub fn new(
callbacks: Vec<Arc<dyn CustomLogger>>,
litellm_call_id: String,
model: String,
metadata: RequestMetadata,
) -> Self {
let now = epoch_seconds();
Self {
callbacks,
id: litellm_call_id.clone(),
litellm_call_id,
model,
custom_llm_provider: DEFAULT_PROVIDER.to_string(),
usage: Usage::default(),
response_cost: 0.0,
start_time: now,
end_time: now,
metadata,
dropped: 0,
}
}
/// Number of logging callbacks that failed to enqueue so far (test/observ.).
#[allow(dead_code)]
pub fn dropped(&self) -> u64 {
self.dropped
}
/// Observe one realtime event. O(1): updates accumulated state only; never
/// buffers frames. Safe to call on every event in either direction.
pub fn observe(&mut self, event: &RealtimeEvent) {
match event.event_type.as_str() {
"session.created" | "session.updated" => self.on_session(event),
"response.done" => self.on_response_done(event),
_ => {}
}
}
/// `session.created` / `session.updated` → capture upstream id + model.
/// Per the request-id rule, the OpenAI session id becomes BOTH `id` and
/// `litellm_call_id`, replacing the gateway-generated fallback.
fn on_session(&mut self, event: &RealtimeEvent) {
let session = event.data.get("session").and_then(Value::as_object);
if let Some(id) = session.and_then(|s| s.get("id")).and_then(Value::as_str)
&& !id.is_empty()
{
self.id = id.to_string();
self.litellm_call_id = id.to_string();
}
if let Some(model) = session.and_then(|s| s.get("model")).and_then(Value::as_str)
&& !model.is_empty()
{
self.model = model.to_string();
}
}
/// `response.done` → add this response's usage to the cumulative totals.
fn on_response_done(&mut self, event: &RealtimeEvent) {
let usage = event
.data
.get("response")
.and_then(Value::as_object)
.and_then(|r| r.get("usage"))
.and_then(Value::as_object);
let Some(usage) = usage else { return };
let input = usage.get("input_tokens").and_then(Value::as_u64);
let output = usage.get("output_tokens").and_then(Value::as_u64);
let total = usage.get("total_tokens").and_then(Value::as_u64);
if let Some(input) = input {
self.usage.prompt_tokens += input;
}
if let Some(output) = output {
self.usage.completion_tokens += output;
}
// Prefer the upstream-reported total; otherwise derive it.
match total {
Some(total) => self.usage.total_tokens += total,
None => {
self.usage.total_tokens += input.unwrap_or(0) + output.unwrap_or(0);
}
}
}
/// Set the per-session response cost ($). Cost computation is Python-side in
/// the proxy; the gateway forwards 0.0 by default and lets the proxy price.
/// Public API (exercised in tests) for the future path where the gateway
/// prices realtime sessions itself.
#[allow(dead_code)]
pub fn set_response_cost(&mut self, cost: f64) {
self.response_cost = cost;
}
/// Build the `StandardLoggingPayload` from accumulated state.
pub fn build_payload(&self) -> StandardLoggingPayload {
StandardLoggingPayload {
id: self.id.clone(),
litellm_call_id: self.litellm_call_id.clone(),
call_type: "realtime".to_string(),
model: self.model.clone(),
custom_llm_provider: self.custom_llm_provider.clone(),
response_cost: self.response_cost,
prompt_tokens: self.usage.prompt_tokens,
completion_tokens: self.usage.completion_tokens,
total_tokens: self.usage.total_tokens,
start_time: self.start_time,
end_time: self.end_time,
stream: true,
metadata: StandardLoggingMetadata {
user_api_key_hash: self.metadata.user_api_key_hash.clone(),
user_api_key_user_id: self.metadata.user_api_key_user_id.clone(),
user_api_key_team_id: self.metadata.user_api_key_team_id.clone(),
..Default::default()
},
messages: None,
}
}
/// Finish the session: stamp the end time and fan the payload out to every
/// callback. On a logger enqueue error we bump a non-fatal counter (the
/// realtime session has already ended; a dropped log must never propagate).
pub async fn log_messages(&mut self, status: SessionStatus) {
self.end_time = epoch_seconds();
let payload = self.build_payload();
let timing = CallbackTiming::new(payload.start_time, payload.end_time);
let runner = CustomLoggerRunner::new(self.callbacks.clone());
match status {
SessionStatus::Success => {
let response = CallbackValue::new("realtime", serde_json::Value::Null);
let report = runner
.async_log_success_event(
&ModelCallDetails::from_standard_logging_payload(payload),
&response,
timing,
)
.await;
self.dropped += report.dropped as u64;
}
SessionStatus::Failure => {
let error = LoggingError {
message: "realtime session ended in failure".to_string(),
kind: "RealtimeSessionError".to_string(),
};
let response = CallbackValue::new(
"error",
serde_json::json!({
"message": error.message,
"kind": error.kind,
}),
);
let report = runner
.async_log_failure_event(
&ModelCallDetails::from_standard_logging_payload(payload)
.with_failure_error(error),
Some(&response),
timing,
)
.await;
self.dropped += report.dropped as u64;
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::integrations::custom_logger::LogError;
use crate::integrations::custom_logger::LogFuture;
use std::sync::atomic::{AtomicU64, Ordering};
fn event(raw: &str) -> RealtimeEvent {
serde_json::from_str(raw).expect("valid event json")
}
/// A test logger that records the last payload it saw.
#[derive(Default)]
struct CapturingLogger {
calls: AtomicU64,
last_model: std::sync::Mutex<Option<String>>,
last_total_tokens: AtomicU64,
}
impl CustomLogger for CapturingLogger {
fn async_log_success_event<'a>(
&'a self,
model_call_details: &'a ModelCallDetails,
_response_obj: &'a CallbackValue,
_timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async move {
let payload = model_call_details
.standard_logging_payload
.as_ref()
.expect("standard logging payload");
self.calls.fetch_add(1, Ordering::SeqCst);
*self.last_model.lock().unwrap() = Some(payload.model.clone());
self.last_total_tokens
.store(payload.total_tokens, Ordering::SeqCst);
Ok(())
})
}
}
#[tokio::test]
async fn observe_accumulates_model_and_tokens_then_logs() {
let logger = Arc::new(CapturingLogger::default());
let callbacks: Vec<Arc<dyn CustomLogger>> = vec![logger.clone()];
let mut streaming = RealTimeStreaming::new(
callbacks,
"call_abc".to_string(),
"gpt-realtime".to_string(),
RequestMetadata {
user_api_key_hash: Some("hash123".to_string()),
user_api_key_user_id: Some("user-1".to_string()),
user_api_key_team_id: Some("team-1".to_string()),
},
);
streaming.observe(&event(
r#"{"type":"session.created","session":{"id":"sess_001","model":"gpt-realtime-2025"}}"#,
));
streaming.observe(&event(
r#"{"type":"response.done","response":{"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}}"#,
));
// A second response.done accumulates.
streaming.observe(&event(
r#"{"type":"response.done","response":{"usage":{"input_tokens":3,"output_tokens":2,"total_tokens":5}}}"#,
));
let payload = streaming.build_payload();
assert_eq!(payload.model, "gpt-realtime-2025");
// Request-id rule: session.created's id becomes BOTH id and
// litellm_call_id (replacing the "call_abc" gateway fallback), so the
// SpendLogs request_id is always the OpenAI session id.
assert_eq!(payload.id, "sess_001");
assert_eq!(payload.litellm_call_id, "sess_001");
assert_eq!(payload.prompt_tokens, 13);
assert_eq!(payload.completion_tokens, 7);
assert_eq!(payload.total_tokens, 20);
assert_eq!(payload.response_cost, 0.0);
assert_eq!(payload.call_type, "realtime");
assert_eq!(payload.custom_llm_provider, "openai");
assert_eq!(
payload.metadata.user_api_key_hash.as_deref(),
Some("hash123")
);
streaming.log_messages(SessionStatus::Success).await;
assert_eq!(logger.calls.load(Ordering::SeqCst), 1);
assert_eq!(
logger.last_model.lock().unwrap().as_deref(),
Some("gpt-realtime-2025")
);
assert_eq!(logger.last_total_tokens.load(Ordering::SeqCst), 20);
assert_eq!(streaming.dropped(), 0);
}
#[test]
fn blank_session_id_and_model_keep_the_gateway_fallbacks() {
let mut streaming = RealTimeStreaming::new(
Vec::new(),
"call_fallback".to_string(),
"gpt-realtime".to_string(),
RequestMetadata::default(),
);
streaming.observe(&event(
r#"{"type":"session.created","session":{"id":"","model":""}}"#,
));
let payload = streaming.build_payload();
assert_eq!(payload.id, "call_fallback");
assert_eq!(payload.litellm_call_id, "call_fallback");
assert_eq!(payload.model, "gpt-realtime");
streaming.observe(&event(
r#"{"type":"session.updated","session":{"id":"sess_002","model":""}}"#,
));
let payload = streaming.build_payload();
assert_eq!(payload.id, "sess_002");
assert_eq!(payload.litellm_call_id, "sess_002");
assert_eq!(payload.model, "gpt-realtime");
}
#[test]
fn payload_serializes_with_camelcase_times_and_realtime_call_type() {
let mut streaming = RealTimeStreaming::new(
Vec::new(),
"call_xyz".to_string(),
"gpt-realtime".to_string(),
RequestMetadata::default(),
);
streaming.observe(&event(
r#"{"type":"response.done","response":{"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}"#,
));
streaming.set_response_cost(0.0042);
let payload = streaming.build_payload();
let json = serde_json::to_string(&payload).expect("serialize payload");
assert!(json.contains("\"startTime\""), "missing startTime: {json}");
assert!(json.contains("\"endTime\""), "missing endTime: {json}");
assert!(
json.contains("\"call_type\":\"realtime\""),
"missing call_type realtime: {json}"
);
assert!(
json.contains("\"response_cost\""),
"missing response_cost: {json}"
);
assert_eq!(payload.response_cost, 0.0042);
}
/// A logger whose enqueue always fails should bump the dropped counter, not
/// panic or propagate.
#[tokio::test]
async fn failing_logger_bumps_dropped_counter() {
struct FailingLogger;
impl CustomLogger for FailingLogger {
fn async_log_success_event<'a>(
&'a self,
_model_call_details: &'a ModelCallDetails,
_response_obj: &'a CallbackValue,
_timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async { Err(LogError::channel_full()) })
}
fn async_log_failure_event<'a>(
&'a self,
_model_call_details: &'a ModelCallDetails,
_response_obj: Option<&'a CallbackValue>,
_timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async { Err(LogError::channel_closed()) })
}
}
let callbacks: Vec<Arc<dyn CustomLogger>> = vec![Arc::new(FailingLogger)];
let mut streaming = RealTimeStreaming::new(
callbacks,
"call_1".to_string(),
"gpt-realtime".to_string(),
RequestMetadata::default(),
);
streaming.log_messages(SessionStatus::Success).await;
assert_eq!(streaming.dropped(), 1);
}
}

43
litellm-rust/crates/ai-gateway/src/routes/AGENTS.md
View file

@ -1,43 +0,0 @@
# routes/ — the route template
Every route follows the **same shape** so the layout is predictable. The rule:
> **Each route module exposes `pub fn router() -> Router<AppState>`.**
> `routes/mod.rs::app` merges them all and applies state once. Adding a route is:
> create the module, then add one `.merge(<name>::router())` line.
## Default: one file
A route is a single file containing `router()` + its handler(s) (handlers stay
private). This is the norm — don't split until it hurts.
```
pub fn router() -> Router<AppState> { Router::new().route(PATH, get(handle)) }
async fn handle(...) -> impl IntoResponse { ... }
```
`health.rs` is the example.
## Split out `service` when there's real logic
When a route has business logic worth testing without axum, put it in a sibling
`service` (a file, or a folder if the route grows). The route file stays the
**axum surface** (router + handler + any socket/SSE adapter); `service` is plain
Rust with **no axum types**, and its job is to pick the deployment and call the
`core` route entrypoint (see `messages/service.rs` calling
`litellm_core::messages::messages`). Never build a provider request, resolve a
key, or perform the provider call here. `realtime/` is the older example:
```
realtime/
mod.rs # axum surface: router() + handler + the WS<->events adapter
service.rs # pure logic: select deployment + call provider (no axum) — testable
```
Split `service` further (or add `transport`, `repo`, …) only once a single file
genuinely gets hard to read.
## Invariants
- **Auth is an extractor, not a manual call.** A handler requires auth by adding
`crate::auth::RequireMasterKey` to its arguments; it runs during extraction.
Never re-implement the check per route.
- **Handlers contain no business logic; `service` contains no axum types.**
- **No provider handlers in this crate.** Transforms, auth headers, and the
provider HTTP call live in `core/src/<route>/`.
- A route owns its paths in its own `router()`; `mod.rs` only merges.
- Cross-cutting concerns (logging, CORS, timeouts) → Tower layers in `mod.rs`,
not duplicated in handlers.

24
litellm-rust/crates/ai-gateway/src/routes/health.rs
View file

@ -1,24 +0,0 @@
//! Health probes. Simple-route template: a `router()` plus its handlers, in one file.
use axum::Router;
use axum::http::StatusCode;
use axum::routing::get;
use crate::state::AppState;
/// This route's contribution to the app router.
pub fn router() -> Router<AppState> {
Router::new()
.route("/health/liveness", get(liveness))
.route("/health/readiness", get(readiness))
}
/// The process is up.
async fn liveness() -> StatusCode {
StatusCode::OK
}
/// The server is ready to accept traffic.
async fn readiness() -> StatusCode {
StatusCode::OK
}

532
litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs
View file

@ -1,532 +0,0 @@
//! `POST /v1/messages`, the Anthropic Messages HTTP surface.
mod service;
use axum::Router;
use axum::body::Body;
use axum::extract::{Json, State};
use axum::http::StatusCode;
use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE, HeaderMap, HeaderValue};
use axum::response::{IntoResponse, Response};
use axum::routing::post;
use litellm_core::Error;
use serde_json::{Map, Value};
use crate::auth::RequireMasterKey;
use crate::constants::{MESSAGES_HEADERS_NOT_FORWARDED, MESSAGES_ROUTE_PATH};
use crate::state::AppState;
/// This route's contribution to the app router.
pub fn router() -> Router<AppState> {
Router::new().route(MESSAGES_ROUTE_PATH, post(handle))
}
#[tracing::instrument(
name = "messages_gateway_route",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
async fn handle(
_auth: RequireMasterKey,
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> Result<Response, MessagesRouteError> {
let extra_headers = forwarded_headers(&headers)?;
match service::run(&state.router, body, extra_headers)
.await
.map_err(MessagesRouteError::from)?
{
service::MessagesResponse::Json(body) => Ok(Json(body).into_response()),
service::MessagesResponse::Stream(upstream) => stream_response(upstream),
}
}
fn stream_response(upstream: reqwest::Response) -> Result<Response, MessagesRouteError> {
let content_type = upstream
.headers()
.get(CONTENT_TYPE)
.cloned()
.unwrap_or_else(|| HeaderValue::from_static("text/event-stream"));
let mut response = Response::builder()
.status(
StatusCode::from_u16(upstream.status().as_u16()).map_err(|error| {
MessagesRouteError(Error::InvalidResponse(format!(
"invalid upstream response status: {error}"
)))
})?,
)
.header(CONTENT_TYPE, content_type);
if let Some(value) = upstream.headers().get(CACHE_CONTROL) {
response = response.header(CACHE_CONTROL, value);
}
response
.body(Body::from_stream(upstream.bytes_stream()))
.map_err(|error| {
MessagesRouteError(Error::InvalidResponse(format!(
"failed to build streaming response: {error}"
)))
})
}
fn forwarded_headers(headers: &HeaderMap) -> Result<Option<Map<String, Value>>, Error> {
let forwarded = headers
.iter()
.filter(|(name, _)| {
!MESSAGES_HEADERS_NOT_FORWARDED
.iter()
.any(|excluded| name.as_str().eq_ignore_ascii_case(excluded))
})
.map(|(name, value)| {
let value = value.to_str().map_err(|_| {
Error::InvalidRequest(format!("invalid value for header {}", name.as_str()))
})?;
Ok((name.to_string(), Value::String(value.to_string())))
})
.collect::<Result<Map<_, _>, Error>>()?;
Ok((!forwarded.is_empty()).then_some(forwarded))
}
#[derive(Debug)]
struct MessagesRouteError(Error);
impl From<Error> for MessagesRouteError {
fn from(error: Error) -> Self {
Self(error)
}
}
impl IntoResponse for MessagesRouteError {
fn into_response(self) -> Response {
let (status, message) = match self.0 {
Error::InvalidRequest(message) => (StatusCode::BAD_REQUEST, message),
Error::InvalidProvider(_) | Error::Routing(_) => (
StatusCode::NOT_FOUND,
"no messages deployment is configured for this model".to_string(),
),
Error::Auth(_)
| Error::MissingApiKey { .. }
| Error::MissingAzureAiCredentials
| Error::MissingAzureDocumentIntelligenceCredentials
| Error::MissingReductoApiKey => (
StatusCode::BAD_GATEWAY,
"messages provider authentication failed".to_string(),
),
Error::Http { .. }
| Error::Network(_)
| Error::Connect(_)
| Error::InvalidResponse(_)
| Error::InvalidType { .. }
| Error::MissingField(_)
| Error::MissingDocumentUrl => (
StatusCode::BAD_GATEWAY,
"messages provider request failed".to_string(),
),
// The gateway has no Python implementation to decline to, so a
// request the core cannot serve is reported to the caller. The
// reason is a fixed internal string, never provider content.
Error::Unsupported(reason) => (
StatusCode::BAD_REQUEST,
format!("messages request is not supported: {reason}"),
),
};
(
status,
Json(serde_json::json!({"error": {"message": message}})),
)
.into_response()
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use axum::body::Body;
use axum::http::Request;
use axum::http::StatusCode;
use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE};
use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter};
use serde_json::json;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tower::ServiceExt;
use super::super::app;
use crate::io::realtime_pool::RealtimePool;
use crate::state::AppState;
fn state(model: &str, api_base: String, master_key: Option<&str>) -> AppState {
state_with_provider(model, model, api_base, master_key)
}
fn state_with_provider(
model_alias: &str,
provider_model: &str,
api_base: String,
master_key: Option<&str>,
) -> AppState {
AppState {
router: Arc::new(ModelRouter::new(vec![Deployment {
model_name: model_alias.to_string(),
litellm_params: LiteLLMParams {
model: format!("anthropic/{provider_model}"),
api_key: Some("upstream-key".to_string()),
api_base: Some(api_base),
},
}])),
master_key: master_key.map(Arc::from),
loggers: Arc::new(Vec::new()),
realtime_pool: RealtimePool::disabled(),
}
}
async fn upstream(listener: TcpListener) -> (String, tokio::task::JoinHandle<String>) {
let address = listener.local_addr().expect("listener has address");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let mut request = Vec::new();
let mut buffer = [0_u8; 4096];
loop {
let read = socket.read(&mut buffer).await.expect("reads request");
request.extend_from_slice(&buffer[..read]);
if request.windows(4).any(|window| window == b"\r\n\r\n") {
break;
}
}
let request = String::from_utf8(request).expect("request is utf8");
let content_length = request
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
let header_end = request.find("\r\n\r\n").expect("request has headers") + 4;
let mut full_request = request.into_bytes();
while full_request.len().saturating_sub(header_end) < content_length {
let read = socket.read(&mut buffer).await.expect("reads body");
full_request.extend_from_slice(&buffer[..read]);
}
let request = String::from_utf8(full_request).expect("request is utf8");
let body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-test"}"#;
let response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
body.len(),
body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes response");
request
});
(format!("http://{address}"), server)
}
async fn streaming_upstream(
listener: TcpListener,
status: u16,
content_type: &'static str,
body: &'static str,
) -> (String, tokio::task::JoinHandle<String>) {
let address = listener.local_addr().expect("listener has address");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let mut request = Vec::new();
let mut buffer = [0_u8; 4096];
loop {
let read = socket.read(&mut buffer).await.expect("reads request");
request.extend_from_slice(&buffer[..read]);
if request.windows(4).any(|window| window == b"\r\n\r\n") {
break;
}
}
let request_text = String::from_utf8(request).expect("request is utf8");
let content_length = request_text
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
let header_end = request_text.find("\r\n\r\n").expect("request has headers") + 4;
let mut full_request = request_text.into_bytes();
while full_request.len().saturating_sub(header_end) < content_length {
let read = socket.read(&mut buffer).await.expect("reads body");
full_request.extend_from_slice(&buffer[..read]);
}
let response = format!(
"HTTP/1.1 {status} OK\r\ncontent-type: {content_type}\r\ncache-control: no-cache\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
body.len()
);
socket
.write_all(response.as_bytes())
.await
.expect("writes response");
String::from_utf8(full_request).expect("request is utf8")
});
(format!("http://{address}"), server)
}
#[tokio::test]
async fn route_constructs_anthropic_upstream_request() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let (api_base, server) = upstream(listener).await;
let app = app(state("claude-test", api_base, Some("master-key")));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer master-key")
.header("x-api-key", "request-upstream-key")
.header("anthropic-beta", "beta-feature")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "claude-test",
"max_tokens": 16,
"messages": [{"role": "user", "content": "hello"}]
})
.to_string(),
))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::OK);
let body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body reads");
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&body).expect("json")["id"],
"msg_1"
);
let upstream_request = server.await.expect("upstream task completes");
let (head, body) = upstream_request
.split_once("\r\n\r\n")
.expect("upstream request has body");
let head = head.to_ascii_lowercase();
assert!(head.contains("x-api-key: request-upstream-key"));
assert!(head.contains("anthropic-beta: beta-feature"));
assert!(!head.contains("authorization: bearer master-key"));
let body: serde_json::Value = serde_json::from_str(body).expect("upstream body is json");
assert_eq!(body["model"], "claude-test");
assert_eq!(body["messages"][0]["content"], "hello");
}
#[tokio::test]
async fn route_substitutes_model_alias_with_provider_model_upstream() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let (api_base, server) = upstream(listener).await;
let app = app(state_with_provider(
"production",
"claude-sonnet-4-5",
api_base,
Some("master-key"),
));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer master-key")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "production",
"max_tokens": 16,
"messages": [{"role": "user", "content": "hello"}]
})
.to_string(),
))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::OK);
let upstream_request = server.await.expect("upstream task completes");
let (_, upstream_body) = upstream_request
.split_once("\r\n\r\n")
.expect("upstream request has body");
let upstream_body: serde_json::Value =
serde_json::from_str(upstream_body).expect("upstream body is json");
assert_eq!(upstream_body["model"], "claude-sonnet-4-5");
assert_ne!(upstream_body["model"], "production");
}
#[tokio::test]
async fn route_streams_anthropic_events_without_buffering_or_reordering() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let events = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
let (api_base, server) =
streaming_upstream(listener, 200, "text/event-stream", events).await;
let app = app(state("claude-test", api_base, Some("master-key")));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer master-key")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "claude-test",
"max_tokens": 16,
"stream": true,
"messages": [{"role": "user", "content": "hello"}]
})
.to_string(),
))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(CONTENT_TYPE)
.unwrap()
.to_str()
.unwrap(),
"text/event-stream"
);
assert_eq!(
response
.headers()
.get(CACHE_CONTROL)
.unwrap()
.to_str()
.unwrap(),
"no-cache"
);
let response_body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body reads");
assert_eq!(response_body, events.as_bytes());
let upstream_request = server.await.expect("upstream task completes");
let (_, upstream_body) = upstream_request
.split_once("\r\n\r\n")
.expect("upstream request has body");
assert_eq!(
serde_json::from_str::<serde_json::Value>(upstream_body)
.expect("upstream body is json")["stream"],
true
);
}
#[tokio::test]
async fn route_maps_streaming_upstream_errors_before_starting_response() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let (api_base, server) = streaming_upstream(
listener,
429,
"application/json",
r#"{"error":"rate limited"}"#,
)
.await;
let app = app(state("claude-test", api_base, Some("master-key")));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer master-key")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "claude-test",
"max_tokens": 16,
"stream": true,
"messages": [{"role": "user", "content": "hello"}]
})
.to_string(),
))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
let response_body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body reads");
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&response_body).expect("error is json")["error"]
["message"],
"messages provider request failed"
);
server.await.expect("upstream task completes");
}
#[tokio::test]
async fn route_rejects_missing_master_key() {
let app = app(state(
"claude-test",
"http://127.0.0.1:1".to_string(),
Some("master-key"),
));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("content-type", "application/json")
.body(Body::from("{}"))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
#[tokio::test]
async fn route_rejects_invalid_master_key() {
let app = app(state(
"claude-test",
"http://127.0.0.1:1".to_string(),
Some("master-key"),
));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer wrong-key")
.header("content-type", "application/json")
.body(Body::from("{}"))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
#[tokio::test]
async fn route_rejects_malformed_json_without_panicking() {
let app = app(state(
"claude-test",
"http://127.0.0.1:1".to_string(),
Some("master-key"),
));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer master-key")
.header("content-type", "application/json")
.body(Body::from("{not-json"))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
}
}

71
litellm-rust/crates/ai-gateway/src/routes/messages/service.rs
View file

@ -1,71 +0,0 @@
use std::sync::Arc;
use litellm_core::Error;
use litellm_core::constants::ANTHROPIC_MESSAGES_PROVIDER;
use litellm_core::messages::types::MessagesRequest;
use litellm_core::messages::{messages, messages_stream};
use litellm_core::router::Router;
use serde_json::{Map, Value};
pub(crate) enum MessagesResponse {
Json(Value),
Stream(reqwest::Response),
}
#[tracing::instrument(
name = "messages_gateway_service",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
pub async fn run(
router: &Arc<Router>,
body: Value,
extra_headers: Option<Map<String, Value>>,
) -> Result<MessagesResponse, Error> {
let model = body
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|model| !model.is_empty())
.ok_or_else(|| Error::InvalidRequest("messages body requires a model".to_string()))?;
let deployment = router
.get_available_deployment(model)
.ok_or_else(|| Error::Routing(format!("no deployment available for model '{model}'")))?;
let provider_model = deployment.litellm_params.model.as_str();
let upstream_model = provider_model
.split_once('/')
.map_or(provider_model, |(_, model)| model);
let custom_llm_provider = if provider_model.contains('/') {
None
} else {
Some(ANTHROPIC_MESSAGES_PROVIDER)
};
let mut body = body;
body.as_object_mut()
.ok_or_else(|| Error::InvalidRequest("messages body must be an object".to_string()))?
.insert(
"model".to_string(),
Value::String(upstream_model.to_string()),
);
let request = MessagesRequest {
model: provider_model,
body,
api_key: deployment.litellm_params.api_key.as_deref(),
api_base: deployment.litellm_params.api_base.as_deref(),
custom_llm_provider,
extra_headers,
timeout: None,
};
if request.body.get("stream").and_then(Value::as_bool) == Some(true) {
return messages_stream(request).await.map(MessagesResponse::Stream);
}
let response = messages(request).await?;
serde_json::to_value(response)
.map(MessagesResponse::Json)
.map_err(|err| {
Error::InvalidResponse(format!("failed to serialize messages response: {err}"))
})
}

25
litellm-rust/crates/ai-gateway/src/routes/mod.rs
View file

@ -1,25 +0,0 @@
//! HTTP routes.
//!
//! **Template:** every route module exposes `pub fn router() -> Router<AppState>`
//! that mounts its own paths; [`app`] merges them. A trivial route is a single
//! file (`health.rs`); a non-trivial one is a folder (`realtime/`) with
//! `handler` (entry) + `service` (logic) + `transport` (adapters). See AGENTS.md.
pub mod health;
pub mod messages;
pub mod realtime;
pub mod responses;
use axum::Router;
use crate::state::AppState;
/// Assemble the application router by merging every route module's `router()`.
pub fn app(state: AppState) -> Router {
Router::new()
.merge(health::router())
.merge(messages::router())
.merge(realtime::router())
.merge(responses::router())
.with_state(state)
}

87
litellm-rust/crates/ai-gateway/src/routes/realtime/README.md
View file

@ -1,87 +0,0 @@
# Realtime route (`GET /v1/realtime`)
Proxies OpenAI's realtime WebSocket. `mod.rs` is the axum surface (handler +
socket↔events adapter); `service.rs` is the pure logic (select a deployment, then
splice client ↔ upstream). The pool itself lives in
`crates/providers/src/realtime_pool.rs`.
## Connection pooling
### The problem
The gateway's realtime overhead lives **entirely in session establishment**. On each
client connect it dials a *fresh* upstream WS to OpenAI and waits for
`session.created` before it can serve. Measured at 5000 calls / 500 concurrency, the
fresh-dial session phase is **~360 ms** vs **~7 ms** direct; dial, first-audio, and
streaming add ~0. So the one lever is removing that per-connect handshake from the
critical path.
### The idea
Keep a few upstream OpenAI sockets **already connected and already past
`session.created`** (buffered). On a client connect, hand off a warm socket — relay
its buffered `session.created` instantly (a local `Vec::pop`, sub-millisecond) and
splice exactly as a fresh dial would. A background task keeps the pool topped up. On
a miss or dead socket we fall back to fresh-dial: the pool is a latency optimization,
never a correctness dependency.
```
┌───────────────────────────────────────┐
client connect ──────► │ routes/realtime → service::run │
│ pool.take(key) │
│ hit → relay buffered │
│ session.created, then splice │
│ miss → fresh dial (original path) │
└───────────────┬───────────────────────┘
│ replenish (async, concurrent)
┌───────────────▼───────────────────────┐
background task ─────► │ RealtimePool: per-key warm sockets │
│ each = { ws, buffered session.created}│
│ liveness-checked before handoff │
└─────────────────────────────────────────┘
```
A warm session is indistinguishable from a fresh one: OpenAI sends `session.created`
unprompted on connect, we pre-read exactly that one frame and relay it on handoff,
and we send nothing else on the socket before a client exists — so the client's first
`session.update` behaves identically either way.
### Sizing
Each warm socket serves **exactly one** session (realtime isn't multiplexed), so the
pool is sized to the **peak concurrent connects per instance**, not total live
connections:
```
REALTIME_POOL_SIZE ≈ peak_concurrency / instance_count
```
e.g. 500 concurrency over 10 instances → ~50–64 per instance. The replenisher dials
the missing sockets **concurrently**, so a drained pool refills in ~one handshake
window and keeps supply close to the connect rate. Over-provisioning just burns idle
upstream sockets, which is why warm sockets are short-lived
(`REALTIME_POOL_MAX_IDLE_SECS`).
### Config
| env | default | meaning |
| ----------------------------- | ------- | --------------------------------------------------------------- |
| `REALTIME_POOL_SIZE` | `4` | target warm sockets per key. `0` disables pooling (fresh-dial). |
| `REALTIME_POOL_MAX_IDLE_SECS` | `30` | max time a warm socket sits before it's closed and replaced. |
### Notes
- **Miss / dead socket → fresh dial.** Burst beyond warm supply, or a socket that
died, never blocks or fails — it falls back to the original path. The pool can only
make a connect faster, never slower or more fragile.
- **Auth scope.** The pool key includes `api_key`, so a warm socket is only handed to
a request resolving to the same key — no cross-tenant reuse.
- **Idle billing.** Warm sockets are liveness-checked at handoff and capped at
`REALTIME_POOL_MAX_IDLE_SECS` to bound idle billing and dodge OpenAI's idle timeout.
- **Replenish backoff.** If a key's warm-up dials all fail (invalid credentials, an
unreachable upstream), the replenisher puts that key into exponential backoff
(500 ms → 30 s cap) instead of re-dialing it every tick. This bounds connection
attempts against a broken key so it can't exhaust upstream rate limits and degrade
valid cold-path traffic; the backoff resets the moment a dial succeeds.
Benchmarks and repro: `../../benchmarks/realtime/README.md`.

166
litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs
View file

@ -1,166 +0,0 @@
//! `GET /v1/realtime` (WebSocket).
//!
//! This file is the **axum surface**: `router()`, the handler, and the small
//! socket↔events adapter. The pure logic (no axum) lives in [`service`]. Auth is
//! the `RequireMasterKey` extractor, so the handler stays thin.
mod service;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use crate::io::realtime_pool::RealtimePool;
use axum::Router;
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::extract::{Query, State};
use axum::http::StatusCode;
use axum::response::Response;
use axum::routing::get;
use futures_util::{SinkExt, StreamExt};
use litellm_core::realtime::types::RealtimeEvent;
use litellm_core::router::Router as ModelRouter;
use serde::Deserialize;
use crate::auth::RequireMasterKey;
use crate::integrations::custom_logger::CustomLogger;
use crate::integrations::types::RequestMetadata;
use crate::realtime::streaming::{RealTimeStreaming, SessionStatus};
use crate::state::AppState;
/// Process-local monotonic counter, mixed into the per-session call id so two
/// sessions opened in the same nanosecond still get distinct ids.
static CALL_SEQ: AtomicU64 = AtomicU64::new(0);
/// Generate a per-connection `litellm_call_id`. No external uuid dep: epoch
/// nanos + a process-local sequence is unique enough for log correlation.
fn new_call_id() -> String {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0);
let seq = CALL_SEQ.fetch_add(1, Ordering::Relaxed);
format!("rt-{nanos:x}-{seq:x}")
}
/// This route's contribution to the app router.
pub fn router() -> Router<AppState> {
Router::new().route("/v1/realtime", get(handle))
}
#[derive(Debug, Deserialize)]
struct RealtimeQuery {
model: String,
}
/// Auth runs via the `RequireMasterKey` extractor. We validate the model BEFORE
/// the upgrade so failures are clean HTTP (400/404), not a socket that opens then
/// closes, then hand the socket to `bridge`.
async fn handle(
_auth: RequireMasterKey,
ws: WebSocketUpgrade,
State(state): State<AppState>,
Query(query): Query<RealtimeQuery>,
) -> Result<Response, (StatusCode, String)> {
if query.model.trim().is_empty() {
return Err((
StatusCode::BAD_REQUEST,
"missing 'model' query param".to_string(),
));
}
if !state.router.has_deployment(&query.model) {
return Err((
StatusCode::NOT_FOUND,
format!("no deployment for model '{}'", query.model),
));
}
let router = state.router.clone();
let pool = state.realtime_pool.clone();
let loggers = state.loggers.clone();
let master_key = state.master_key.clone();
let model = query.model;
Ok(ws.on_upgrade(move |socket| bridge(socket, router, pool, loggers, master_key, model)))
}
/// Adapt the axum socket (text frames) to the typed-event `Stream`/`Sink` the
/// service wants, keeping axum types out of `service`.
///
/// This is also the realtime-logging seam: every upstream→client event (the
/// direction carrying `session.created` and `response.done` with usage) is fed
/// to a [`RealTimeStreaming`] collector via the splice's `observe` callback. The
/// observe is O(1) and never buffers frames. When the splice returns (any of the
/// three break paths — client disconnect, upstream close, idle timeout), we flush
/// one logging payload to the registered callbacks.
async fn bridge(
socket: WebSocket,
router: Arc<ModelRouter>,
pool: Arc<RealtimePool>,
loggers: Arc<Vec<Arc<dyn CustomLogger>>>,
master_key: Option<Arc<str>>,
model: String,
) {
let (ws_sink, ws_stream) = socket.split();
// Attribute the spend log to the key that authenticated this session (the
// master key — the gateway is master-key auth). A non-null user_api_key_hash
// is required for the Python spend logger to write a SpendLogs row.
//
// SECURITY: hash the key — never send the raw credential. This field fans out
// to spend logs and every callback integration; the SHA-256 (matching the
// proxy's hash_token) keeps the plaintext master key out of all of them while
// still matching the key's hash in LiteLLM_SpendLogs.
let metadata = RequestMetadata {
user_api_key_hash: master_key.as_deref().map(crate::auth::hash_token),
..RequestMetadata::default()
};
// Owned by THIS task only. The splice observes it via a synchronous `&mut`
// callback (below), so there is no Arc/Mutex/atomic on the per-frame hot
// path — just a monomorphized FnMut mutating stack-local fields. This is
// what lets observe scale: 10K concurrent sessions = 10K independent
// collectors, zero cross-task synchronization.
let mut collector = RealTimeStreaming::new(
loggers.as_ref().clone(),
new_call_id(),
model.clone(),
metadata,
);
let client_in = ws_stream.filter_map(|message| async move {
match message {
Ok(Message::Text(text)) => serde_json::from_str::<RealtimeEvent>(&text).ok(),
_ => None,
}
});
// Plain forwarding sink — no observe here anymore.
let client_out = ws_sink.with(|event: RealtimeEvent| async move {
Ok::<Message, axum::Error>(Message::Text(
serde_json::to_string(&event).unwrap_or_default(),
))
});
futures_util::pin_mut!(client_in, client_out);
// The observe closure borrows `&mut collector` for the duration of the
// splice; the borrow ends when `run` returns, freeing the collector for the
// single post-session `log_messages` flush. `run` picks a pooled (warm) or
// fresh upstream — observe fires on the upstream arm either way.
let result = service::run(
&router,
&pool,
&model,
None,
|event: &RealtimeEvent| collector.observe(event),
client_in,
client_out,
)
.await;
let status = if result.is_ok() {
SessionStatus::Success
} else {
SessionStatus::Failure
};
collector.log_messages(status).await;
}

77
litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs
View file

@ -1,77 +0,0 @@
//! Business logic: select a deployment with the (pure) core router, then call the
//! provider splice. The seam between `core::router` (selection only) and
//! `io` (the actual WebSocket I/O).
//!
//! On connect we try a pre-warmed upstream from the pool (handshake already paid,
//! `session.created` buffered) and relay it instantly. On a pool miss or dead warm
//! socket we fresh-dial exactly as before — the pool is never on the critical path
//! for correctness, only latency.
use std::time::Duration;
use crate::io::realtime_pool::{RealtimePool, upstream_key};
use futures_util::{Sink, Stream};
use litellm_core::error::Error;
use litellm_core::realtime::types::RealtimeEvent;
use litellm_core::router::Router;
/// Select a deployment for `model` and splice the client stream to the provider.
///
/// `pool` supplies a pre-warmed upstream when one is available; otherwise we
/// fresh-dial. A disabled pool always misses, so this collapses to the original
/// fresh-dial behavior.
pub async fn run<In, Out>(
router: &Router,
pool: &RealtimePool,
model: &str,
idle_timeout: Option<Duration>,
observe: impl FnMut(&RealtimeEvent) + Send,
client_in: In,
client_out: Out,
) -> Result<(), Error>
where
In: Stream<Item = RealtimeEvent> + Unpin + Send,
Out: Sink<RealtimeEvent> + Unpin + Send,
<Out as Sink<RealtimeEvent>>::Error: std::fmt::Display,
{
let deployment = router
.get_available_deployment(model)
.ok_or_else(|| Error::Routing(format!("no deployment available for model '{model}'")))?;
let params = &deployment.litellm_params;
// Strip a leading `openai/` so the OpenAI-only realtime fn gets the bare model.
let provider_model = params
.model
.strip_prefix("openai/")
.unwrap_or(&params.model);
// Warm path: take a pooled upstream (handshake already paid) and relay its
// buffered session.created immediately. On miss/dead socket fall through.
if let Some(key) = upstream_key(
provider_model,
params.api_key.as_deref(),
params.api_base.as_deref(),
) && let Some(handoff) = pool.take(&key)
{
return crate::io::realtime::realtime_warm(
provider_model,
handoff,
idle_timeout,
observe,
client_in,
client_out,
)
.await;
}
// Cold path: fresh dial (the original behavior).
crate::io::realtime::realtime(
provider_model,
params.api_key.as_deref(),
params.api_base.as_deref(),
idle_timeout,
observe,
client_in,
client_out,
)
.await
}

348
litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs
View file

@ -1,348 +0,0 @@
mod service;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use axum::Router;
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::extract::{Query, State};
use axum::http::StatusCode;
use axum::response::Response;
use axum::routing::get;
use futures_util::{Sink, SinkExt, StreamExt};
use litellm_core::responses::types::{ResponsesErrorFrame, ResponsesWsEvent, ResponsesWsEventType};
use litellm_core::router::Router as ModelRouter;
use serde::Deserialize;
use crate::auth::RequireMasterKey;
use crate::integrations::custom_logger::CustomLogger;
use crate::integrations::types::RequestMetadata;
use crate::state::AppState;
static CALL_SEQ: AtomicU64 = AtomicU64::new(0);
fn new_call_id() -> String {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_nanos())
.unwrap_or(0);
let sequence = CALL_SEQ.fetch_add(1, Ordering::Relaxed);
format!("respws-{nanos:x}-{sequence:x}")
}
pub fn router() -> Router<AppState> {
Router::new()
.route("/v1/responses", get(handle))
.route("/responses", get(handle))
}
#[derive(Debug, Deserialize)]
struct ResponsesQuery {
model: Option<String>,
}
async fn handle(
_auth: RequireMasterKey,
ws: WebSocketUpgrade,
State(state): State<AppState>,
Query(query): Query<ResponsesQuery>,
) -> Result<Response, (StatusCode, String)> {
if let Some(model) = query.model.as_deref() {
validate_model(&state.router, model)?;
}
let router = state.router.clone();
let loggers = state.loggers.clone();
let master_key = state.master_key.clone();
Ok(ws.on_upgrade(move |socket| bridge(socket, router, loggers, master_key, query.model)))
}
fn validate_model(router: &ModelRouter, model: &str) -> Result<(), (StatusCode, String)> {
if model.trim().is_empty() {
return Err((
StatusCode::BAD_REQUEST,
"missing 'model' query param".to_string(),
));
}
let Some(deployment) = router.get_available_deployment(model) else {
return Err((
StatusCode::NOT_FOUND,
format!("no deployment for model '{model}'"),
));
};
if deployment.litellm_params.model.contains('/')
&& !deployment.litellm_params.model.starts_with("openai/")
{
return Err((
StatusCode::BAD_REQUEST,
"Responses WebSocket route supports OpenAI deployments only".to_string(),
));
}
Ok(())
}
async fn send_error_and_close<S>(sink: &mut S, message: String)
where
S: futures_util::Sink<Message> + Unpin,
S::Error: std::fmt::Display,
{
if let Ok(payload) = serde_json::to_string(&ResponsesErrorFrame::invalid_request(message)) {
let _ = sink.send(Message::Text(payload)).await;
}
let _ = sink
.send(Message::Close(Some(axum::extract::ws::CloseFrame {
code: 1008,
reason: "Pre-call error".into(),
})))
.await;
let _ = sink.close().await;
}
struct ResponseClientSink {
sink: futures_util::stream::SplitSink<WebSocket, Message>,
}
impl Sink<ResponsesWsEvent> for ResponseClientSink {
type Error = axum::Error;
fn poll_ready(
mut self: std::pin::Pin<&mut Self>,
context: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
std::pin::Pin::new(&mut self.sink).poll_ready(context)
}
fn start_send(
mut self: std::pin::Pin<&mut Self>,
item: ResponsesWsEvent,
) -> Result<(), Self::Error> {
let payload = serde_json::to_string(&item).map_err(axum::Error::new)?;
std::pin::Pin::new(&mut self.sink).start_send(Message::Text(payload))
}
fn poll_flush(
mut self: std::pin::Pin<&mut Self>,
context: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
std::pin::Pin::new(&mut self.sink).poll_flush(context)
}
fn poll_close(
mut self: std::pin::Pin<&mut Self>,
context: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
std::pin::Pin::new(&mut self.sink).poll_close(context)
}
}
impl ResponseClientSink {
async fn close_with_code(&mut self, code: u16, reason: &'static str) {
let _ = self
.sink
.send(Message::Close(Some(axum::extract::ws::CloseFrame {
code,
reason: reason.into(),
})))
.await;
let _ = self.sink.close().await;
}
}
async fn bridge(
socket: WebSocket,
router: Arc<ModelRouter>,
loggers: Arc<Vec<Arc<dyn CustomLogger>>>,
master_key: Option<Arc<str>>,
requested_model: Option<String>,
) {
let (mut ws_sink, ws_stream) = socket.split();
let (model, first_frame, stream) = if let Some(model) = requested_model {
(model, None, ws_stream)
} else {
let mut stream = ws_stream;
let first = match stream.next().await {
Some(Ok(Message::Text(text))) => {
match serde_json::from_str::<ResponsesWsEvent>(&text) {
Ok(event) => event,
Err(_) => {
send_error_and_close(
&mut ws_sink,
"Invalid JSON in response.create event".to_string(),
)
.await;
return;
}
}
}
_ => {
send_error_and_close(&mut ws_sink, "Missing response.create event".to_string())
.await;
return;
}
};
let Some(model) = first.model().filter(|value| !value.trim().is_empty()) else {
send_error_and_close(
&mut ws_sink,
"Missing model in response.create event".to_string(),
)
.await;
return;
};
if first.event_type != ResponsesWsEventType::ResponseCreate {
send_error_and_close(
&mut ws_sink,
"First frame must be a response.create event".to_string(),
)
.await;
return;
}
(model.to_string(), Some(first), stream)
};
if let Err((status, message)) = validate_model(&router, &model) {
let _ = status;
let _ = message;
send_error_and_close(&mut ws_sink, "Unknown model deployment".to_string()).await;
return;
}
let call_id = new_call_id();
let metadata = RequestMetadata {
user_api_key_hash: master_key.as_deref().map(crate::auth::hash_token),
..RequestMetadata::default()
};
let client_in = Box::pin(stream.filter_map(|message| async move {
match message {
Ok(Message::Text(text)) => serde_json::from_str::<ResponsesWsEvent>(&text).ok(),
_ => None,
}
}));
let mut client_out = ResponseClientSink { sink: ws_sink };
let result = service::run(
&router,
&model,
first_frame,
None,
loggers,
call_id,
metadata,
client_in,
&mut client_out,
)
.await;
if result.is_err() {
client_out
.close_with_code(1011, "Internal server error")
.await;
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::io::realtime_pool::RealtimePool;
use crate::state::AppState;
use axum::body::Body;
use axum::http::Request;
use litellm_core::router::Router as ModelRouter;
use serde_json::json;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use tower::ServiceExt;
struct RecordingSink {
messages: Vec<Message>,
}
impl Sink<Message> for RecordingSink {
type Error = std::convert::Infallible;
fn poll_ready(
self: Pin<&mut Self>,
_context: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
self.messages.push(item);
Ok(())
}
fn poll_flush(
self: Pin<&mut Self>,
_context: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn poll_close(
self: Pin<&mut Self>,
_context: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
}
#[tokio::test]
async fn pre_call_error_matches_python_frame_and_close() {
let mut sink = RecordingSink {
messages: Vec::new(),
};
send_error_and_close(&mut sink, "missing model".to_string()).await;
let Message::Text(payload) = &sink.messages[0] else {
panic!("expected error text frame");
};
assert_eq!(
serde_json::from_str::<serde_json::Value>(payload).expect("error json"),
json!({
"type": "error",
"error": {
"type": "invalid_request_error",
"message": "missing model"
}
})
);
assert_eq!(
sink.messages[1],
Message::Close(Some(axum::extract::ws::CloseFrame {
code: 1008,
reason: "Pre-call error".into(),
}))
);
}
fn state() -> AppState {
AppState {
router: Arc::new(ModelRouter::default()),
master_key: Some(Arc::from("master-key")),
loggers: Arc::new(Vec::new()),
realtime_pool: RealtimePool::disabled(),
}
}
#[tokio::test]
async fn auth_rejects_responses_upgrade_before_handler() {
let request = Request::builder()
.uri("/responses?model=known")
.body(Body::empty())
.expect("request");
let response = router()
.with_state(state())
.oneshot(request)
.await
.expect("response");
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
#[test]
fn unknown_query_model_is_rejected_before_upgrade() {
assert_eq!(
validate_model(&ModelRouter::default(), "unknown").expect_err("unknown model"),
(
StatusCode::NOT_FOUND,
"no deployment for model 'unknown'".to_string()
)
);
}
}

156
litellm-rust/crates/ai-gateway/src/routes/responses/service.rs
View file

@ -1,156 +0,0 @@
use std::sync::Arc;
use std::time::Duration;
use futures_util::{Sink, Stream};
use litellm_core::Error;
use litellm_core::call_lifecycle::{CallLifecycle, CallLifecycleContext};
use litellm_core::responses::instrumentation::{
ResponsesWsCallbackPayload, ResponsesWsInstrumentation, ResponsesWsLogOutcome,
ResponsesWsMetadata,
};
use litellm_core::responses::types::ResponsesWsEvent;
use crate::integrations::custom_logger::{
CallbackTiming, CallbackValue, CustomLogger, CustomLoggerRunner, LoggingError, ModelCallDetails,
};
use crate::integrations::types::RequestMetadata;
#[allow(clippy::too_many_arguments)]
pub async fn run<In, Out>(
router: &litellm_core::router::Router,
model: &str,
first_frame: Option<ResponsesWsEvent>,
idle_timeout: Option<Duration>,
loggers: Arc<Vec<Arc<dyn CustomLogger>>>,
call_id: String,
metadata: RequestMetadata,
client_in: In,
client_out: Out,
) -> Result<(), Error>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
Out::Error: std::fmt::Display,
{
let deployment = router
.get_available_deployment(model)
.ok_or_else(|| Error::Routing(format!("no deployment available for model '{model}'")))?;
let params = &deployment.litellm_params;
let provider_model = params
.model
.strip_prefix("openai/")
.unwrap_or(&params.model);
if params.model.contains('/') && !params.model.starts_with("openai/") {
return Err(Error::InvalidProvider(
"Responses WebSocket route supports OpenAI deployments only".to_string(),
));
}
let instrumentation = Arc::new(ResponsesWsInstrumentation::new(
call_id.clone(),
model,
ResponsesWsMetadata {
user_api_key_hash: metadata.user_api_key_hash,
user_api_key_user_id: metadata.user_api_key_user_id,
user_api_key_team_id: metadata.user_api_key_team_id,
},
));
let observer_instrumentation = Arc::clone(&instrumentation);
let context = CallLifecycleContext::new("responses_websocket", model, "openai", call_id);
let result = CallLifecycle::default()
.run(context, (), instrumentation.as_ref(), |_| async move {
crate::io::responses_ws::async_responses_websocket(
provider_model,
params.api_key.as_deref(),
params.api_base.as_deref(),
first_frame,
idle_timeout,
move |event| {
observer_instrumentation.observe(event);
},
client_in,
client_out,
)
.await
})
.await;
let outcome = instrumentation.take_or_build_outcome(result.is_ok());
dispatch_outcome(loggers, outcome).await;
result
}
async fn dispatch_outcome(
loggers: Arc<Vec<Arc<dyn CustomLogger>>>,
outcome: ResponsesWsLogOutcome,
) {
let runner = CustomLoggerRunner::new(loggers.as_ref().clone());
match outcome {
ResponsesWsLogOutcome::Success { payload, callback } => {
let (details, response, start_time, end_time) = logging_values(payload, callback, None);
let _ = runner
.async_log_success_event(
&details,
&response,
CallbackTiming::new(start_time, end_time),
)
.await;
}
ResponsesWsLogOutcome::Failure {
payload,
callback,
error_message,
error_kind,
} => {
let error = LoggingError {
message: error_message,
kind: error_kind,
};
let (details, response, start_time, end_time) =
logging_values(payload, callback, Some(error));
let _ = runner
.async_log_failure_event(
&details,
Some(&response),
CallbackTiming::new(start_time, end_time),
)
.await;
}
}
}
fn logging_values(
payload: litellm_core::responses::instrumentation::ResponsesWsLogPayload,
callback: ResponsesWsCallbackPayload,
error: Option<LoggingError>,
) -> (ModelCallDetails, CallbackValue, f64, f64) {
let start_time = payload.start_time;
let end_time = payload.end_time;
let callback = CallbackValue::new(callback.object, callback.value);
let details = ModelCallDetails::from_standard_logging_payload(
crate::integrations::types::StandardLoggingPayload {
id: payload.id,
litellm_call_id: payload.litellm_call_id,
call_type: payload.call_type,
model: payload.model,
custom_llm_provider: payload.custom_llm_provider,
response_cost: payload.response_cost,
prompt_tokens: payload.usage.prompt_tokens,
completion_tokens: payload.usage.completion_tokens,
total_tokens: payload.usage.total_tokens,
start_time: payload.start_time,
end_time: payload.end_time,
stream: payload.stream,
metadata: crate::integrations::types::StandardLoggingMetadata {
user_api_key_hash: payload.metadata.user_api_key_hash,
user_api_key_user_id: payload.metadata.user_api_key_user_id,
user_api_key_team_id: payload.metadata.user_api_key_team_id,
..Default::default()
},
messages: None,
},
);
let details = match error {
Some(error) => details.with_failure_error(error),
None => details,
};
(details, callback, start_time, end_time)
}

21
litellm-rust/crates/ai-gateway/src/state.rs
View file

@ -1,21 +0,0 @@
use std::sync::Arc;
use crate::io::realtime_pool::RealtimePool;
use litellm_core::router::Router;
use crate::integrations::custom_logger::CustomLogger;
/// Shared application state handed to every route handler.
#[derive(Clone)]
pub struct AppState {
pub router: Arc<Router>,
/// The gateway master key. Any caller presenting it as a bearer token may
/// invoke the gateway. `None` → auth not configured (routes fail closed).
pub master_key: Option<Arc<str>>,
/// Logging callbacks fanned out at the end of each realtime session.
pub loggers: Arc<Vec<Arc<dyn CustomLogger>>>,
/// Pre-warmed upstream realtime connection pool. Disabled
/// (`RealtimePool::disabled()`) when `REALTIME_POOL_SIZE=0`, in which case
/// every realtime connect fresh-dials exactly as before.
pub realtime_pool: Arc<RealtimePool>,
}

100
litellm-rust/crates/ai-gateway/src/trace_parity.rs
View file

@ -1,100 +0,0 @@
//! Harness-only in-process adapters. Never mounted as production routes.
use std::sync::Arc;
use axum::body::{Body, to_bytes};
use axum::http::header::{AUTHORIZATION, CONTENT_TYPE};
use axum::http::{Request, StatusCode};
use litellm_core::Error;
use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter};
use serde::Serialize;
use serde_json::Value;
use tower::ServiceExt;
use tracing::instrument::WithSubscriber;
use crate::io::realtime_pool::RealtimePool;
use crate::routes;
use crate::state::AppState;
#[derive(Debug, Serialize)]
pub struct GatewayResponse {
pub status: u16,
pub body: Value,
}
#[derive(Debug, Serialize)]
pub struct TracedGatewayResponse {
pub response: Option<GatewayResponse>,
pub error: Option<String>,
pub trace: Vec<litellm_core::observability::FunctionTraceEvent>,
}
pub async fn traced_request(
path: String,
model_alias: String,
provider_model: String,
api_base: String,
body: Value,
) -> TracedGatewayResponse {
let trace = litellm_core::observability::FunctionTrace::default();
let result = request(path, model_alias, provider_model, api_base, body)
.with_subscriber(trace.dispatcher())
.await;
let events = trace.events();
match result {
Ok(response) => TracedGatewayResponse {
response: Some(response),
error: None,
trace: events,
},
Err(error) => TracedGatewayResponse {
response: None,
error: Some(error.to_string()),
trace: events,
},
}
}
pub async fn request(
path: String,
model_alias: String,
provider_model: String,
api_base: String,
body: Value,
) -> Result<GatewayResponse, Error> {
let state = AppState {
router: Arc::new(ModelRouter::new(vec![Deployment {
model_name: model_alias,
litellm_params: LiteLLMParams {
model: provider_model,
api_key: Some("trace-provider-key".to_string()),
api_base: Some(api_base),
},
}])),
master_key: Some(Arc::from("trace-master-key")),
loggers: Arc::new(Vec::new()),
realtime_pool: RealtimePool::disabled(),
};
let request = Request::builder()
.method("POST")
.uri(path)
.header(AUTHORIZATION, "Bearer trace-master-key")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(body.to_string()))
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
let response = match routes::app(state).oneshot(request).await {
Ok(response) => response,
Err(error) => match error {},
};
let status: StatusCode = response.status();
let bytes = to_bytes(response.into_body(), usize::MAX)
.await
.map_err(|error| Error::InvalidResponse(error.to_string()))?;
let body = serde_json::from_slice(&bytes).map_err(|error| {
Error::InvalidResponse(format!("gateway returned invalid JSON: {error}"))
})?;
Ok(GatewayResponse {
status: status.as_u16(),
body,
})
}

53
litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs
View file

@ -1,53 +0,0 @@
//! Guards the wiring, not just the helper: a `wss://` dial through the public
//! API has to resolve its own crypto provider, in a test binary where nothing
//! has installed a process-wide one, and has to leave it uninstalled.
use std::time::Duration;
use futures_util::{sink, stream};
use litellm_ai_gateway::io::responses_ws::async_responses_websocket;
use tokio::net::TcpListener;
async fn dead_tls_server() -> u16 {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind a loopback port");
let port = listener
.local_addr()
.expect("read the bound address")
.port();
tokio::spawn(async move {
while let Ok((stream, _peer)) = listener.accept().await {
drop(stream);
}
});
port
}
#[tokio::test]
async fn dialing_wss_returns_an_error_instead_of_panicking() {
let port = dead_tls_server().await;
let result = async_responses_websocket(
"gpt-5",
Some("test-key"),
Some(&format!("wss://127.0.0.1:{port}/")),
None,
Some(Duration::from_secs(10)),
|_| {},
stream::empty(),
sink::drain(),
)
.await;
assert!(
result.is_err(),
"a plain TCP server cannot finish a TLS handshake"
);
assert!(
rustls::crypto::CryptoProvider::get_default().is_none(),
"the dial settles its provider on its own connector, not process-wide"
);
}

25
litellm-rust/crates/auth-aws/Cargo.toml Normal file
View file

@ -0,0 +1,25 @@
[package]
name = "litellm-auth-aws"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-auth.workspace = true
moka = { workspace = true, features = ["sync"] }
serde_json.workspace = true
sha2.workspace = true
thiserror.workspace = true
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"] }
aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"] }
aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"] }
aws-sigv4 = "1.5.1"
aws-types = "1.4.0"
aws-smithy-runtime-api = "1.13.0"
[dev-dependencies]
reqwest.workspace = true
tokio.workspace = true

949
litellm-rust/crates/auth-aws/src/aws.rs Normal file
View file

@ -0,0 +1,949 @@
use std::collections::BTreeMap;
use std::sync::OnceLock;
use std::time::Duration;
use std::time::{SystemTime, UNIX_EPOCH};
use moka::sync::Cache;
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use aws_credential_types::Credentials;
use aws_credential_types::provider::ProvideCredentials;
use aws_sigv4::http_request::{
SignableBody, SignableRequest, SigningParams, SigningSettings, sign,
};
use aws_sigv4::sign::v4;
use aws_smithy_runtime_api::client::identity::Identity;
use super::Error;
use super::constants::{
AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION, AWS_REGION_NAME,
AWS_ROLE_ARN, AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN,
AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE,
BEDROCK_SERVICE, DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX,
SIGV4_COMPUTED_HEADER_NAMES,
};
const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60);
const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600);
static STATIC_CREDENTIALS_CACHE: OnceLock<Cache<String, Credentials>> = OnceLock::new();
static AMBIENT_CREDENTIALS_CACHE: OnceLock<Cache<String, Credentials>> = OnceLock::new();
fn credential_cache_ttl(flow: &AwsAuthFlow) -> Option<Duration> {
match flow {
AwsAuthFlow::StaticKeys { .. } => Some(STATIC_CREDENTIALS_TTL),
AwsAuthFlow::DefaultChain => Some(AMBIENT_CREDENTIALS_TTL),
AwsAuthFlow::WebIdentity { .. }
| AwsAuthFlow::AssumeRole { .. }
| AwsAuthFlow::Profile { .. }
| AwsAuthFlow::SessionToken { .. } => None,
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct AwsAuthConfig {
pub access_key_id: Option<String>,
pub secret_access_key: Option<String>,
pub session_token: Option<String>,
pub region_name: Option<String>,
pub session_name: Option<String>,
pub profile_name: Option<String>,
pub role_name: Option<String>,
pub web_identity_token: Option<String>,
pub sts_endpoint: Option<String>,
pub external_id: Option<String>,
}
impl AwsAuthConfig {
fn with_environment(self, env_lookup: &(dyn Fn(&str) -> Option<String> + Sync)) -> Self {
Self {
access_key_id: self.access_key_id.or_else(|| env_lookup(AWS_ACCESS_KEY_ID)),
secret_access_key: self
.secret_access_key
.or_else(|| env_lookup(AWS_SECRET_ACCESS_KEY)),
session_token: self.session_token.or_else(|| env_lookup(AWS_SESSION_TOKEN)),
region_name: self.region_name.or_else(|| env_lookup(AWS_REGION_NAME)),
session_name: self.session_name.or_else(|| env_lookup(AWS_SESSION_NAME)),
profile_name: self.profile_name.or_else(|| env_lookup(AWS_PROFILE_NAME)),
role_name: self.role_name.or_else(|| env_lookup(AWS_ROLE_NAME)),
web_identity_token: self
.web_identity_token
.or_else(|| env_lookup(AWS_WEB_IDENTITY_TOKEN)),
sts_endpoint: self.sts_endpoint.or_else(|| env_lookup(AWS_STS_ENDPOINT)),
external_id: self.external_id.or_else(|| env_lookup(AWS_EXTERNAL_ID)),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum AwsAuthFlow {
WebIdentity {
token: String,
role: String,
session_name: String,
},
AssumeRole {
role: String,
session_name: Option<String>,
},
Profile {
name: String,
},
SessionToken {
access_key_id: String,
secret_access_key: String,
session_token: String,
},
StaticKeys {
access_key_id: String,
secret_access_key: String,
region_name: String,
},
DefaultChain,
}
fn cache_key(config: &AwsAuthConfig, flow: &AwsAuthFlow) -> String {
let mut hasher = Sha256::new();
hasher.update(format!("{config:?}:{flow:?}"));
format!("{:x}", hasher.finalize())
}
fn static_credentials_cache() -> &'static Cache<String, Credentials> {
STATIC_CREDENTIALS_CACHE.get_or_init(|| {
Cache::builder()
.max_capacity(200)
.time_to_live(STATIC_CREDENTIALS_TTL)
.build()
})
}
fn ambient_credentials_cache() -> &'static Cache<String, Credentials> {
AMBIENT_CREDENTIALS_CACHE.get_or_init(|| {
Cache::builder()
.max_capacity(200)
.time_to_live(AMBIENT_CREDENTIALS_TTL)
.build()
})
}
fn get_cached_credentials(key: &str) -> Option<Credentials> {
static_credentials_cache()
.get(key)
.or_else(|| ambient_credentials_cache().get(key))
}
fn set_cached_credentials(key: String, credentials: Credentials, ttl: Duration) {
if ttl == STATIC_CREDENTIALS_TTL {
static_credentials_cache().insert(key, credentials);
} else {
ambient_credentials_cache().insert(key, credentials);
}
}
fn role_identity(arn: &str) -> Option<(&str, &str, &str)> {
let mut parts = arn.splitn(6, ':');
let ("arn", partition, _, _, account, resource) = (
parts.next()?,
parts.next()?,
parts.next()?,
parts.next()?,
parts.next()?,
parts.next()?,
) else {
return None;
};
let role = if let Some(role) = resource.strip_prefix("role/") {
role.rsplit('/').next()?
} else {
resource.strip_prefix("assumed-role/")?.split('/').next()?
};
Some((partition, account, role))
}
fn same_role_arns(target: &str, caller: &str) -> bool {
role_identity(target) == role_identity(caller)
}
pub fn classify_auth(
config: AwsAuthConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> AwsAuthFlow {
let config = config.with_environment(env_lookup);
if let (Some(token), Some(role), Some(session_name)) = (
config.web_identity_token.clone(),
config.role_name.clone(),
config.session_name.clone(),
) {
return AwsAuthFlow::WebIdentity {
token,
role,
session_name,
};
}
if let Some(role) = config.role_name.clone() {
return AwsAuthFlow::AssumeRole {
role,
session_name: config.session_name.clone(),
};
}
if let Some(name) = config.profile_name {
return AwsAuthFlow::Profile { name };
}
if let (Some(access_key_id), Some(secret_access_key), Some(session_token)) = (
config.access_key_id.clone(),
config.secret_access_key.clone(),
config.session_token,
) {
return AwsAuthFlow::SessionToken {
access_key_id,
secret_access_key,
session_token,
};
}
if let (Some(access_key_id), Some(secret_access_key), Some(region_name)) = (
config.access_key_id,
config.secret_access_key,
config.region_name,
) {
return AwsAuthFlow::StaticKeys {
access_key_id,
secret_access_key,
region_name,
};
}
AwsAuthFlow::DefaultChain
}
pub async fn resolve_credentials(
config: AwsAuthConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Credentials, Error> {
let resolved = config.clone().with_environment(env_lookup);
let flow = classify_auth(config, env_lookup);
match flow {
AwsAuthFlow::SessionToken {
access_key_id,
secret_access_key,
session_token,
} => Ok(Credentials::new(
access_key_id,
secret_access_key,
Some(session_token),
None,
"litellm-static-session",
)),
AwsAuthFlow::StaticKeys {
access_key_id,
secret_access_key,
region_name,
} => {
let flow = AwsAuthFlow::StaticKeys {
access_key_id: access_key_id.clone(),
secret_access_key: secret_access_key.clone(),
region_name,
};
let key = cache_key(&resolved, &flow);
if let Some(credentials) = get_cached_credentials(&key) {
return Ok(credentials);
}
let credentials = Credentials::new(
access_key_id,
secret_access_key,
None,
None,
"litellm-static",
);
set_cached_credentials(
key,
credentials.clone(),
credential_cache_ttl(&flow).unwrap_or(STATIC_CREDENTIALS_TTL),
);
Ok(credentials)
}
AwsAuthFlow::Profile { name } => {
let provider = aws_config::profile::ProfileFileCredentialsProvider::builder()
.profile_name(name)
.build();
provider
.provide_credentials()
.await
.map_err(|error| Error::AwsProfile(error.to_string()))
}
AwsAuthFlow::AssumeRole { role, session_name } => {
if is_already_running_as_role(&role, &resolved).await? {
let ambient_flow = AwsAuthFlow::DefaultChain;
let key = cache_key(&resolved, &ambient_flow);
if let Some(credentials) = get_cached_credentials(&key) {
return Ok(credentials);
}
let provider =
aws_config::default_provider::credentials::DefaultCredentialsChain::builder()
.build()
.await;
let credentials = provider
.provide_credentials()
.await
.map_err(|error| Error::AwsDefaultChain(error.to_string()))?;
set_cached_credentials(
key,
credentials.clone(),
credential_cache_ttl(&ambient_flow).unwrap_or(AMBIENT_CREDENTIALS_TTL),
);
return Ok(credentials);
}
let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
if let Some(region) = resolved.region_name.clone() {
loader = loader.region(aws_types::region::Region::new(region));
}
if let Some(endpoint) = resolved.sts_endpoint.clone() {
loader = loader.endpoint_url(endpoint);
}
if let (Some(access_key_id), Some(secret_access_key)) =
(resolved.access_key_id, resolved.secret_access_key)
{
loader = loader.credentials_provider(Credentials::new(
access_key_id,
secret_access_key,
resolved.session_token,
None,
"litellm-role-source",
));
}
let sdk_config = loader.load().await;
let builder = aws_config::sts::AssumeRoleProvider::builder(role);
let builder = match session_name {
Some(name) => builder.session_name(name),
None => builder.session_name(default_session_name()),
};
let builder = match resolved.external_id {
Some(id) => builder.external_id(id),
None => builder,
};
let provider = builder.configure(&sdk_config).build().await;
provider
.provide_credentials()
.await
.map_err(|error| Error::AwsAssumeRole(error.to_string()))
}
AwsAuthFlow::WebIdentity {
token,
role,
session_name,
} => {
let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
if let Some(region) = resolved.region_name {
loader = loader.region(aws_types::region::Region::new(region));
}
if let Some(endpoint) = resolved.sts_endpoint {
loader = loader.endpoint_url(endpoint);
}
let sdk_config = loader.load().await;
let client = aws_sdk_sts::Client::new(&sdk_config);
let response = client
.assume_role_with_web_identity()
.role_arn(role)
.role_session_name(session_name)
.web_identity_token(token)
.send()
.await
.map_err(|error| Error::AwsWebIdentity(error.to_string()))?;
let credentials = response
.credentials()
.ok_or(Error::AwsMissingWebIdentityCredentials)?;
let expiration = SystemTime::try_from(*credentials.expiration())
.map_err(|error| Error::AwsWebIdentityExpiration(error.to_string()))?;
Ok(Credentials::new(
credentials.access_key_id(),
credentials.secret_access_key(),
Some(credentials.session_token().to_string()),
Some(expiration),
"litellm-web-identity",
))
}
AwsAuthFlow::DefaultChain => {
let key = cache_key(&resolved, &AwsAuthFlow::DefaultChain);
if let Some(credentials) = get_cached_credentials(&key) {
return Ok(credentials);
}
let provider =
aws_config::default_provider::credentials::DefaultCredentialsChain::builder()
.build()
.await;
let credentials = provider
.provide_credentials()
.await
.map_err(|error| Error::AwsDefaultChain(error.to_string()))?;
set_cached_credentials(
key,
credentials.clone(),
credential_cache_ttl(&AwsAuthFlow::DefaultChain).unwrap_or(AMBIENT_CREDENTIALS_TTL),
);
Ok(credentials)
}
}
}
async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> Result<bool, Error> {
if role_identity(role).is_none() {
return Ok(false);
}
if let (Ok(current_role), Ok(token_file)) = (
std::env::var(AWS_ROLE_ARN),
std::env::var(AWS_WEB_IDENTITY_TOKEN_FILE),
) && !token_file.is_empty()
{
return Ok(same_role_arns(role, &current_role));
}
let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
if let Some(region) = config.region_name.clone() {
loader = loader.region(aws_types::region::Region::new(region));
}
if let Some(endpoint) = config.sts_endpoint.clone() {
loader = loader.endpoint_url(endpoint);
}
let sdk_config = loader.load().await;
let response = match aws_sdk_sts::Client::new(&sdk_config)
.get_caller_identity()
.send()
.await
{
Ok(response) => response,
Err(_) => return Ok(false),
};
Ok(response
.arn()
.is_some_and(|caller| same_role_arns(role, caller)))
}
fn default_session_name() -> String {
let seconds = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |duration| duration.as_secs());
format!("{DEFAULT_SESSION_NAME_PREFIX}-{seconds}")
}
/// The subset of `headers` SigV4 should cover.
///
/// Python signs only these and reattaches the rest afterwards, so a forwarded
/// client header cannot change the canonical request and invalidate the
/// signature. Signing everything instead makes the request 403 on a header the
/// caller supplied, on a deployment that works on the Python path.
pub fn aws_signature_headers(headers: &BTreeMap<String, String>) -> BTreeMap<String, String> {
headers
.iter()
.filter(|(name, _)| {
let name = name.to_ascii_lowercase();
AWS_SIGNED_HEADER_NAMES.contains(&name.as_str())
|| name.starts_with("x-amz-")
|| name.starts_with("x-amzn-")
})
.map(|(name, value)| (name.clone(), value.clone()))
.collect()
}
/// Whether the signer produces `name` itself.
///
/// Python's reattach loop skips these, so a caller-supplied copy never reaches
/// the wire next to the computed one.
pub fn is_sigv4_computed_header(name: &str) -> bool {
SIGV4_COMPUTED_HEADER_NAMES.contains(&name.to_ascii_lowercase().as_str())
}
pub fn sign_bedrock_post(
url: &str,
body: &[u8],
headers: &BTreeMap<String, String>,
region: &str,
credentials: &Credentials,
signing_time: SystemTime,
) -> Result<BTreeMap<String, String>, Error> {
let identity: Identity = credentials.clone().into();
let params = v4::SigningParams::builder()
.identity(&identity)
.region(region)
.name(BEDROCK_SERVICE)
.time(signing_time)
.settings(SigningSettings::default())
.build()
.map(SigningParams::from)
.map_err(|error| Error::AwsSigningParameters(error.to_string()))?;
let header_refs = headers
.iter()
.map(|(name, value)| (name.as_str(), value.as_str()));
let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body))
.map_err(|error| Error::AwsSignableRequest(error.to_string()))?;
let (instructions, _) = sign(request, &params)
.map_err(|error| Error::AwsSigning(error.to_string()))?
.into_parts();
Ok(instructions
.headers()
.map(|(name, value)| {
let normalized_name = match name {
"authorization" => "Authorization",
"x-amz-date" => "X-Amz-Date",
"x-amz-security-token" => "X-Amz-Security-Token",
_ => name,
};
(normalized_name.to_string(), value.to_string())
})
.collect())
}
/// Model-id and region parsing shared by every Bedrock route.
pub fn bedrock_model_id_and_region(model: &str) -> (String, Option<String>) {
let mut stripped = model;
for prefix in ["bedrock/converse/", "bedrock/", "converse/"] {
if let Some(value) = stripped.strip_prefix(prefix) {
stripped = value;
break;
}
}
let mut region = None;
if let Some((candidate, remainder)) = stripped.split_once('/')
&& is_bedrock_region(candidate)
{
region = Some(candidate.to_string());
stripped = remainder;
}
for prefix in ["nova-2/", "nova/"] {
if let Some(value) = stripped.strip_prefix(prefix) {
stripped = value;
break;
}
}
if region.is_none() {
// Python splits the whole ARN and takes field 3, the region. Stripping
// `arn:` first shifts every field down one, so the region is field 2
// here; field 3 is the account id.
region = stripped
.strip_prefix("arn:")
.and_then(|value| value.split(':').nth(2))
.filter(|value| !value.is_empty())
.map(str::to_string);
}
(stripped.to_string(), region)
}
fn is_bedrock_region(value: &str) -> bool {
value.len() > 3
&& value.contains('-')
&& value
.chars()
.all(|char| char.is_ascii_alphanumeric() || char == '-')
}
pub fn resolve_bedrock_region(
model_region: Option<&str>,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> String {
if let Some(region) = optional_params
.get("aws_region_name")
.and_then(Value::as_str)
{
return region.to_string();
}
if let Some(region) = model_region {
return region.to_string();
}
env_lookup(AWS_REGION_NAME)
.or_else(|| env_lookup(AWS_REGION))
.unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string())
}
pub fn aws_auth_config(
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> AwsAuthConfig {
let value = |key: &str| {
optional_params
.get(key)
.and_then(Value::as_str)
.map(str::to_string)
};
let env = |key: &str| env_lookup(key);
AwsAuthConfig {
access_key_id: value("aws_access_key_id").or_else(|| env("AWS_ACCESS_KEY_ID")),
secret_access_key: value("aws_secret_access_key").or_else(|| env("AWS_SECRET_ACCESS_KEY")),
session_token: value("aws_session_token").or_else(|| env("AWS_SESSION_TOKEN")),
region_name: value("aws_region_name").or_else(|| env(AWS_REGION_NAME)),
session_name: value("aws_session_name").or_else(|| env("AWS_SESSION_NAME")),
profile_name: value("aws_profile_name").or_else(|| env("AWS_PROFILE_NAME")),
role_name: value("aws_role_name").or_else(|| env("AWS_ROLE_NAME")),
web_identity_token: value("aws_web_identity_token")
.or_else(|| env("AWS_WEB_IDENTITY_TOKEN")),
sts_endpoint: value("aws_sts_endpoint").or_else(|| env("AWS_STS_ENDPOINT")),
external_id: value("aws_external_id").or_else(|| env("AWS_EXTERNAL_ID")),
}
}
/// Credentials a host resolved through its own chain and handed down verbatim.
///
/// A host with its own resolution (LiteLLM's Python `BaseAWSLLM`, which reads
/// profiles, STS and boto sessions) passes the result here so the core signs
/// with exactly those. Without this the core would re-derive from ambient
/// state, where an unrelated `AWS_ROLE_NAME` or `AWS_PROFILE_NAME` in the
/// environment outranks explicit keys in [`classify_auth`] and the two sides
/// would sign as different principals.
pub fn host_supplied_credentials(optional_params: &Map<String, Value>) -> Option<Credentials> {
let value = |key: &str| {
optional_params
.get(key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
};
let access_key_id = value("aws_access_key_id")?;
let secret_access_key = value("aws_secret_access_key")?;
Some(Credentials::new(
access_key_id,
secret_access_key,
value("aws_session_token").map(str::to_string),
None,
"litellm-host-supplied",
))
}
#[cfg(test)]
mod tests {
use super::*;
fn no_env(_: &str) -> Option<String> {
None
}
fn parity_inputs() -> (String, Vec<u8>, BTreeMap<String, String>) {
(
"https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke"
.to_string(),
br#"{"input":"hello"}"#.to_vec(),
BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]),
)
}
#[test]
fn reads_the_region_field_of_a_model_arn_not_the_account_id() {
// Python's `_get_aws_region_from_model_arn` splits the whole ARN and
// takes field 3. Stripping `arn:` first shifts every field down one, so
// the region is field 2 here. Taking field 3 after the strip returns
// the account id, which is not a region at all.
let (_, region) = bedrock_model_id_and_region(
"bedrock/arn:aws:bedrock:us-west-2:123456789012:foundation-model/anthropic.claude-v2",
);
assert_eq!(region.as_deref(), Some("us-west-2"));
}
#[test]
fn classification_preserves_python_precedence() {
let config = AwsAuthConfig {
access_key_id: Some("ak".into()),
secret_access_key: Some("sk".into()),
session_token: Some("token".into()),
region_name: Some("us-east-1".into()),
session_name: Some("session".into()),
profile_name: Some("profile".into()),
role_name: Some("role".into()),
web_identity_token: Some("oidc".into()),
..Default::default()
};
assert!(matches!(
classify_auth(config, &no_env),
AwsAuthFlow::WebIdentity { .. }
));
}
#[test]
fn classification_covers_fallthroughs() {
let env = |key: &str| match key {
AWS_PROFILE_NAME => Some("profile".into()),
_ => None,
};
assert!(matches!(
classify_auth(AwsAuthConfig::default(), &env),
AwsAuthFlow::Profile { .. }
));
assert!(matches!(
classify_auth(
AwsAuthConfig {
access_key_id: Some("ak".into()),
secret_access_key: Some("sk".into()),
session_token: Some("token".into()),
..Default::default()
},
&no_env
),
AwsAuthFlow::SessionToken { .. }
));
assert!(matches!(
classify_auth(
AwsAuthConfig {
access_key_id: Some("ak".into()),
secret_access_key: Some("sk".into()),
region_name: Some("us-east-1".into()),
..Default::default()
},
&no_env
),
AwsAuthFlow::StaticKeys { .. }
));
assert_eq!(
classify_auth(AwsAuthConfig::default(), &no_env),
AwsAuthFlow::DefaultChain
);
}
#[tokio::test]
async fn static_credentials_do_not_use_network() {
let credentials = resolve_credentials(
AwsAuthConfig {
access_key_id: Some("ak".into()),
secret_access_key: Some("sk".into()),
region_name: Some("us-east-1".into()),
..Default::default()
},
&no_env,
)
.await
.expect("static credentials");
assert_eq!(credentials.access_key_id(), "ak");
assert_eq!(credentials.session_token(), None);
}
#[test]
fn cache_policy_matches_python_flows() {
assert_eq!(
credential_cache_ttl(&AwsAuthFlow::StaticKeys {
access_key_id: "ak".into(),
secret_access_key: "sk".into(),
region_name: "us-east-1".into(),
}),
Some(STATIC_CREDENTIALS_TTL)
);
assert_eq!(
credential_cache_ttl(&AwsAuthFlow::DefaultChain),
Some(AMBIENT_CREDENTIALS_TTL)
);
assert_eq!(
credential_cache_ttl(&AwsAuthFlow::SessionToken {
access_key_id: "ak".into(),
secret_access_key: "sk".into(),
session_token: "token".into(),
}),
None
);
assert_eq!(
credential_cache_ttl(&AwsAuthFlow::Profile {
name: "profile".into()
}),
None
);
assert_eq!(
credential_cache_ttl(&AwsAuthFlow::AssumeRole {
role: "arn:aws:iam::123456789012:role/demo".into(),
session_name: None,
}),
None
);
assert_eq!(
credential_cache_ttl(&AwsAuthFlow::WebIdentity {
token: "token".into(),
role: "arn:aws:iam::123456789012:role/demo".into(),
session_name: "session".into(),
}),
None
);
}
#[test]
fn cache_round_trip_preserves_credentials() {
let key = format!("cache-test-{}", std::process::id());
let credentials = Credentials::new("cache-ak", "cache-sk", None, None, "test");
set_cached_credentials(key.clone(), credentials.clone(), STATIC_CREDENTIALS_TTL);
assert_eq!(
get_cached_credentials(&key).map(|value| value.access_key_id().to_string()),
Some("cache-ak".to_string())
);
}
#[test]
fn same_role_comparison_matches_partition_account_and_role() {
assert!(same_role_arns(
"arn:aws:iam::123456789012:role/path/demo",
"arn:aws:sts::123456789012:assumed-role/demo/session"
));
assert!(!same_role_arns(
"arn:aws:iam::123456789012:role/demo",
"arn:aws:iam::999999999999:role/demo"
));
assert!(!same_role_arns(
"arn:aws:iam::123456789012:role/demo",
"arn:aws-cn:iam::123456789012:role/demo"
));
assert!(!same_role_arns(
"arn:aws:iam::123456789012:user/demo",
"arn:aws:iam::123456789012:role/demo"
));
}
#[test]
fn a_forwarded_client_header_is_not_folded_into_the_signature() {
// Python signs only the AWS header set, so a header a caller forwarded
// cannot change the canonical request. Signing it instead makes the
// request 403 the moment anything on the wire rewrites or drops it.
let (url, body, mut headers) = parity_inputs();
headers.insert("x-request-id".to_string(), "abc-123".to_string());
headers.insert("Accept-Encoding".to_string(), "gzip".to_string());
headers.insert("x-amzn-trace-id".to_string(), "Root=1-abc".to_string());
let signable = aws_signature_headers(&headers);
assert!(!signable.contains_key("x-request-id"));
assert!(!signable.contains_key("Accept-Encoding"));
// The AWS-prefixed one is genuinely part of the signature.
assert!(signable.contains_key("x-amzn-trace-id"));
assert!(signable.contains_key("Content-Type"));
let credentials = Credentials::new(
"AKIDEXAMPLE",
"wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
None,
None,
"test",
);
let signed = sign_bedrock_post(
&url,
&body,
&signable,
"us-east-1",
&credentials,
SystemTime::UNIX_EPOCH,
)
.expect("signs");
let authorization = signed
.get("Authorization")
.expect("carries an authorization header");
assert!(
!authorization.contains("x-request-id"),
"forwarded header reached SignedHeaders: {authorization}"
);
assert!(
!authorization.contains("accept-encoding"),
"forwarded header reached SignedHeaders: {authorization}"
);
}
#[test]
fn signing_matches_botocore_golden_vector() {
let (url, body, headers) = parity_inputs();
let credentials = Credentials::new(
"AKIDEXAMPLE",
"wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
Some("session-token".to_string()),
None,
"test",
);
let signed = sign_bedrock_post(
&url,
&body,
&headers,
"us-east-1",
&credentials,
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
)
.expect("golden signature");
assert_eq!(
signed.get("X-Amz-Date").map(String::as_str),
Some("20240102T030405Z")
);
assert_eq!(
signed.get("X-Amz-Security-Token").map(String::as_str),
Some("session-token")
);
assert_eq!(
signed.get("Authorization").map(String::as_str),
Some(
"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464"
)
);
}
#[test]
fn signing_without_session_token_omits_security_header() {
let (url, body, headers) = parity_inputs();
let credentials = Credentials::new(
"AKIDEXAMPLE",
"wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
None,
None,
"test",
);
let signed = sign_bedrock_post(
&url,
&body,
&headers,
"us-east-1",
&credentials,
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
)
.expect("signature");
assert!(!signed.contains_key("X-Amz-Security-Token"));
}
#[ignore]
#[tokio::test]
async fn live_bedrock_invoke_model_returns_200() -> Result<(), Box<dyn std::error::Error>> {
let access_key_id = std::env::var("AWS_BEDROCK_TEST_ACCESS_KEY_ID")?;
let secret_access_key = std::env::var("AWS_BEDROCK_TEST_SECRET_ACCESS_KEY")?;
let body = br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":1,"messages":[{"role":"user","content":[{"type":"text","text":"ping"}]}]}"#.to_vec();
let headers =
BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]);
let credentials = resolve_credentials(
AwsAuthConfig {
access_key_id: Some(access_key_id),
secret_access_key: Some(secret_access_key),
region_name: Some("us-west-2".to_string()),
..Default::default()
},
&no_env,
)
.await?;
let client = reqwest::Client::new();
let mut failures = Vec::new();
for region in ["us-west-2", "us-east-1"] {
let url = format!(
"https://bedrock-runtime.{region}.amazonaws.com/model/us.anthropic.claude-opus-4-8/invoke"
);
let signed_headers = sign_bedrock_post(
&url,
&body,
&headers,
region,
&credentials,
SystemTime::now(),
)?;
let mut request = client.post(&url).body(body.clone());
for (name, value) in &headers {
request = request.header(name, value);
}
for (name, value) in signed_headers {
request = request.header(name, value);
}
let response = request.send().await?;
let status = response.status();
let response_body = response.text().await?;
let snippet: String = response_body.chars().take(240).collect();
println!("region={region} status={status} response={snippet}");
if status == reqwest::StatusCode::OK {
return Ok(());
}
failures.push(format!("{region}: {status} {snippet}"));
}
panic!(
"no Bedrock region returned HTTP 200: {}",
failures.join("; ")
);
}
}

43
litellm-rust/crates/auth-aws/src/constants.rs Normal file
View file

@ -0,0 +1,43 @@
pub const AWS_ACCESS_KEY_ID: &str = "AWS_ACCESS_KEY_ID";
pub const AWS_SECRET_ACCESS_KEY: &str = "AWS_SECRET_ACCESS_KEY";
pub const AWS_SESSION_TOKEN: &str = "AWS_SESSION_TOKEN";
pub const AWS_REGION_NAME: &str = "AWS_REGION_NAME";
pub const AWS_REGION: &str = "AWS_REGION";
pub const AWS_SESSION_NAME: &str = "AWS_SESSION_NAME";
pub const AWS_PROFILE_NAME: &str = "AWS_PROFILE_NAME";
pub const AWS_ROLE_NAME: &str = "AWS_ROLE_NAME";
pub const AWS_WEB_IDENTITY_TOKEN: &str = "AWS_WEB_IDENTITY_TOKEN";
pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN";
pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
pub const AWS_STS_ENDPOINT: &str = "AWS_STS_ENDPOINT";
pub const AWS_EXTERNAL_ID: &str = "AWS_EXTERNAL_ID";
pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK";
/// Headers SigV4 covers, beyond the `x-amz-` / `x-amzn-` prefixes. Mirrors
/// Python's `_filter_headers_for_aws_signature` allowlist.
pub const AWS_SIGNED_HEADER_NAMES: &[&str] = &[
"host",
"content-type",
"date",
"x-amz-date",
"x-amz-security-token",
"x-amz-content-sha256",
"x-amz-algorithm",
"x-amz-credential",
"x-amz-signedheaders",
"x-amz-signature",
];
/// Headers the signer emits itself. Mirrors Python's `SIGV4_COMPUTED_HEADERS`,
/// which the reattach loop skips so a caller's copy cannot ride alongside the
/// computed one.
pub const SIGV4_COMPUTED_HEADER_NAMES: &[&str] = &[
"authorization",
"x-amz-date",
"x-amz-security-token",
"date",
];
pub const BEDROCK_SERVICE: &str = "bedrock";
pub const DEFAULT_SESSION_NAME_PREFIX: &str = "litellm-session";
pub const DEFAULT_BEDROCK_REGION: &str = "us-west-2";
pub const BEDROCK_RUNTIME_ENDPOINT_TEMPLATE: &str =
"https://bedrock-runtime.{region}.amazonaws.com";

46
litellm-rust/crates/auth-aws/src/error.rs Normal file
View file

@ -0,0 +1,46 @@
use thiserror::Error as ThisError;
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
pub enum Error {
#[error("AWS profile credentials failed: {0}")]
AwsProfile(String),
#[error("AWS default credentials failed: {0}")]
AwsDefaultChain(String),
#[error("AWS role credentials failed: {0}")]
AwsAssumeRole(String),
#[error("AWS web identity credentials failed: {0}")]
AwsWebIdentity(String),
#[error("AWS web identity expiration was invalid: {0}")]
AwsWebIdentityExpiration(String),
#[error("AWS signing parameters failed: {0}")]
AwsSigningParameters(String),
#[error("AWS signable request failed: {0}")]
AwsSignableRequest(String),
#[error("AWS request signing failed: {0}")]
AwsSigning(String),
#[error("AWS web identity response had no credentials")]
AwsMissingWebIdentityCredentials,
}
impl From<Error> for litellm_auth::Error {
fn from(error: Error) -> Self {
Self::ProviderAuthentication(error.to_string())
}
}
#[cfg(test)]
mod tests {
use super::Error;
#[test]
fn converts_to_shared_auth_error_without_losing_context() {
let error = litellm_auth::Error::from(Error::AwsProfile("profile not found".into()));
assert_eq!(
error,
litellm_auth::Error::ProviderAuthentication(
"AWS profile credentials failed: profile not found".into()
)
);
}
}

6
litellm-rust/crates/auth-aws/src/lib.rs Normal file
View file

@ -0,0 +1,6 @@
mod aws;
pub mod constants;
mod error;
pub use aws::*;
pub use error::Error;

21
litellm-rust/crates/auth-azure/Cargo.toml Normal file
View file

@ -0,0 +1,21 @@
[package]
name = "litellm-auth-azure"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-auth.workspace = true
moka.workspace = true
serde_json.workspace = true
sha2.workspace = true
strum.workspace = true
url.workspace = true
azure_core = "1.0.0"
azure_identity = { version = "1.0.0", features = ["tokio"] }
[dev-dependencies]
tokio.workspace = true

6
litellm-rust/crates/core/src/providers/azure_ai/auth/credential_provider_cache.rs → litellm-rust/crates/auth-azure/src/credential_provider_cache.rs
View file

@ -4,7 +4,7 @@ use std::sync::Arc;
use azure_core::credentials::TokenCredential;
use moka::future::Cache;
use crate::AuthError;
use litellm_auth::Error;
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub(crate) struct AzureCredentialProviderCacheKey {
@ -31,9 +31,9 @@ impl AzureCredentialProviderCache {
&self,
key: AzureCredentialProviderCacheKey,
create: F,
) -> Result<Arc<dyn TokenCredential>, AuthError>
) -> Result<Arc<dyn TokenCredential>, Error>
where
F: Future<Output = Result<Arc<dyn TokenCredential>, AuthError>>,
F: Future<Output = Result<Arc<dyn TokenCredential>, Error>>,
{
self.entries
.try_get_with(key, create)

7
litellm-rust/crates/auth-azure/src/lib.rs Normal file
View file

@ -0,0 +1,7 @@
mod credential_provider_cache;
mod native;
mod resolve;
mod types;
pub use resolve::AzureAuthService;
pub use types::AzureAuthInputs;

51
litellm-rust/crates/core/src/providers/azure_ai/auth/native.rs → litellm-rust/crates/auth-azure/src/native.rs
View file

@ -1,4 +1,3 @@
use crate::auth::error::AuthConfigurationError;
use std::sync::Arc;
use std::time::{Duration, UNIX_EPOCH};
@ -13,8 +12,8 @@ use azure_identity::{
};
use sha2::{Digest, Sha256};
use crate::AuthError;
use crate::auth::{InputSource, ResolvedCredential, SecretValue, Sourced};
use litellm_auth::Error;
use litellm_auth::{InputSource, ResolvedCredential, SecretValue, Sourced};
use super::credential_provider_cache::{
AzureCredentialProviderCache, AzureCredentialProviderCacheKey,
@ -62,7 +61,7 @@ pub(crate) struct ValidatedAzureRequest {
}
impl ValidatedAzureRequest {
pub(crate) fn new(request: NativeAzureRequest) -> Result<Self, AuthError> {
pub(crate) fn new(request: NativeAzureRequest) -> Result<Self, Error> {
validate_authority(&request)?;
let credential_source = validate_sources(&request)?;
Ok(Self {
@ -120,7 +119,7 @@ impl NativeAzureTokenAcquirer {
pub(crate) async fn acquire(
&self,
request: ValidatedAzureRequest,
) -> Result<ResolvedCredential, AuthError> {
) -> Result<ResolvedCredential, Error> {
let scope = request.request.scope().to_string();
let key = request.request.cache_key();
let transport = self.transport.clone();
@ -134,7 +133,7 @@ impl NativeAzureTokenAcquirer {
let token = credential
.get_token(&[scope.as_str()], None)
.await
.map_err(|error| AuthError::AzureTokenAcquisition(error.to_string()))?;
.map_err(|error| Error::AzureTokenAcquisition(error.to_string()))?;
let expires_on = u64::try_from(token.expires_on.unix_timestamp())
.ok()
.map(|seconds| UNIX_EPOCH + Duration::from_secs(seconds));
@ -239,7 +238,7 @@ impl NativeAzureRequest {
}
}
fn validate_authority(request: &NativeAzureRequest) -> Result<(), AuthError> {
fn validate_authority(request: &NativeAzureRequest) -> Result<(), Error> {
let authority = match request {
NativeAzureRequest::ClientSecret { authority, .. }
| NativeAzureRequest::ClientAssertion { authority, .. }
@ -251,8 +250,7 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), AuthError> {
let Some(authority) = authority else {
return Ok(());
};
let url = url::Url::parse(authority.value())
.map_err(|_| AuthError::Configuration(AuthConfigurationError::InvalidAzureAuthority))?;
let url = url::Url::parse(authority.value()).map_err(|_| Error::InvalidAzureAuthority)?;
if url.scheme() != "https"
|| url.host_str().is_none()
|| !url.username().is_empty()
@ -261,14 +259,12 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), AuthError> {
|| url.fragment().is_some()
|| !matches!(url.path(), "" | "/")
{
return Err(AuthError::Configuration(
AuthConfigurationError::InvalidAzureAuthority,
));
return Err(Error::InvalidAzureAuthority);
}
Ok(())
}
fn validate_sources(request: &NativeAzureRequest) -> Result<InputSource, AuthError> {
fn validate_sources(request: &NativeAzureRequest) -> Result<InputSource, Error> {
match request {
NativeAzureRequest::ClientSecret {
tenant_id,
@ -356,7 +352,7 @@ fn is_request_controlled<T>(value: &Sourced<T>, optional: Option<&Sourced<String
|| optional.is_some_and(|value| value.source() == InputSource::Request)
}
fn trusted_only(sources: &[InputSource]) -> Result<InputSource, AuthError> {
fn trusted_only(sources: &[InputSource]) -> Result<InputSource, Error> {
if sources.contains(&InputSource::Request) {
return mixed_sources();
}
@ -371,16 +367,14 @@ fn trusted_source(sources: &[InputSource]) -> InputSource {
}
}
fn mixed_sources<T>() -> Result<T, AuthError> {
Err(AuthError::Configuration(
AuthConfigurationError::MixedAzureCredentialSources,
))
fn mixed_sources<T>() -> Result<T, Error> {
Err(Error::MixedAzureCredentialSources)
}
fn build_credential(
request: NativeAzureRequest,
transport: Option<azure_core::http::Transport>,
) -> Result<Arc<dyn TokenCredential>, AuthError> {
) -> Result<Arc<dyn TokenCredential>, Error> {
match request {
NativeAzureRequest::ClientSecret {
tenant_id,
@ -439,11 +433,7 @@ fn build_credential(
NativeAzureRequest::DeveloperTools { .. } => DeveloperToolsCredential::new(None)
.map(|credential| credential as Arc<dyn TokenCredential>),
}
.map_err(|error| {
AuthError::Configuration(AuthConfigurationError::AzureCredentialInitialization(
error.to_string(),
))
})
.map_err(|error| Error::AzureCredentialInitialization(error.to_string()))
}
fn client_options(
@ -494,7 +484,7 @@ mod tests {
use azure_core::{Bytes, Result};
use super::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest};
use crate::auth::{InputSource, SecretValue, Sourced};
use litellm_auth::{InputSource, SecretValue, Sourced};
fn deployment<T>(value: T) -> Sourced<T> {
Sourced::new(value, InputSource::Deployment)
@ -659,9 +649,7 @@ mod tests {
assert!(matches!(
error,
crate::AuthError::Configuration(
crate::auth::error::AuthConfigurationError::MixedAzureCredentialSources
)
litellm_auth::Error::MixedAzureCredentialSources
));
}
@ -691,12 +679,7 @@ mod tests {
authority,
))
.unwrap_err();
assert!(matches!(
error,
crate::AuthError::Configuration(
crate::auth::error::AuthConfigurationError::InvalidAzureAuthority
)
));
assert!(matches!(error, litellm_auth::Error::InvalidAzureAuthority));
}
}
}

93
litellm-rust/crates/core/src/providers/azure_ai/auth/resolve.rs → litellm-rust/crates/auth-azure/src/resolve.rs
View file

@ -1,6 +1,5 @@
use crate::AuthError;
use crate::auth::error::AuthConfigurationError;
use crate::auth::{
use litellm_auth::Error;
use litellm_auth::{
CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential,
SecretValue, Sourced, TokenProviderHandle,
};
@ -37,7 +36,7 @@ pub(crate) enum AzureCredentialPlan {
}
/// Rust counterpart to Python's `get_azure_ad_token`, not `BaseAzureLLM`.
pub(crate) struct AzureAuthService {
pub struct AzureAuthService {
native: Arc<dyn AzureTokenAcquirer>,
}
@ -45,14 +44,14 @@ trait AzureTokenAcquirer: Send + Sync {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + '_>>;
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>>;
}
impl AzureTokenAcquirer for NativeAzureTokenAcquirer {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + '_>> {
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>> {
Box::pin(NativeAzureTokenAcquirer::acquire(self, request))
}
}
@ -71,17 +70,17 @@ impl AzureAuthService {
Self { native }
}
pub(crate) async fn get_azure_ad_token(
pub async fn get_azure_ad_token(
&self,
inputs: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Option<Sourced<ResolvedCredential>>, AuthError> {
) -> Result<Option<Sourced<ResolvedCredential>>, Error> {
match select_auth_plan(inputs, env_lookup)? {
AzureCredentialPlan::Supplied(credential) => Ok(Some(credential)),
AzureCredentialPlan::Caller(caller) => {
let credential = caller.acquire().await?;
if credential.secret().expose().is_empty() {
return Err(AuthError::EmptyAzureToken);
return Err(Error::EmptyAzureToken);
}
Ok(Some(Sourced::new(credential, InputSource::Deployment)))
}
@ -94,7 +93,7 @@ impl AzureAuthService {
} => {
let assertion = resolve_reference(inputs, env_lookup, reference.value())
.await?
.ok_or(AuthError::UnresolvedOidcReference)?;
.ok_or(Error::UnresolvedOidcReference)?;
let request = ValidatedAzureRequest::new(NativeAzureRequest::ClientAssertion {
tenant_id,
client_id,
@ -126,7 +125,7 @@ impl AzureAuthService {
Err(error) => failures.push(error),
}
}
Err(AuthError::CredentialChain(failures))
Err(Error::CredentialChain(failures))
}
AzureCredentialPlan::Missing => Ok(None),
}
@ -136,7 +135,7 @@ impl AzureAuthService {
pub(crate) fn select_auth_plan(
inputs: &AzureAuthInputs,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<AzureCredentialPlan, AuthError> {
) -> Result<AzureCredentialPlan, Error> {
let token = configured_secret(&inputs.azure_ad_token, AZURE_AD_TOKEN_ENV, env_lookup);
let tenant_id = configured_string(&inputs.tenant_id, AZURE_TENANT_ID_ENV, env_lookup);
let client_id = configured_string(&inputs.client_id, AZURE_CLIENT_ID_ENV, env_lookup);
@ -157,7 +156,7 @@ pub(crate) fn select_auth_plan(
.map(|selector| Sourced::new(selector, value.source()))
})
.transpose()
.map_err(|_| AuthError::Configuration(AuthConfigurationError::InvalidAzureSelector))?;
.map_err(|_| Error::InvalidAzureSelector)?;
let federated_token_file = configured_string(
&inputs.federated_token_file,
AZURE_FEDERATED_TOKEN_FILE_ENV,
@ -229,7 +228,7 @@ fn select_native_plan(
scope: Sourced<String>,
authority: Option<Sourced<String>>,
refresh_source: InputSource,
) -> Result<AzureCredentialPlan, AuthError> {
) -> Result<AzureCredentialPlan, Error> {
let selected = selector.unwrap_or_else(|| {
Sourced::new(
{
@ -247,9 +246,7 @@ fn select_native_plan(
let selection_source = selected.source();
match selected.into_value() {
AzureCredentialType::ClientSecretCredential => Err(AuthError::Configuration(
AuthConfigurationError::MissingClientSecretFields,
)),
AzureCredentialType::ClientSecretCredential => Err(Error::MissingClientSecretFields),
AzureCredentialType::WorkloadIdentityCredential => {
Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new(
workload_request(tenant_id, client_id, federated_token_file, scope, authority)?,
@ -331,17 +328,11 @@ fn workload_request(
token_file_path: Option<Sourced<String>>,
scope: Sourced<String>,
authority: Option<Sourced<String>>,
) -> Result<NativeAzureRequest, AuthError> {
) -> Result<NativeAzureRequest, Error> {
Ok(NativeAzureRequest::WorkloadIdentity {
tenant_id: tenant_id.ok_or(AuthError::Configuration(
AuthConfigurationError::MissingWorkloadTenant,
))?,
client_id: client_id.ok_or(AuthError::Configuration(
AuthConfigurationError::MissingWorkloadClient,
))?,
token_file_path: token_file_path.ok_or(AuthError::Configuration(
AuthConfigurationError::MissingWorkloadTokenFile,
))?,
tenant_id: tenant_id.ok_or(Error::MissingWorkloadTenant)?,
client_id: client_id.ok_or(Error::MissingWorkloadClient)?,
token_file_path: token_file_path.ok_or(Error::MissingWorkloadTokenFile)?,
scope,
authority,
})
@ -383,7 +374,7 @@ async fn resolve_reference(
inputs: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
reference: &CredentialRef,
) -> Result<Option<SecretValue>, AuthError> {
) -> Result<Option<SecretValue>, Error> {
let lookup = match reference {
CredentialRef::Explicit(secret) => return Ok(Some(secret.clone())),
CredentialRef::Env(name) => env_lookup(name)
@ -395,9 +386,7 @@ async fn resolve_reference(
let resolver = inputs
.credential_resolver
.as_ref()
.ok_or(AuthError::Configuration(
AuthConfigurationError::MissingHostResolver,
))?;
.ok_or(Error::MissingHostResolver)?;
resolver.resolve(reference).await?
}
};
@ -409,15 +398,13 @@ async fn resolve_reference(
fn oidc_reference(
token: &Option<Sourced<SecretValue>>,
) -> Result<Option<Sourced<CredentialRef>>, AuthError> {
) -> Result<Option<Sourced<CredentialRef>>, Error> {
let Some(token) = token.as_ref() else {
return Ok(None);
};
let value = token.value().expose();
if token.source() == InputSource::Request && value.starts_with("oidc/") {
return Err(AuthError::Configuration(
AuthConfigurationError::RequestAzureCredentialReference,
));
return Err(Error::RequestAzureCredentialReference);
}
if let Some(name) = value.strip_prefix("oidc/env/") {
return non_empty_reference(name, "OIDC environment reference")
@ -439,18 +426,14 @@ fn oidc_reference(
)));
}
if value.starts_with("oidc/") {
return Err(AuthError::Configuration(
AuthConfigurationError::UnsupportedOidcReference,
));
return Err(Error::UnsupportedOidcReference);
}
Ok(None)
}
fn non_empty_reference(value: &str, kind: &str) -> Result<String, AuthError> {
fn non_empty_reference(value: &str, kind: &str) -> Result<String, Error> {
if value.is_empty() {
return Err(AuthError::Configuration(
AuthConfigurationError::EmptyReference(kind.to_string()),
));
return Err(Error::EmptyReference(kind.to_string()));
}
Ok(value.to_string())
}
@ -466,14 +449,14 @@ mod tests {
AzureAuthService, AzureCredentialPlan, AzureTokenAcquirer, oidc_reference,
resolve_reference, select_auth_plan,
};
use crate::AuthError;
use crate::auth::ResolvedCredential;
use crate::auth::{
use crate::native::ValidatedAzureRequest;
use crate::types::AzureAuthInputs;
use litellm_auth::Error;
use litellm_auth::ResolvedCredential;
use litellm_auth::{
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialRef,
CredentialResolver, CredentialResolverHandle, InputSource, SecretValue, Sourced,
};
use crate::providers::azure_ai::auth::native::ValidatedAzureRequest;
use crate::providers::azure_ai::auth::types::AzureAuthInputs;
#[derive(Debug)]
struct FileResolver;
@ -487,9 +470,8 @@ mod tests {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> std::pin::Pin<
Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + '_>,
> {
) -> std::pin::Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>>
{
let kind = request.kind();
self.requests.lock().unwrap().push(kind);
Box::pin(async move {
@ -499,7 +481,7 @@ mod tests {
expires_on: None,
})
} else {
Err(AuthError::AzureTokenAcquisition(format!("{kind} failed")))
Err(Error::AzureTokenAcquisition(format!("{kind} failed")))
}
})
}
@ -612,12 +594,7 @@ mod tests {
})
.unwrap_err();
assert!(matches!(
error,
AuthError::Configuration(
crate::auth::error::AuthConfigurationError::RequestAzureCredentialReference
)
));
assert!(matches!(error, Error::RequestAzureCredentialReference));
}
#[tokio::test]
@ -678,6 +655,6 @@ mod tests {
.await
.unwrap_err();
assert!(matches!(error, AuthError::CredentialChain(errors) if errors.len() == 2));
assert!(matches!(error, Error::CredentialChain(errors) if errors.len() == 2));
}
}

19
litellm-rust/crates/core/src/providers/azure_ai/auth/types.rs → litellm-rust/crates/auth-azure/src/types.rs
View file

@ -1,10 +1,9 @@
use crate::auth::error::AuthConfigurationError;
use serde_json::{Map, Value};
use std::collections::BTreeMap;
use strum::EnumString;
use crate::AuthError;
use crate::auth::{
use litellm_auth::Error;
use litellm_auth::{
CredentialResolverHandle, InputSource, SecretValue, Sourced, TokenProviderHandle,
};
@ -54,14 +53,14 @@ pub struct AzureAuthInputs {
impl AzureAuthInputs {
#[cfg(test)]
pub fn from_optional_params(params: &Map<String, Value>) -> Result<Self, AuthError> {
pub fn from_optional_params(params: &Map<String, Value>) -> Result<Self, Error> {
Self::from_sourced_optional_params(params, &BTreeMap::new())
}
pub fn from_sourced_optional_params(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
) -> Result<Self, AuthError> {
) -> Result<Self, Error> {
Ok(Self {
azure_ad_token: secret_config(params, sources, "azure_ad_token")?,
azure_ad_token_provider: None,
@ -88,15 +87,13 @@ fn string_config(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
name: &str,
) -> Result<ConfigValue<String>, AuthError> {
) -> Result<ConfigValue<String>, Error> {
let source = source_for(sources, name);
match params.get(name) {
None => Ok(ConfigValue::Absent),
Some(Value::Null) => Ok(ConfigValue::ExplicitNone(source)),
Some(Value::String(value)) => Ok(ConfigValue::Value(Sourced::new(value.clone(), source))),
Some(_) => Err(AuthError::Configuration(
AuthConfigurationError::InvalidFieldType(name.to_string()),
)),
Some(_) => Err(Error::InvalidFieldType(name.to_string())),
}
}
@ -104,7 +101,7 @@ fn secret_config(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
name: &str,
) -> Result<ConfigValue<SecretValue>, AuthError> {
) -> Result<ConfigValue<SecretValue>, Error> {
Ok(match string_config(params, sources, name)? {
ConfigValue::Absent => ConfigValue::Absent,
ConfigValue::ExplicitNone(source) => ConfigValue::ExplicitNone(source),
@ -123,7 +120,7 @@ mod tests {
use std::collections::BTreeMap;
use super::{AzureAuthInputs, AzureCredentialType, ConfigValue};
use crate::auth::{InputSource, Sourced};
use litellm_auth::{InputSource, Sourced};
#[test]
fn selector_parsing_is_exact() {

16
litellm-rust/crates/auth-gcp/Cargo.toml Normal file
View file

@ -0,0 +1,16 @@
[package]
name = "litellm-auth-gcp"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-auth.workspace = true
moka.workspace = true
serde_json.workspace = true
sha2.workspace = true
tokio.workspace = true
gcp_auth = "0.12.7"

72
litellm-rust/crates/core/src/auth/vertex.rs → litellm-rust/crates/auth-gcp/src/lib.rs
View file

@ -9,9 +9,8 @@ use moka::future::Cache;
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use crate::auth::error::AuthConfigurationError;
use crate::auth::http::apply_credential;
use crate::auth::{AuthError, CredentialPlacement, InputSource, SecretValue, Sourced};
use litellm_auth::http::apply_credential;
use litellm_auth::{CredentialPlacement, Error, InputSource, SecretValue, Sourced};
const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token";
@ -24,17 +23,17 @@ const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION";
const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION";
#[derive(Clone, Debug, Default)]
pub(crate) struct VertexConfig {
pub struct VertexConfig {
credentials: Option<Sourced<SecretValue>>,
project_id: Option<String>,
location: Option<String>,
}
impl VertexConfig {
pub(crate) fn from_sourced_optional_params(
pub fn from_sourced_optional_params(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
) -> Result<Self, AuthError> {
) -> Result<Self, Error> {
Ok(Self {
credentials: optional_credentials(
params,
@ -46,16 +45,16 @@ impl VertexConfig {
})
}
pub(crate) fn project_id(&self) -> Option<&str> {
pub fn project_id(&self) -> Option<&str> {
self.project_id.as_deref()
}
pub(crate) fn location(&self) -> Option<&str> {
pub fn location(&self) -> Option<&str> {
self.location.as_deref()
}
}
pub(crate) struct VertexEnvironment {
pub struct VertexEnvironment {
pub headers: Vec<(String, String)>,
pub project_id: String,
}
@ -65,7 +64,7 @@ struct VertexAccessToken {
project_id: String,
}
pub(crate) fn get_vertex_ai_project(
pub fn get_vertex_ai_project(
config: &VertexConfig,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<String> {
@ -75,7 +74,7 @@ pub(crate) fn get_vertex_ai_project(
.or_else(|| non_empty_env(env_lookup, VERTEXAI_PROJECT_ENV))
}
pub(crate) fn get_vertex_ai_location(
pub fn get_vertex_ai_location(
config: &VertexConfig,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<String> {
@ -87,7 +86,7 @@ pub(crate) fn get_vertex_ai_location(
}
#[derive(Clone)]
pub(crate) struct VertexAuth {
pub struct VertexAuth {
providers: Cache<CredentialCacheKey, Arc<dyn VertexTokenSource>>,
loader: Arc<dyn VertexProviderLoader>,
}
@ -106,14 +105,13 @@ impl VertexAuth {
}
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) async fn validate_environment(
pub async fn validate_environment(
&self,
headers: Vec<(String, String)>,
api_key: Option<&str>,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<VertexEnvironment, AuthError> {
) -> Result<VertexEnvironment, Error> {
let has_authorization = headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case("Authorization"));
@ -161,7 +159,7 @@ impl VertexAuth {
&self,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<VertexAccessToken, AuthError> {
) -> Result<VertexAccessToken, Error> {
let provider = self.load_provider(config, env_lookup).await?;
let (token, project_id) = tokio::try_join!(provider.token(), provider.project_id())?;
Ok(VertexAccessToken { token, project_id })
@ -171,7 +169,7 @@ impl VertexAuth {
&self,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Arc<dyn VertexTokenSource>, AuthError> {
) -> Result<Arc<dyn VertexTokenSource>, Error> {
let source = credential_source(config, env_lookup);
let key = source.cache_key();
self.providers
@ -190,7 +188,7 @@ trait VertexProviderLoader: Send + Sync {
fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>>;
}
type VertexAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, AuthError>> + Send + 'a>>;
type VertexAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
struct GcpTokenSource(Arc<dyn TokenProvider>);
@ -250,7 +248,7 @@ impl VertexProviderLoader for GcpProviderLoader {
}
}
fn validate_request_credentials(configured: &str) -> Result<&str, AuthError> {
fn validate_request_credentials(configured: &str) -> Result<&str, Error> {
let token_uri = serde_json::from_str::<Value>(configured)
.ok()
.and_then(|credentials| {
@ -260,7 +258,7 @@ fn validate_request_credentials(configured: &str) -> Result<&str, AuthError> {
.map(str::to_string)
});
if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) {
return Err(AuthConfigurationError::RequestVertexTokenEndpoint.into());
return Err(Error::RequestVertexTokenEndpoint);
}
Ok(configured)
}
@ -322,7 +320,7 @@ fn optional_credentials(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
names: &[&str],
) -> Result<Option<Sourced<SecretValue>>, AuthError> {
) -> Result<Option<Sourced<SecretValue>>, Error> {
for name in names {
let source = source_for(sources, name);
match params.get(*name) {
@ -337,17 +335,10 @@ fn optional_credentials(
.map(SecretValue::new)
.map(|value| Sourced::new(value, source))
.map(Some)
.map_err(|error| {
AuthError::Configuration(AuthConfigurationError::InvalidFieldType(format!(
"{}: {error}",
names[0]
)))
});
.map_err(|error| Error::InvalidFieldType(format!("{}: {error}", names[0])));
}
Some(_) => {
return Err(AuthError::Configuration(
AuthConfigurationError::InvalidFieldType(names[0].to_string()),
));
return Err(Error::InvalidFieldType(names[0].to_string()));
}
}
}
@ -358,19 +349,14 @@ fn source_for(sources: &BTreeMap<String, InputSource>, name: &str) -> InputSourc
sources.get(name).copied().unwrap_or_default()
}
fn optional_string(
params: &Map<String, Value>,
names: &[&str],
) -> Result<Option<String>, AuthError> {
fn optional_string(params: &Map<String, Value>, names: &[&str]) -> Result<Option<String>, Error> {
for name in names {
match params.get(*name) {
None | Some(Value::Null) => continue,
Some(Value::String(value)) if value.trim().is_empty() => continue,
Some(Value::String(value)) => return Ok(Some(value.clone())),
Some(_) => {
return Err(AuthError::Configuration(
AuthConfigurationError::InvalidFieldType(names[0].to_string()),
));
return Err(Error::InvalidFieldType(names[0].to_string()));
}
}
}
@ -383,8 +369,8 @@ fn non_empty_env(env_lookup: &dyn Fn(&str) -> Option<String>, name: &str) -> Opt
.filter(|value| !value.is_empty())
}
fn auth_acquisition_error(error: gcp_auth::Error) -> AuthError {
AuthError::VertexTokenAcquisition(error.to_string())
fn auth_acquisition_error(error: gcp_auth::Error) -> Error {
Error::VertexTokenAcquisition(error.to_string())
}
#[cfg(test)]
@ -538,15 +524,11 @@ mod tests {
);
assert!(matches!(
validate_request_credentials(r#"{"token_uri":"http://127.0.0.1/token"}"#),
Err(AuthError::Configuration(
AuthConfigurationError::RequestVertexTokenEndpoint
))
Err(Error::RequestVertexTokenEndpoint)
));
assert!(matches!(
validate_request_credentials("{}"),
Err(AuthError::Configuration(
AuthConfigurationError::RequestVertexTokenEndpoint
))
Err(Error::RequestVertexTokenEndpoint)
));
}

15
litellm-rust/crates/auth/Cargo.toml Normal file
View file

@ -0,0 +1,15 @@
[package]
name = "litellm-auth"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
serde.workspace = true
subtle.workspace = true
thiserror.workspace = true
veil.workspace = true
[dev-dependencies]
tokio.workspace = true

18
litellm-rust/crates/core/src/auth/credential.rs → litellm-rust/crates/auth/src/credential.rs
View file

@ -5,7 +5,7 @@ use std::sync::Arc;
use veil::Redact;
use crate::AuthError;
use crate::Error;
use super::{ResolvedCredential, SecretValue, TokenProviderHandle};
@ -48,7 +48,7 @@ pub enum CredentialLookup {
}
pub type CredentialLookupFuture<'a> =
Pin<Box<dyn Future<Output = Result<CredentialLookup, AuthError>> + Send + 'a>>;
Pin<Box<dyn Future<Output = Result<CredentialLookup, Error>> + Send + 'a>>;
pub trait CredentialResolver: std::fmt::Debug + Send + Sync {
fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a>;
@ -62,7 +62,7 @@ impl CredentialResolverHandle {
Self(resolver)
}
pub async fn resolve(&self, reference: &CredentialRef) -> Result<CredentialLookup, AuthError> {
pub async fn resolve(&self, reference: &CredentialRef) -> Result<CredentialLookup, Error> {
self.0.resolve(reference).await
}
}
@ -84,7 +84,7 @@ impl CredentialPlan {
pub async fn resolve(
&self,
resolver: &CredentialResolverHandle,
) -> Result<CredentialPlanResolution, AuthError> {
) -> Result<CredentialPlanResolution, Error> {
match self {
Self::Static(CredentialRef::Explicit(secret)) => Ok(
CredentialPlanResolution::Resolved(ResolvedCredential::Static(secret.clone())),
@ -103,7 +103,7 @@ impl CredentialPlan {
Self::Caller(caller) => {
let credential = caller.acquire().await?;
if credential.secret().expose().is_empty() {
return Err(AuthError::EmptyCallerCredential);
return Err(Error::EmptyCallerCredential);
}
Ok(CredentialPlanResolution::Resolved(credential))
}
@ -119,8 +119,8 @@ mod tests {
CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution,
CredentialRef, CredentialResolver, CredentialResolverHandle,
};
use crate::AuthError;
use crate::auth::SecretValue;
use crate::Error;
use crate::SecretValue;
#[derive(Debug)]
struct HostResolver;
@ -164,7 +164,7 @@ mod tests {
impl CredentialResolver for FailingResolver {
fn resolve<'a>(&'a self, _reference: &'a CredentialRef) -> CredentialLookupFuture<'a> {
Box::pin(async { Err(AuthError::UnresolvedOidcReference) })
Box::pin(async { Err(Error::UnresolvedOidcReference) })
}
}
@ -178,6 +178,6 @@ mod tests {
.await
.expect_err("acquisition errors cannot become fallback");
assert_eq!(error, AuthError::UnresolvedOidcReference);
assert_eq!(error, Error::UnresolvedOidcReference);
}
}

120
litellm-rust/crates/auth/src/error.rs Normal file
View file

@ -0,0 +1,120 @@
use thiserror::Error as ThisError;
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
pub enum Error {
#[error("invalid authentication configuration: credential header already exists")]
ExistingCredentialHeader,
#[error(
"invalid authentication configuration: credential plan is not allowed by the provider auth policy"
)]
DisallowedCredentialPlan,
#[error("invalid authentication configuration: credential cannot be empty")]
EmptyCredential,
#[error("invalid authentication configuration: invalid Azure credential selector")]
InvalidAzureSelector,
#[error(
"invalid authentication configuration: ClientSecretCredential requires tenant_id, client_id, and client_secret"
)]
MissingClientSecretFields,
#[error("invalid authentication configuration: WorkloadIdentityCredential requires tenant_id")]
MissingWorkloadTenant,
#[error("invalid authentication configuration: WorkloadIdentityCredential requires client_id")]
MissingWorkloadClient,
#[error(
"invalid authentication configuration: WorkloadIdentityCredential requires azure_federated_token_file"
)]
MissingWorkloadTokenFile,
#[error(
"invalid authentication configuration: credential reference requires a host credential resolver"
)]
MissingHostResolver,
#[error(
"invalid authentication configuration: caller credential plan requires provider-specific inputs"
)]
MissingCallerInputs,
#[error("invalid authentication configuration: credential header {0} already exists")]
DuplicateHeader(&'static str),
#[error("invalid authentication configuration: {0} must be a string or null")]
InvalidFieldType(String),
#[error("invalid authentication configuration: unsupported OIDC reference")]
UnsupportedOidcReference,
#[error("invalid authentication configuration: {0} cannot be empty")]
EmptyReference(String),
#[error("invalid authentication configuration: Azure credential initialization failed: {0}")]
AzureCredentialInitialization(String),
#[error(
"invalid authentication configuration: Azure authority must be an HTTPS origin without credentials, query, or fragment"
)]
InvalidAzureAuthority,
#[error(
"invalid authentication configuration: request-controlled Azure auth inputs cannot be combined with host credentials"
)]
MixedAzureCredentialSources,
#[error(
"invalid authentication configuration: request-controlled Azure credential references are not allowed"
)]
RequestAzureCredentialReference,
#[error(
"invalid authentication configuration: host credentials cannot be sent to a request-controlled Azure endpoint"
)]
RequestAzureCredentialDestination,
#[error(
"invalid authentication configuration: credentials cannot be sent to a request-controlled Vertex AI endpoint"
)]
RequestVertexCredentialDestination,
#[error(
"invalid authentication configuration: request-controlled Vertex credentials must use the canonical Google OAuth token endpoint"
)]
RequestVertexTokenEndpoint,
#[error("credential acquisition failed: {0}")]
AzureTokenAcquisition(String),
#[error("credential acquisition failed: Vertex AI credentials: {0}")]
VertexTokenAcquisition(String),
#[error("{0}")]
ProviderAuthentication(String),
#[error("credential acquisition failed: {}", .0.iter().map(ToString::to_string).collect::<Vec<_>>().join("; "))]
CredentialChain(Vec<Error>),
#[error("credential caller failed: credential caller returned an empty credential")]
EmptyCallerCredential,
#[error("credential caller failed: Azure AD token provider returned an empty token")]
EmptyAzureToken,
#[error("credential acquisition failed: Azure OIDC reference did not resolve to a value")]
UnresolvedOidcReference,
#[error(
"Missing {provider} API Key - Set `api_key` or the {environment_variable} environment variable"
)]
MissingApiKey {
provider: &'static str,
environment_variable: &'static str,
},
#[error(
"Missing {provider} API Base - Set {environment_variable} environment variable or pass api_base parameter"
)]
MissingApiBase {
provider: &'static str,
environment_variable: &'static str,
},
#[error(
"Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. Expected format: https://<resource-name>.services.ai.azure.com/anthropic"
)]
MissingAzureApiBase,
#[error("invalid authentication header")]
InvalidHeader,
}
#[cfg(test)]
mod tests {
use super::Error;
#[test]
fn missing_api_key_names_provider_and_environment_variable() {
assert_eq!(
Error::MissingApiKey {
provider: "Anthropic",
environment_variable: "ANTHROPIC_API_KEY",
}
.to_string(),
"Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY environment variable"
);
}
}

15
litellm-rust/crates/core/src/auth/http.rs → litellm-rust/crates/auth/src/http.rs
View file

@ -1,5 +1,4 @@
use crate::AuthError;
use crate::auth::error::AuthConfigurationError;
use crate::Error;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CredentialPlacement {
@ -16,23 +15,19 @@ impl CredentialPlacement {
}
}
pub(crate) fn apply_credential(
pub fn apply_credential(
headers: Vec<(String, String)>,
credential: &str,
placement: CredentialPlacement,
) -> Result<Vec<(String, String)>, AuthError> {
) -> Result<Vec<(String, String)>, Error> {
if credential.trim().is_empty() {
return Err(AuthError::Configuration(
AuthConfigurationError::EmptyCredential,
));
return Err(Error::EmptyCredential);
}
if headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case(placement.header_name()))
{
return Err(AuthError::Configuration(
AuthConfigurationError::DuplicateHeader(placement.header_name()),
));
return Err(Error::DuplicateHeader(placement.header_name()));
}
let value = match placement {
CredentialPlacement::Bearer => format!("Bearer {credential}"),

7
litellm-rust/crates/core/src/auth/mod.rs → litellm-rust/crates/auth/src/lib.rs
View file

@ -1,8 +1,6 @@
mod credential;
pub mod error;
pub(crate) mod vertex;
pub use error::AuthError;
pub(crate) mod http;
mod error;
pub mod http;
mod policy;
mod secret;
mod token;
@ -51,6 +49,7 @@ pub use credential::{
CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
credential_default_fields, credential_index,
};
pub use error::Error;
pub use http::{CredentialPlacement, RequestAuth};
pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
pub use secret::SecretValue;

23
litellm-rust/crates/core/src/auth/policy.rs → litellm-rust/crates/auth/src/policy.rs
View file

@ -1,5 +1,4 @@
use crate::AuthError;
use crate::auth::error::AuthConfigurationError;
use crate::Error;
use super::http::apply_credential;
use super::{CredentialPlacement, ResolvedCredential};
@ -46,22 +45,18 @@ impl ProviderAuthPolicy {
headers: Vec<(String, String)>,
kind: CredentialPlanKind,
credential: &ResolvedCredential,
) -> Result<Vec<(String, String)>, AuthError> {
) -> Result<Vec<(String, String)>, Error> {
if self.has_existing_credential(&headers) {
return match self.existing_header_behavior {
ExistingHeaderBehavior::Preserve => Ok(headers),
ExistingHeaderBehavior::Reject => Err(AuthError::Configuration(
AuthConfigurationError::ExistingCredentialHeader,
)),
ExistingHeaderBehavior::Reject => Err(Error::ExistingCredentialHeader),
};
}
let rule =
self.rules
.iter()
.find(|rule| rule.kind == kind)
.ok_or(AuthError::Configuration(
AuthConfigurationError::DisallowedCredentialPlan,
))?;
let rule = self
.rules
.iter()
.find(|rule| rule.kind == kind)
.ok_or(Error::DisallowedCredentialPlan)?;
apply_credential(headers, credential.secret().expose(), rule.placement)
}
}
@ -69,7 +64,7 @@ impl ProviderAuthPolicy {
#[cfg(test)]
mod tests {
use super::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
use crate::auth::{CredentialPlacement, ResolvedCredential, SecretValue};
use crate::{CredentialPlacement, ResolvedCredential, SecretValue};
const RULES: &[CredentialRule] = &[CredentialRule {
kind: CredentialPlanKind::Static,

0
litellm-rust/crates/core/src/auth/secret.rs → litellm-rust/crates/auth/src/secret.rs
View file

6
litellm-rust/crates/core/src/auth/token.rs → litellm-rust/crates/auth/src/token.rs
View file

@ -5,7 +5,7 @@ use std::time::SystemTime;
use veil::Redact;
use crate::AuthError;
use crate::Error;
use super::secret::SecretValue;
@ -27,7 +27,7 @@ impl ResolvedCredential {
}
pub type TokenFuture<'a> =
Pin<Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + 'a>>;
Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + 'a>>;
pub trait TokenProvider: std::fmt::Debug + Send + Sync {
fn acquire(&self) -> TokenFuture<'_>;
@ -41,7 +41,7 @@ impl TokenProviderHandle {
Self(caller)
}
pub async fn acquire(&self) -> Result<ResolvedCredential, AuthError> {
pub async fn acquire(&self) -> Result<ResolvedCredential, Error> {
self.0.acquire().await
}
}

14
litellm-rust/crates/cache-memory/Cargo.toml Normal file
View file

@ -0,0 +1,14 @@
[package]
name = "litellm-cache-memory"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
serde_json.workspace = true
[dev-dependencies]
rstest.workspace = true
tokio.workspace = true

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

@ -0,0 +1,254 @@
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap};
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs,
Error,
};
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
const DEFAULT_TTL: Duration = Duration::from_secs(600);
type ValueMeasure<V> = Arc<dyn Fn(&V) -> Result<usize, Error> + Send + Sync>;
type ValueValidator<V> = Arc<dyn Fn(&V) -> Result<(), Error> + Send + Sync>;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CacheWrite {
Stored,
Disabled,
TooLarge,
}
struct CacheState<V> {
values: HashMap<String, V>,
expirations: HashMap<String, Duration>,
expiration_heap: BinaryHeap<Reverse<(Duration, String)>>,
}
pub struct InMemoryCache<V: Clone> {
state: Mutex<CacheState<V>>,
max_size_in_memory: usize,
default_ttl: Duration,
max_entry_bytes: Option<usize>,
measure_value: Option<ValueMeasure<V>>,
validate_value: Option<ValueValidator<V>>,
now: Arc<dyn Fn() -> Duration + Send + Sync>,
}
impl<V: Clone> Default for InMemoryCache<V> {
fn default() -> Self {
Self::new(None, None)
}
}
impl<V: Clone> InMemoryCache<V> {
pub fn new(max_size_in_memory: Option<usize>, default_ttl: Option<Duration>) -> Self {
Self::with_clock(max_size_in_memory, default_ttl, || {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
})
}
pub fn with_clock(
max_size_in_memory: Option<usize>,
default_ttl: Option<Duration>,
now: impl Fn() -> Duration + Send + Sync + 'static,
) -> Self {
Self::with_clock_and_size_measurement(max_size_in_memory, default_ttl, None, None, now)
}
pub fn with_clock_and_size_measurement(
max_size_in_memory: Option<usize>,
default_ttl: Option<Duration>,
max_entry_bytes: Option<usize>,
measure_value: Option<ValueMeasure<V>>,
now: impl Fn() -> Duration + Send + Sync + 'static,
) -> Self {
Self {
state: Mutex::new(CacheState {
values: HashMap::new(),
expirations: HashMap::new(),
expiration_heap: BinaryHeap::new(),
}),
max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
max_entry_bytes,
measure_value,
validate_value: None,
now: Arc::new(now),
}
}
pub fn set_cache(
&self,
key: impl Into<String>,
value: V,
ttl: Option<Duration>,
) -> Result<CacheWrite, Error> {
if self.max_size_in_memory == 0 {
return Ok(CacheWrite::Disabled);
}
if let Some(validate) = &self.validate_value {
validate(&value)?;
}
if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value)
&& measure(&value)? > limit
{
return Ok(CacheWrite::TooLarge);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now);
let key = key.into();
state.values.insert(key.clone(), value);
let expiration = state.expirations.get(&key).copied();
if expiration.is_none_or(|expiration| expiration < now) {
let expiration = now + ttl.unwrap_or(self.default_ttl);
state.expirations.insert(key.clone(), expiration);
state.expiration_heap.push(Reverse((expiration, key)));
}
Ok(CacheWrite::Stored)
}
pub fn get_cache(&self, key: &str) -> Result<Option<V>, Error> {
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
if state
.expirations
.get(key)
.is_some_and(|expiration| *expiration < now)
{
Self::remove(&mut state, key);
}
Ok(state.values.get(key).cloned())
}
pub fn expires_at(&self, key: &str) -> Result<Option<Duration>, Error> {
Ok(self
.state
.lock()
.map_err(|_| Error::Unavailable)?
.expirations
.get(key)
.copied())
}
pub fn delete_cache(&self, key: &str) -> Result<(), Error> {
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::remove(&mut state, key);
Ok(())
}
pub fn flush_cache(&self) -> Result<(), Error> {
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
state.values.clear();
state.expirations.clear();
state.expiration_heap.clear();
Ok(())
}
fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration) {
while let Some(Reverse((expiration, key))) = state.expiration_heap.peek().cloned() {
if state.expirations.get(&key).copied() != Some(expiration) {
state.expiration_heap.pop();
} else if expiration <= now {
state.expiration_heap.pop();
Self::remove(state, &key);
} else {
break;
}
}
while state.values.len() >= capacity {
let Some(Reverse((expiration, key))) = state.expiration_heap.pop() else {
break;
};
if state.expirations.get(&key).copied() == Some(expiration) {
Self::remove(state, &key);
}
}
}
fn remove(state: &mut CacheState<V>, key: &str) {
state.values.remove(key);
state.expirations.remove(key);
}
}
impl InMemoryCache<CacheEntry> {
pub fn response_cache(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self {
Self::response_cache_with_clock(capacity, ttl, max_entry_bytes, || {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
})
}
pub fn response_cache_with_clock(
capacity: usize,
ttl: Duration,
max_entry_bytes: usize,
now: impl Fn() -> Duration + Send + Sync + 'static,
) -> Self {
let mut cache = Self::with_clock_and_size_measurement(
Some(capacity),
Some(ttl),
Some(max_entry_bytes),
Some(Arc::new(|entry: &CacheEntry| {
serde_json::to_vec(entry)
.map(|bytes| bytes.len())
.map_err(|_| Error::InvalidEntry)
})),
now,
);
cache.validate_value = Some(Arc::new(|entry: &CacheEntry| {
entry
.timestamp
.is_finite()
.then_some(())
.ok_or(Error::InvalidEntry)
}));
cache
}
}
impl BaseCache for InMemoryCache<CacheEntry> {
type Value = CacheEntry;
fn default_ttl(&self) -> Duration {
self.default_ttl
}
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> {
let ttl = self.get_ttl(&kwargs);
self.set_cache(key, value, Some(ttl)).map(|_| ())
}
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
self.get_cache(key)
}
fn delete_cache(&self, key: &str) -> Result<(), Error> {
self.delete_cache(key)
}
fn flush_cache(&self) -> Result<(), Error> {
self.flush_cache()
}
fn disconnect(&self) -> CacheFuture<'_, ()> {
Box::pin(async { Ok(()) })
}
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
Box::pin(async {
Ok(CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "In-memory cache connection test successful".into(),
error: None,
})
})
}
}

3
litellm-rust/crates/cache-memory/src/lib.rs Normal file
View file

@ -0,0 +1,3 @@
mod cache;
pub use cache::{CacheWrite, InMemoryCache};

158
litellm-rust/crates/cache-memory/tests/cache.rs Normal file
View file

@ -0,0 +1,158 @@
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use litellm_cache::{BaseCache, CacheConnectionStatus, CacheEntry, Error};
use litellm_cache_memory::{CacheWrite, InMemoryCache};
use rstest::{fixture, rstest};
#[fixture]
fn clock() -> Arc<AtomicU64> {
Arc::new(AtomicU64::new(100))
}
fn cache(clock: Arc<AtomicU64>, capacity: usize) -> InMemoryCache<String> {
InMemoryCache::with_clock(Some(capacity), Some(Duration::from_secs(60)), move || {
Duration::from_secs(clock.load(Ordering::SeqCst))
})
}
#[rstest]
fn default_explicit_and_override_ttls_follow_python_rules(clock: Arc<AtomicU64>) {
let cache = cache(clock.clone(), 4);
cache.set_cache("key", "first".into(), None).unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(160))
);
cache
.set_cache("key", "second".into(), Some(Duration::from_secs(10)))
.unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(160))
);
clock.store(160, Ordering::SeqCst);
assert_eq!(cache.get_cache("key").unwrap(), Some("second".into()));
clock.store(161, Ordering::SeqCst);
assert_eq!(cache.get_cache("key").unwrap(), None);
cache
.set_cache("key", "third".into(), Some(Duration::from_secs(10)))
.unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(171))
);
}
#[rstest]
fn write_at_expiry_boundary_refreshes_ttl(clock: Arc<AtomicU64>) {
let cache = cache(clock.clone(), 4);
cache
.set_cache("key", "first".into(), Some(Duration::from_secs(10)))
.unwrap();
clock.store(110, Ordering::SeqCst);
cache
.set_cache("key", "second".into(), Some(Duration::from_secs(10)))
.unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(120))
);
clock.store(115, Ordering::SeqCst);
assert_eq!(cache.get_cache("key").unwrap(), Some("second".into()));
}
#[rstest]
fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Arc<AtomicU64>) {
let cache = cache(clock, 2);
cache
.set_cache("early", "a".into(), Some(Duration::from_secs(10)))
.unwrap();
cache
.set_cache("late", "b".into(), Some(Duration::from_secs(20)))
.unwrap();
cache.delete_cache("early").unwrap();
cache
.set_cache("new", "c".into(), Some(Duration::from_secs(30)))
.unwrap();
assert_eq!(cache.get_cache("late").unwrap(), Some("b".into()));
cache
.set_cache("last", "d".into(), Some(Duration::from_secs(40)))
.unwrap();
assert_eq!(cache.get_cache("late").unwrap(), None);
}
#[test]
fn disabled_size_limited_and_synchronized_response_writes_are_observable() {
let disabled = InMemoryCache::<CacheEntry>::response_cache(0, Duration::from_secs(60), 80);
assert_eq!(
disabled
.set_cache(
"a",
CacheEntry {
timestamp: 1.0,
response: serde_json::json!("x")
},
None
)
.unwrap(),
CacheWrite::Disabled
);
let cache = InMemoryCache::<CacheEntry>::response_cache(2, Duration::from_secs(60), 80);
assert_eq!(
cache
.set_cache(
"large",
CacheEntry {
timestamp: 1.0,
response: serde_json::json!("x".repeat(100))
},
None
)
.unwrap(),
CacheWrite::TooLarge
);
cache
.set_cache(
"small",
CacheEntry {
timestamp: 1.0,
response: serde_json::json!("ok"),
},
None,
)
.unwrap();
assert!(cache.get_cache("small").unwrap().is_some());
assert_eq!(
cache
.set_cache(
"invalid",
CacheEntry {
timestamp: f64::NAN,
response: serde_json::json!("bad"),
},
None,
)
.unwrap_err(),
Error::InvalidEntry
);
cache.delete_cache("small").unwrap();
cache.flush_cache().unwrap();
}
#[tokio::test]
async fn connection_test_matches_python_result_contract() {
let cache = InMemoryCache::<CacheEntry>::default();
let result = BaseCache::test_connection(&cache).await.unwrap();
assert_eq!(result.status, CacheConnectionStatus::Success);
assert_eq!(result.message, "In-memory cache connection test successful");
assert_eq!(result.error, None);
assert_eq!(
serde_json::to_value(result).unwrap(),
serde_json::json!({
"status": "success",
"message": "In-memory cache connection test successful"
})
);
}

15
litellm-rust/crates/cache-redis/Cargo.toml Normal file
View file

@ -0,0 +1,15 @@
[package]
name = "litellm-cache-redis"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
redis = "1.7.0"
serde_json.workspace = true
tokio.workspace = true
[dev-dependencies]
redis-test = "1.0.4"

315
litellm-rust/crates/cache-redis/src/cache.rs Normal file
View file

@ -0,0 +1,315 @@
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::Duration;
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs,
Error,
};
use redis::Commands;
const DEFAULT_TTL: Duration = Duration::from_secs(600);
const KEY_PREFIX: &str = "litellm-cache:";
pub struct RedisCache<C = redis::Connection> {
connection: Arc<Mutex<C>>,
default_ttl: Duration,
}
impl RedisCache<redis::Connection> {
pub fn new(url: &str, default_ttl: Option<Duration>) -> Result<Self, Error> {
let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?;
let connection = client.get_connection().map_err(|_| Error::Unavailable)?;
Ok(Self::with_connection(connection, default_ttl))
}
}
impl<C> RedisCache<C>
where
C: redis::ConnectionLike + Send + 'static,
{
fn with_connection(connection: C, default_ttl: Option<Duration>) -> Self {
Self {
connection: Arc::new(Mutex::new(connection)),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
}
}
fn connection(&self) -> Result<MutexGuard<'_, C>, Error> {
self.connection.lock().map_err(|_| Error::Unavailable)
}
fn namespaced_key(key: &str) -> String {
format!("{KEY_PREFIX}{key}")
}
fn namespaced_pattern() -> &'static str {
const PATTERN: &str = "litellm-cache:*";
PATTERN
}
fn encode(value: &CacheEntry) -> Result<Vec<u8>, Error> {
serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
}
fn decode(value: Vec<u8>) -> Result<CacheEntry, Error> {
serde_json::from_slice(&value).map_err(|_| Error::InvalidEntry)
}
fn ttl_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}
fn run_blocking<T, F>(connection: Arc<Mutex<C>>, operation: F) -> CacheFuture<'static, T>
where
T: Send + 'static,
F: FnOnce(&mut C) -> Result<T, Error> + Send + 'static,
{
Box::pin(async move {
tokio::task::spawn_blocking(move || {
let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
operation(&mut connection)
})
.await
.map_err(|_| Error::Unavailable)?
})
}
}
impl<C> BaseCache for RedisCache<C>
where
C: redis::ConnectionLike + Send + 'static,
{
type Value = CacheEntry;
fn default_ttl(&self) -> Duration {
self.default_ttl
}
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> {
let payload = Self::encode(&value)?;
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
self.connection()?
.set_ex::<_, _, ()>(Self::namespaced_key(key), payload, ttl)
.map_err(|_| Error::Unavailable)
}
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
self.connection()?
.get::<_, Option<Vec<u8>>>(Self::namespaced_key(key))
.map_err(|_| Error::Unavailable)?
.map(Self::decode)
.transpose()
}
fn delete_cache(&self, key: &str) -> Result<(), Error> {
self.connection()?
.del::<_, ()>(Self::namespaced_key(key))
.map_err(|_| Error::Unavailable)
}
fn flush_cache(&self) -> Result<(), Error> {
let mut connection = self.connection()?;
let keys = connection
.scan_match(Self::namespaced_pattern())
.map_err(|_| Error::Unavailable)?
.collect::<redis::RedisResult<Vec<String>>>()
.map_err(|_| Error::Unavailable)?;
if keys.is_empty() {
return Ok(());
}
connection
.del::<_, usize>(keys)
.map(|_| ())
.map_err(|_| Error::Unavailable)
}
fn async_set_cache<'a>(
&'a self,
key: &'a str,
value: Self::Value,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
let payload = Self::encode(&value);
let key = Self::namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection
.set_ex::<_, _, ()>(key, payload?, ttl)
.map_err(|_| Error::Unavailable)
})
}
fn async_get_cache<'a>(
&'a self,
key: &'a str,
_: &'a CacheKwargs,
) -> CacheFuture<'a, Option<Self::Value>> {
let key = Self::namespaced_key(key);
Box::pin(async move {
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection
.get::<_, Option<Vec<u8>>>(key)
.map_err(|_| Error::Unavailable)
})
.await?
.map(Self::decode)
.transpose()
})
}
fn async_set_cache_pipeline<'a>(
&'a self,
cache_list: Vec<(String, Self::Value)>,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
let entries = cache_list
.into_iter()
.map(|(key, value)| {
Self::encode(&value).map(|payload| (Self::namespaced_key(&key), payload))
})
.collect::<Result<Vec<_>, _>>();
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
for (key, payload) in entries? {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(|_| Error::Unavailable)?;
}
Ok(())
})
}
fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> {
let key = Self::namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
})
}
fn disconnect(&self) -> CacheFuture<'_, ()> {
Box::pin(async { Ok(()) })
}
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
Box::pin(async move {
Self::run_blocking(Arc::clone(&self.connection), |connection| {
redis::cmd("PING")
.query::<String>(connection)
.map_err(|_| Error::Unavailable)
})
.await?;
Ok(CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Redis cache connection test successful".into(),
error: None,
})
})
}
}
#[cfg(test)]
mod tests {
use super::RedisCache;
use litellm_cache::{BaseCache, CacheEntry, CacheKwargs};
use redis_test::{MockCmd, MockRedisConnection};
use serde_json::json;
use std::time::Duration;
fn entry() -> CacheEntry {
CacheEntry {
timestamp: 123.0,
response: json!({"choices": [{"text": "cached"}]}),
}
}
#[test]
fn cache_entries_round_trip_through_json() {
let entry = entry();
let encoded = RedisCache::<redis::Connection>::encode(&entry).unwrap();
assert_eq!(
RedisCache::<redis::Connection>::decode(encoded).unwrap(),
entry
);
}
#[test]
fn invalid_json_is_rejected() {
assert!(RedisCache::<redis::Connection>::decode(b"not json".to_vec()).is_err());
}
#[test]
fn ttl_seconds_rounds_up_and_keeps_expiration_positive() {
assert_eq!(
RedisCache::<redis::Connection>::ttl_seconds(Duration::ZERO),
1
);
assert_eq!(
RedisCache::<redis::Connection>::ttl_seconds(Duration::from_millis(1500)),
2
);
assert_eq!(
RedisCache::<redis::Connection>::ttl_seconds(Duration::from_secs(15)),
15
);
}
#[test]
fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() {
let value = entry();
let payload = RedisCache::<redis::Connection>::encode(&value).unwrap();
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SETEX")
.arg("litellm-cache:key")
.arg(600)
.arg(payload.clone()),
Ok("OK"),
),
MockCmd::new(redis::cmd("GET").arg("litellm-cache:key"), Ok(payload)),
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None);
cache
.set_cache("key", value.clone(), CacheKwargs::default())
.unwrap();
assert_eq!(
cache.get_cache("key", &CacheKwargs::default()).unwrap(),
Some(value)
);
cache.delete_cache("key").unwrap();
}
#[test]
fn flush_scans_and_deletes_only_cache_keys() {
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SCAN")
.cursor_arg(0)
.arg("MATCH")
.arg("litellm-cache:*"),
Ok(redis_test::redis_value!(["0", ["litellm-cache:key"]])),
),
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None);
cache.flush_cache().unwrap();
}
#[tokio::test]
async fn test_connection_runs_ping_off_executor() {
let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None);
assert_eq!(
cache.test_connection().await.unwrap().status,
litellm_cache::CacheConnectionStatus::Success
);
}
}

3
litellm-rust/crates/cache-redis/src/lib.rs Normal file
View file

@ -0,0 +1,3 @@
mod cache;
pub use cache::RedisCache;

6
litellm-rust/crates/cache-redis/tests/cache.rs Normal file
View file

@ -0,0 +1,6 @@
use litellm_cache_redis::RedisCache;
#[test]
fn constructor_rejects_invalid_urls() {
assert!(RedisCache::new("not a redis url", None).is_err());
}

11
litellm-rust/crates/config/Cargo.toml → litellm-rust/crates/cache/Cargo.toml vendored
View file

@ -1,16 +1,15 @@
[package]
name = "litellm-config"
name = "litellm-cache"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-core.workspace = true
pyo3 = { workspace = true, features = ["auto-initialize"], optional = true }
serde.workspace = true
serde_json.workspace = true
sha2.workspace = true
thiserror.workspace = true
[features]
default = []
python = ["dep:pyo3"]
[dev-dependencies]
rstest.workspace = true

98
litellm-rust/crates/cache/src/base_cache.rs vendored Normal file
View file

@ -0,0 +1,98 @@
use std::future::Future;
use std::pin::Pin;
use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::Error;
pub type CacheFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
#[derive(Clone, Debug, Default, PartialEq)]
pub struct CacheKwargs {
pub ttl: Option<Duration>,
pub extras: Map<String, Value>,
}
#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum CacheConnectionStatus {
Success,
Failed,
}
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
pub struct CacheConnectionResult {
pub status: CacheConnectionStatus,
pub message: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}
pub trait BaseCache: Send + Sync {
type Value: Clone + Send + Sync + 'static;
fn default_ttl(&self) -> Duration {
Duration::from_secs(60)
}
fn get_ttl(&self, kwargs: &CacheKwargs) -> Duration {
kwargs.ttl.unwrap_or_else(|| self.default_ttl())
}
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error>;
fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result<Option<Self::Value>, Error>;
fn async_set_cache<'a>(
&'a self,
key: &'a str,
value: Self::Value,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
Box::pin(async move { self.set_cache(key, value, kwargs) })
}
fn async_get_cache<'a>(
&'a self,
key: &'a str,
kwargs: &'a CacheKwargs,
) -> CacheFuture<'a, Option<Self::Value>> {
Box::pin(async move { self.get_cache(key, kwargs) })
}
fn async_set_cache_pipeline<'a>(
&'a self,
cache_list: Vec<(String, Self::Value)>,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
Box::pin(async move {
for (key, value) in cache_list {
self.set_cache(&key, value, kwargs.clone())?;
}
Ok(())
})
}
fn batch_cache_write<'a>(
&'a self,
key: &'a str,
value: Self::Value,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
self.async_set_cache(key, value, kwargs)
}
fn delete_cache(&self, key: &str) -> Result<(), Error>;
fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> {
Box::pin(async move { self.delete_cache(key) })
}
fn flush_cache(&self) -> Result<(), Error>;
fn disconnect(&self) -> CacheFuture<'_, ()>;
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult>;
}

166
litellm-rust/crates/cache/src/caching.rs vendored Normal file
View file

@ -0,0 +1,166 @@
use std::sync::Arc;
use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sha2::{Digest, Sha256};
use crate::{BaseCache, CacheKwargs, Error};
pub use crate::BaseCache as Cache;
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
pub enum CacheMode {
#[default]
#[serde(rename = "default_on")]
DefaultOn,
#[serde(rename = "default_off")]
DefaultOff,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct CacheKeyField {
pub name: String,
pub value: Option<String>,
pub api_parameter: bool,
pub internal_parameter: bool,
}
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
pub struct CacheKeyInput {
pub fields: Vec<CacheKeyField>,
pub preset: Option<String>,
pub namespace: Option<String>,
pub include_provider_parameters: bool,
}
#[derive(Default)]
pub struct CacheKeyContext {
pub model_group: Option<String>,
pub caching_groups: Vec<(Vec<String>, String)>,
pub file_checksum: Option<String>,
pub file_object_name: Option<String>,
pub metadata_file_name: Option<String>,
pub parameters_file_name: Option<String>,
}
impl CacheKeyContext {
pub fn apply(self, input: &mut CacheKeyInput) {
let group = self.model_group.as_ref().and_then(|model| {
self.caching_groups
.iter()
.find(|(models, _)| models.contains(model))
});
for field in &mut input.fields {
match field.name.as_str() {
"model" => {
field.value = group
.map(|(_, formatted)| formatted.clone())
.or_else(|| self.model_group.clone())
.or_else(|| field.value.take())
}
"file" => {
field.value = self
.file_checksum
.clone()
.or_else(|| self.file_object_name.clone())
.or_else(|| self.metadata_file_name.clone())
.or_else(|| self.parameters_file_name.clone())
}
_ => {}
}
}
}
}
pub fn get_cache_key(input: &CacheKeyInput) -> String {
cache_key(input)
}
pub fn cache_key(input: &CacheKeyInput) -> String {
if let Some(preset) = &input.preset {
return preset.clone();
}
let mut digest = Sha256::new();
for field in &input.fields {
if (field.api_parameter || (input.include_provider_parameters && !field.internal_parameter))
&& let Some(value) = &field.value
{
digest.update(field.name.as_bytes());
digest.update(b": ");
digest.update(value.as_bytes());
}
}
let hash = format!("{:x}", digest.finalize());
input
.namespace
.as_deref()
.filter(|namespace| !namespace.is_empty())
.map_or(hash.clone(), |namespace| format!("{namespace}:{hash}"))
}
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]
pub struct CacheControls {
pub supported_call_type: bool,
pub configured: bool,
pub native_backend: bool,
pub default_on: bool,
pub caching: Option<bool>,
pub no_cache: bool,
pub no_store: bool,
#[serde(default)]
pub use_cache: bool,
}
impl CacheControls {
pub fn reads(self) -> bool {
self.supported_call_type
&& self.configured
&& self.caching.unwrap_or(true)
&& !self.no_cache
&& (self.default_on || self.use_cache)
}
pub fn writes(self) -> bool {
self.supported_call_type
&& self.configured
&& !self.no_store
&& (self.default_on || self.use_cache)
}
}
pub fn should_use_cache(controls: CacheControls) -> bool {
controls.reads() || controls.writes()
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CacheEntry {
pub timestamp: f64,
pub response: Value,
}
impl CacheEntry {
pub fn fresh(&self, now: Duration, max_age: Option<Duration>) -> bool {
self.timestamp.is_finite()
&& max_age.is_none_or(|age| now.as_secs_f64() - self.timestamp <= age.as_secs_f64())
}
}
pub fn get_cache(
cache: &dyn BaseCache<Value = CacheEntry>,
key: &str,
kwargs: &CacheKwargs,
) -> Result<Option<CacheEntry>, Error> {
cache.get_cache(key, kwargs)
}
pub fn set_cache(
cache: &dyn BaseCache<Value = CacheEntry>,
key: &str,
entry: CacheEntry,
kwargs: CacheKwargs,
) -> Result<(), Error> {
cache.set_cache(key, entry, kwargs)
}
pub type CacheBackend = Arc<dyn BaseCache<Value = CacheEntry>>;

7
litellm-rust/crates/cache/src/error.rs vendored Normal file
View file

@ -0,0 +1,7 @@
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum Error {
#[error("cache is unavailable")]
Unavailable,
#[error("invalid cache entry")]
InvalidEntry,
}

12
litellm-rust/crates/cache/src/lib.rs vendored Normal file
View file

@ -0,0 +1,12 @@
mod base_cache;
mod caching;
mod error;
pub use base_cache::{
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheFuture, CacheKwargs,
};
pub use caching::{
Cache, CacheBackend, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput,
CacheMode, cache_key, get_cache, get_cache_key, set_cache, should_use_cache,
};
pub use error::Error;

139
litellm-rust/crates/cache/tests/caching.rs vendored Normal file
View file

@ -0,0 +1,139 @@
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheControls, CacheEntry, CacheFuture, CacheKeyContext,
CacheKeyField, CacheKeyInput, CacheKwargs, Error, cache_key, get_cache_key,
};
use sha2::{Digest, Sha256};
use std::time::Duration;
struct TestCache {
default_ttl: Duration,
}
impl BaseCache for TestCache {
type Value = CacheEntry;
fn default_ttl(&self) -> Duration {
self.default_ttl
}
fn set_cache(&self, _: &str, _: Self::Value, _: CacheKwargs) -> Result<(), Error> {
Ok(())
}
fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
Ok(None)
}
fn delete_cache(&self, _: &str) -> Result<(), Error> {
Ok(())
}
fn flush_cache(&self) -> Result<(), Error> {
Ok(())
}
fn disconnect(&self) -> CacheFuture<'_, ()> {
Box::pin(async { Ok(()) })
}
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
unreachable!()
}
}
#[test]
fn ttl_uses_default_and_allows_per_call_override() {
let cache = TestCache {
default_ttl: Duration::from_secs(60),
};
assert_eq!(
cache.get_ttl(&CacheKwargs::default()),
Duration::from_secs(60)
);
assert_eq!(
cache.get_ttl(&CacheKwargs {
ttl: Some(Duration::from_secs(5)),
..Default::default()
}),
Duration::from_secs(5)
);
}
#[test]
fn keys_match_python_order_groups_files_presets_and_namespaces() {
let mut input = CacheKeyInput {
fields: vec![
CacheKeyField {
name: "model".into(),
value: Some("deployment".into()),
api_parameter: true,
internal_parameter: false,
},
CacheKeyField {
name: "file".into(),
value: None,
api_parameter: true,
internal_parameter: false,
},
],
namespace: Some("team".into()),
..Default::default()
};
CacheKeyContext {
model_group: Some("group".into()),
caching_groups: vec![(vec!["group".into()], "['group']".into())],
file_checksum: Some("checksum".into()),
..Default::default()
}
.apply(&mut input);
assert_eq!(
cache_key(&input),
format!(
"team:{:x}",
Sha256::digest(b"model: ['group']file: checksum")
)
);
input.preset = Some("preset".into());
assert_eq!(get_cache_key(&input), "preset");
}
#[test]
fn cache_controls_honor_default_modes_and_directives() {
let enabled = CacheControls {
supported_call_type: true,
configured: true,
default_on: true,
..Default::default()
};
assert!(enabled.reads());
assert!(enabled.writes());
assert!(
!CacheControls {
default_on: false,
..enabled
}
.reads()
);
assert!(
CacheControls {
default_on: false,
use_cache: true,
..enabled
}
.reads()
);
assert!(
!CacheControls {
no_cache: true,
..enabled
}
.reads()
);
assert!(
!CacheControls {
no_store: true,
..enabled
}
.writes()
);
}

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