scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-30 01:52:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 70

chore: merge main into litellm_responses_precall_block_stream

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-09-23 14:51:00 +00:00
commit 3f6da36df1
67 changed files with 5890 additions and 2553 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.circleci/config.yml
View file

@ -3231,7 +3231,7 @@ workflows:
name: integration-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, database, providers, extensions, sdk, cost, browser]
suite: [management, accounting, database, providers, extensions, mcp, sdk, cost, browser]
filters:
branches:
only:

35
.circleci/scripts/run_integration.sh
View file

@ -111,6 +111,15 @@ upstream_pid=$!
if [ "$suite" = cost ]; then
export INTEGRATION_WORKERS=8
fi
if [ "$suite" = mcp ]; then
export INTEGRATION_WORKERS=4 INTEGRATION_COVERAGE=1
fi
coverage_data="$PWD/$results/coverage/data"
proxy_command=(.venv/bin/python -m integration._support.proxy)
if [ "${INTEGRATION_COVERAGE:-0}" = 1 ]; then
mkdir -p "$(dirname "$coverage_data")"
proxy_command=(.venv/bin/python -m coverage run --rcfile=tests/integration/mcp_coverage.toml -m integration._support.proxy)
fi
start_proxy() {
local port="$1"
local log_name="$2"
@ -133,8 +142,8 @@ start_proxy() {
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
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 STORE_MODEL_IN_DB=True "${cost_map_env[@]}" \
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
.venv/bin/python -m integration._support.proxy --config tests/integration/proxy_config.yaml \
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 COVERAGE_FILE="$coverage_data" \
"${proxy_command[@]}" --config tests/integration/proxy_config.yaml \
--host 127.0.0.1 --port "$port" --num_workers 1 --telemetry False \
--use_prisma_db_push --enforce_prisma_migration_check \
> "$results/$log_name" 2>&1 &
@ -146,7 +155,7 @@ proxy_pid="$launched_pid"
curl --noproxy '*' -sSf -X POST "$INTEGRATION_PROXY_URL/config/update" \
-H "Authorization: Bearer $LITELLM_MASTER_KEY" -H 'Content-Type: application/json' \
-d '{"router_settings": {"num_retries": 0}}' > "$results/seed-router-settings.json"
if [ "$suite" = management ]; then
if [ "$suite" = management ] || [ "$suite" = mcp ]; then
export INTEGRATION_PEER_URL=http://127.0.0.1:4001
start_proxy 4001 peer.log
peer_pid="$launched_pid"
@ -187,3 +196,23 @@ env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
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"
if [ "${INTEGRATION_COVERAGE:-0}" = 1 ]; then
for covered_pid in "$proxy_pid" "$peer_pid"; do
[ -n "$covered_pid" ] || continue
kill -TERM -- "-$covered_pid"
for _ in {1..300}; do
kill -0 "$covered_pid" 2>/dev/null || break
sleep 0.1
done
wait "$covered_pid" 2>/dev/null || true
done
proxy_pid=""
peer_pid=""
COVERAGE_FILE="$coverage_data" .venv/bin/python -m coverage combine --rcfile=tests/integration/mcp_coverage.toml
COVERAGE_FILE="$coverage_data" .venv/bin/python -m coverage report --rcfile=tests/integration/mcp_coverage.toml \
> "$results/coverage/coverage.txt"
COVERAGE_FILE="$coverage_data" .venv/bin/python -m coverage html --rcfile=tests/integration/mcp_coverage.toml \
-d "$results/coverage/html"
tail -n 1 "$results/coverage/coverage.txt"
fi

4
.circleci/scripts/verify_integration_browser.py
View file

@ -31,8 +31,8 @@ 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"]
(Path(__file__).resolve().parents[2] / "tests/e2e/ui/tests/integrationCritical/expected.json").read_text()
)
assert expected and result["stats"]["expected"] == len(expected)
assert all(result["stats"][name] == 0 for name in ("unexpected", "flaky", "skipped"))

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

@ -235,9 +235,7 @@ class Slice:
return True # a `-k` this parser cannot model is assumed to claim everything
if any(term.lower() in relative_path.lower() for term in self.excluded):
return False
return not self.required or any(
term.lower() in name.lower() for term in self.required for name in inner_names
)
return not self.required or any(term.lower() in name.lower() for term in self.required for name in inner_names)
def _strings(node: object) -> Iterable[str]:
@ -307,9 +305,7 @@ def _matchable_names(relative_path: str) -> frozenset[str]:
except (OSError, SyntaxError):
return frozenset({relative_path})
return frozenset({relative_path}) | frozenset(
node.name
for node in ast.walk(tree)
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))
node.name for node in ast.walk(tree) if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))
)
@ -331,9 +327,7 @@ def _deselected_everywhere(allowlist: Allowlist) -> tuple[Finding, ...]:
slices: Final = _slices()
named_by_workflow: Final = _workflow_named_tokens()
globbed: Final = tuple(
path
for path in _test_files()
if any(_token_covers(glob, path) for slice_ in slices for glob in slice_.globs)
path for path in _test_files() if any(_token_covers(glob, path) for slice_ in slices for glob in slice_.globs)
)
return tuple(
Finding(
@ -363,11 +357,7 @@ def _shard_children(root: str, repo_root: pathlib.Path = REPO_ROOT) -> tuple[str
child.relative_to(repo_root).as_posix()
for child in (repo_root / root).iterdir()
if not child.name.startswith(".")
and (
_holds_tests(child)
if child.is_dir()
else child.name.startswith("test_") and child.suffix == ".py"
)
and (_holds_tests(child) if child.is_dir() else child.name.startswith("test_") and child.suffix == ".py")
)
)
@ -499,13 +489,32 @@ def _check_shards() -> int:
return 0
def _integration_groups(runner: pathlib.Path) -> dict[str, tuple[str, ...]]:
module: Final = ast.parse(runner.read_text())
literal: Final = next(
node.value
for node in module.body
if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) and node.target.id == "GROUPS"
)
mapping: Final = literal.args[0] if isinstance(literal, ast.Call) else literal
return {group: tuple(folders) for group, folders in ast.literal_eval(mapping).items()}
def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozenset[str], tuple[Finding, ...]]:
manifest: Final = repo_root / "tests/integration/contracts.json"
if not manifest.exists():
runner: Final = repo_root / "tests/integration/run.py"
if not runner.exists():
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", {}))
groups: Final = _integration_groups(runner)
integration_root: Final = repo_root / "tests/integration"
paths: Final = frozenset(
str(path.relative_to(repo_root))
for folders in groups.values()
for folder in folders
for path in (integration_root / folder).glob("test_*.py")
)
browser_manifest: Final = repo_root / "tests/e2e/ui/tests/integrationCritical/expected.json"
browser_nodes: Final = json.loads(browser_manifest.read_text()) if browser_manifest.exists() else ()
browser_paths: Final = frozenset(node.split("::", 1)[0] for node in browser_nodes)
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", ())
@ -526,15 +535,14 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
)
required: Final = (frozenset({"browser"}) if browser_paths else frozenset()) | frozenset(
group
for group, folders in entries["groups"].items()
for group, folders in groups.items()
if any(any(path.startswith(f"tests/integration/{folder}/") for folder in folders) for path in paths)
)
ungrouped: Final = frozenset(
path
for path in paths
if sum(
any(path.startswith(f"tests/integration/{folder}/") for folder in folders)
for folders in entries["groups"].values()
any(path.startswith(f"tests/integration/{folder}/") for folder in folders) for folders in groups.values()
)
!= 1
)
@ -547,10 +555,6 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
Finding(path, "integration contract is also selected by GitHub Actions")
for path in paths
if any(_token_covers(token, path) for token in gha_tokens)
) + tuple(
Finding(path, "canonical integration test file is missing")
for path in paths
if not (repo_root / path).is_file()
)
browser_commands: Final = tuple(
scalar.value
@ -592,7 +596,7 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
) + tuple(Finding(path, "canonical node must have exactly one integration group") for path in sorted(ungrouped))
if not paths or not invoked or not scheduled:
return frozenset(), findings + (
Finding(str(manifest.relative_to(repo_root)), "dedicated CircleCI runner is missing"),
Finding(str(runner.relative_to(repo_root)), "dedicated CircleCI runner is missing"),
)
return paths | browser_paths, findings + group_findings + browser_findings + exclusion_findings

186
.github/workflows/create-release.yml vendored
View file

@ -1,186 +0,0 @@
name: Create Release
on:
workflow_dispatch:
inputs:
tag:
description: "Release tag (e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, 1.84.0-dev.2, 1.84.0.post1; legacy v1.83.10-stable still accepted)"
required: true
type: string
commit_hash:
description: "Full 40-char commit SHA to target"
required: true
type: string
permissions: {}
jobs:
release:
name: Create Release
runs-on: ubuntu-latest
permissions:
contents: write
steps:
- name: Validate inputs
env:
TAG: ${{ inputs.tag }}
COMMIT_HASH: ${{ inputs.commit_hash }}
run: |
if ! echo "${COMMIT_HASH}" | grep -qE '^[0-9a-f]{40}$'; then
echo "::error::commit_hash must be a full 40-character commit SHA"
exit 1
fi
if ! echo "${TAG}" | grep -qE '^v?[0-9]+\.[0-9]+\.[0-9]+'; then
echo "::error::tag must start with X.Y.Z (optional leading v), e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, or v1.83.10-stable"
exit 1
fi
- name: Create release
env:
TAG: ${{ inputs.tag }}
COMMIT_HASH: ${{ inputs.commit_hash }}
uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
with:
script: |
const tag = process.env.TAG;
const commitHash = process.env.COMMIT_HASH;
// Mark RC / dev / nightly / alpha / beta tags as GitHub pre-releases.
// Accept both PEP 440 (`.dev`) and SemVer (`-dev`) separators so tags
// like `1.84.0.dev2` and `1.84.0-dev.2` are both detected.
// PEP 440 post-releases (e.g. `1.84.0.post1`) and legacy `-stable[.patch.N]`
// are stable maintenance releases, not pre-releases.
const isPrerelease = /(?:rc|nightly|alpha|beta|[-.]dev)/i.test(tag);
// A stable release should only claim the repo "latest" badge when its
// version is >= the current latest. Otherwise a backport (e.g. 1.84.6)
// would steal "latest" from a newer line (e.g. 1.88.1).
const versionKey = (rawTag) => {
const m = String(rawTag).match(/^v?(\d+)\.(\d+)\.(\d+)/);
if (!m) return null;
const maintenance = String(rawTag).match(/(?:\.post|\.patch\.)(\d+)/i);
return [Number(m[1]), Number(m[2]), Number(m[3]), maintenance ? Number(maintenance[1]) : 0];
};
const isAtLeast = (a, b) => {
for (let i = 0; i < a.length; i++) {
if (a[i] !== b[i]) return a[i] > b[i];
}
return true;
};
const cosignSection = [
`## Verify Docker Image Signature`,
``,
`All LiteLLM Docker images are signed with [cosign](https://docs.sigstore.dev/cosign/overview/). Every release is signed with the same key introduced in [commit \`0112e53\`](https://github.com/BerriAI/litellm/commit/0112e53046018d726492c814b3644b7d376029d0).`,
``,
`**Verify using the pinned commit hash (recommended):**`,
``,
`A commit hash is cryptographically immutable, so this is the strongest way to ensure you are using the original signing key:`,
``,
'```bash',
`cosign verify \\`,
` --key https://raw.githubusercontent.com/BerriAI/litellm/0112e53046018d726492c814b3644b7d376029d0/cosign.pub \\`,
` ghcr.io/berriai/litellm:${tag}`,
'```',
``,
`**Verify using the release tag (convenience):**`,
``,
`Tags are protected in this repository and resolve to the same key. This option is easier to read but relies on tag protection rules:`,
``,
'```bash',
`cosign verify \\`,
` --key https://raw.githubusercontent.com/BerriAI/litellm/${tag}/cosign.pub \\`,
` ghcr.io/berriai/litellm:${tag}`,
'```',
``,
`Expected output:`,
``,
'```',
`The following checks were performed on each of these signatures:`,
` - The cosign claims were validated`,
` - The signatures were verified against the specified public key`,
'```',
``,
`---`,
``,
].join('\n');
try {
let makeLatest = "false";
const newVersion = versionKey(tag);
if (!isPrerelease && newVersion) {
let latestVersion = null;
try {
const latest = await github.rest.repos.getLatestRelease({
owner: context.repo.owner,
repo: context.repo.repo,
});
latestVersion = versionKey(latest.data.tag_name);
} catch (error) {
if (error.status !== 404) throw error;
}
makeLatest = (!latestVersion || isAtLeast(newVersion, latestVersion)) ? "true" : "false";
}
try {
await github.rest.git.createRef({
owner: context.repo.owner,
repo: context.repo.repo,
ref: `refs/tags/${tag}`,
sha: commitHash,
});
} catch (error) {
if (error.status !== 422) throw error;
const existing = await github.rest.git.getRef({
owner: context.repo.owner,
repo: context.repo.repo,
ref: `tags/${tag}`,
});
if (existing.data.object.sha !== commitHash) {
throw new Error(`Tag ${tag} already exists at ${existing.data.object.sha}, expected ${commitHash}`);
}
}
const response = await github.rest.repos.createRelease({
draft: true,
generate_release_notes: true,
name: tag,
owner: context.repo.owner,
prerelease: isPrerelease,
repo: context.repo.repo,
tag_name: tag,
});
const updatedBody = cosignSection + (response.data.body ?? '');
await github.rest.repos.updateRelease({
owner: context.repo.owner,
repo: context.repo.repo,
release_id: response.data.id,
tag_name: tag,
body: updatedBody,
draft: false,
});
if (!isPrerelease) {
await github.rest.repos.updateRelease({
owner: context.repo.owner,
repo: context.repo.repo,
release_id: response.data.id,
tag_name: tag,
make_latest: makeLatest,
});
}
} catch (error) {
core.setFailed(error.message);
}
create-branch:
name: Create Release Branch
needs: release
permissions:
contents: write
uses: ./.github/workflows/create-release-branch.yml
with:
tag: ${{ inputs.tag }}
commit_hash: ${{ inputs.commit_hash }}

1
backend/routes/allowlist.py
View file

@ -26,6 +26,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
"/v2/login",
"/v3/login",
"/logout",
"/session/logout",
"/token",
"/onboarding/",
"/audit",

8
litellm/_logging.py
View file

@ -9,7 +9,7 @@ import sys
from collections.abc import Iterator
from datetime import datetime
from logging import Formatter
from typing import Any, Final, TextIO
from typing import Final, TextIO
from urllib.parse import unquote
import litellm
@ -672,13 +672,13 @@ def _try_parse_json_message(message: str) -> dict[str, object] | None:
msg_stripped: Final = message.strip()
if not (msg_stripped.startswith("{") or msg_stripped.startswith("[")):
return None
parsed: Final = safe_json_loads(message, default=None)
parsed: Final[object] = safe_json_loads(message, default=None)
if parsed is None or not isinstance(parsed, dict):
return None
return parsed
def _try_parse_embedded_python_dict(message: str) -> dict[str, Any] | None:
def _try_parse_embedded_python_dict(message: str) -> dict[str, object] | None:
"""
Try to find and parse a Python dict repr (e.g. str(d) or repr(d)) embedded in
the message. Handles patterns like:
@ -702,7 +702,7 @@ def _try_parse_embedded_python_dict(message: str) -> dict[str, Any] | None:
if depth == 0:
substr = message[start : j + 1]
try:
result = ast.literal_eval(substr)
result: object = ast.literal_eval(substr)
if isinstance(result, dict) and len(result) > 0:
return result
except (ValueError, SyntaxError, TypeError):

1
litellm/integrations/otel/logger.py
View file

@ -561,6 +561,7 @@ class OpenTelemetryV2(CustomLogger):
time_to_first_chunk_seconds=call.time_to_first_chunk_seconds,
request_route=request_root_http_route(),
trace=call.trace,
session_id=call.session_id,
)
end_time_ns: Final = to_ns(end_time)
if carrier is not None and carrier.span is not None:

1
litellm/integrations/otel/mappers/genai.py
View file

@ -43,6 +43,7 @@ class GenAIMapper:
GenAI.OPERATION_NAME: lambda d: d.operation.value,
GenAI.PROVIDER_NAME: lambda d: d.provider or None,
GenAI.OUTPUT_TYPE: lambda d: d.output_type.value if d.output_type else None,
GenAI.CONVERSATION_ID: lambda d: d.session_id,
GenAI.REQUEST_MODEL: lambda d: d.request_model or None,
GenAI.REQUEST_TEMPERATURE: lambda d: d.request_params.temperature,
GenAI.REQUEST_TOP_P: lambda d: d.request_params.top_p,

36
litellm/integrations/otel/model/metadata.py
View file

@ -41,7 +41,7 @@ from dataclasses import dataclass, field
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, cast
from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, SESSION_ID_GENERATED_METADATA_KEY
from litellm.integrations.otel.model.semconv import resolve_operation
from litellm.integrations.otel.model.trace_controls import TraceControls, caller_trace_controls
from litellm.integrations.otel.model.utils import as_str, as_str_mapping, to_seconds
@ -226,6 +226,7 @@ class LLMCallEvent:
provisional_span_name: str
time_to_first_chunk_seconds: float | None
trace: TraceControls
session_id: str | None
@classmethod
def from_dict(cls, kwargs: Mapping[str, object]) -> LLMCallEvent:
@ -233,6 +234,7 @@ class LLMCallEvent:
payload: Final = cast("StandardLoggingPayload", raw_payload) if raw_payload else None
operation: Final = resolve_operation(as_str(kwargs.get("call_type")))
model: Final = as_str(kwargs.get("model")) or ""
trace: Final = caller_trace_controls(kwargs)
return cls(
call_id=_call_id(payload, kwargs),
payload=payload,
@ -242,10 +244,40 @@ class LLMCallEvent:
upstream_started=kwargs.get("api_call_start_time") is not None,
provisional_span_name=f"{operation.value} {model}".strip(),
time_to_first_chunk_seconds=time_to_first_chunk_seconds(kwargs),
trace=caller_trace_controls(kwargs),
trace=trace,
session_id=caller_session_id(kwargs, trace),
)
def caller_session_id(kwargs: Mapping[str, object], trace: TraceControls) -> str | None:
"""The conversation id the caller sent (``litellm_session_id``, else the
``session_id`` trace control); ``None`` when the request carried none.
``get_litellm_params`` back-fills ``litellm_session_id`` from ``metadata.trace_id``
(which the proxy stamps with the OTel trace id) and ``missing_session_id: generate``
mints one into the body; neither is a caller conversation, so both are ignored,
while a ``langfuse_session_id`` header still counts under the generate policy.
``StandardLoggingPayload.session_id`` is never read: the payload drops the
generated marker, so a replayed minted id would pass for a caller's."""
params: Final[Mapping[str, object]] = as_str_mapping(kwargs.get("litellm_params")) or MappingProxyType({})
bodies: Final = tuple(
metadata
for key in ("metadata", "litellm_metadata")
if (metadata := as_str_mapping(params.get(key))) is not None
)
from_body: Final = tuple(session for body in bodies if (session := as_str(body.get("session_id"))))
minted: Final = frozenset(
session
for body in bodies
if body.get(SESSION_ID_GENERATED_METADATA_KEY) and (session := as_str(body.get("session_id")))
)
if minted:
return next((session for session in (trace.session_id, *from_body) if session and session not in minted), None)
explicit: Final = as_str(params.get("litellm_session_id"))
echoes_trace_id: Final = explicit is not None and any(as_str(body.get("trace_id")) == explicit for body in bodies)
return (None if echoes_trace_id else explicit) or trace.session_id or None
def time_to_first_chunk_seconds(kwargs: Mapping[str, Any]) -> float | None:
"""Seconds from the upstream request being issued (``api_call_start_time``)
to the first streamed chunk (``completion_start_time``); ``None`` for

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

@ -407,6 +407,7 @@ class LLMCallSpanData:
call_type: str | None = None
request_route: str | None = None
trace: TraceControls = field(default_factory=TraceControls)
session_id: str | None = None
embedding_output: EmbeddingOutput | None = None
@classmethod
@ -417,6 +418,7 @@ class LLMCallSpanData:
time_to_first_chunk_seconds: float | None = None,
request_route: str | None = None,
trace: TraceControls | None = None,
session_id: str | None = None,
) -> LLMCallSpanData:
params: Final = cast(Mapping[str, object], payload.get("model_parameters") or {})
# The single parse of the request's metadata — the request-vs-provider
@ -463,6 +465,7 @@ class LLMCallSpanData:
call_type=call_type or None,
request_route=request_route or context.identity.request_route,
trace=trace or TraceControls(),
session_id=session_id or None,
embedding_output=embedding_output if capture_content else None,
)

2
litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py
View file

@ -67,7 +67,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig):
optional_params["responseFormat"] = self._normalize_response_format(value)
return optional_params
def _normalize_response_format(self, value: Any) -> Any:
def _normalize_response_format(self, value: Any) -> object:
"""Normalize response_format to TwelveLabs format.
TwelveLabs expects:

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

@ -358,13 +358,13 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig):
)
)
config_payload: Final[dict[str, Any]] = {
config_payload: Final[dict[str, object]] = {
"modules": modules if len(modules) > 1 else modules[0],
}
if stream_config:
config_payload["stream"] = stream_config
request_body: Final[dict[str, Any]] = {"config": config_payload}
request_body: Final[dict[str, object]] = {"config": config_payload}
if placeholder_values is not None:
request_body["placeholder_values"] = placeholder_values

316
litellm/model_prices_and_context_window_backup.json
View file

@ -3894,6 +3894,7 @@
"azure_ai/gpt-image-2": {
"cache_read_input_image_token_cost": 2e-06,
"cache_read_input_token_cost": 1.25e-06,
"deprecation_date": "2027-10-21",
"input_cost_per_image_token": 8e-06,
"input_cost_per_token": 5e-06,
"litellm_provider": "azure_ai",
@ -34171,6 +34172,7 @@
"supports_vision": true
},
"groq/llama-guard-3-8b": {
"deprecation_date": "2025-06-06",
"input_cost_per_token": 2e-07,
"litellm_provider": "groq",
"max_input_tokens": 8192,
@ -39858,21 +39860,21 @@
"supports_web_search": false
},
"openrouter/deepseek/deepseek-v4-pro": {
"input_cost_per_token": 8.92272e-07,
"input_cost_per_token": 8.8044e-07,
"input_cost_per_token_cache_hit": 4.4e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"max_tokens": 384000,
"mode": "chat",
"output_cost_per_token": 1.784544e-06,
"output_cost_per_token": 1.76088e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 7.4356e-08,
"cache_read_input_token_cost": 7.337e-08,
"supports_audio_input": false,
"supports_pdf_input": false,
"supports_vision": false,
@ -39884,10 +39886,10 @@
"cache_read_input_token_cost": 6e-09,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"max_tokens": 384000,
"max_output_tokens": 943718,
"max_tokens": 943718,
"mode": "chat",
"off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":1.5e-7,"output_cost_per_token":6e-7,"cache_read_input_token_cost":3e-9},
"off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":3e-7,"output_cost_per_token":0.0000012,"cache_read_input_token_cost":6e-9},
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -39915,7 +39917,7 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 4.4e-08,
"off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8},
"off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":0.00000132,"output_cost_per_token":0.00000396,"cache_read_input_token_cost":4.4e-8},
"supports_audio_input": false,
"supports_pdf_input": false,
"supports_vision": false,
@ -68170,14 +68172,14 @@
"supports_web_search": true
},
"openrouter/~deepseek/deepseek-flash-latest": {
"cache_read_input_token_cost": 3.6e-09,
"input_cost_per_token": 1.2e-07,
"cache_read_input_token_cost": 6e-09,
"input_cost_per_token": 3e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 943718,
"max_tokens": 943718,
"mode": "chat",
"output_cost_per_token": 4.8e-07,
"output_cost_per_token": 1.2e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -68190,14 +68192,15 @@
"supports_web_search": false
},
"openrouter/~deepseek/deepseek-pro-latest": {
"cache_read_input_token_cost": 1.2726e-08,
"input_cost_per_token": 3.9996e-07,
"cache_read_input_token_cost": 4.4e-08,
"input_cost_per_token": 1.32e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 393216,
"max_tokens": 393216,
"max_output_tokens": 384000,
"max_tokens": 384000,
"mode": "chat",
"output_cost_per_token": 1.19988e-06,
"off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8},
"output_cost_per_token": 3.96e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -68210,14 +68213,14 @@
"supports_web_search": false
},
"openrouter/~deepseek/deepseek-v4-flash-latest": {
"cache_read_input_token_cost": 8e-09,
"input_cost_per_token": 3e-08,
"cache_read_input_token_cost": 1.6e-08,
"input_cost_per_token": 4e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 1310720,
"max_output_tokens": 943718,
"max_tokens": 943718,
"mode": "chat",
"output_cost_per_token": 8e-07,
"output_cost_per_token": 6.4e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -68280,13 +68283,13 @@
},
"openrouter/~moonshotai/kimi-latest": {
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 1.4989e-06,
"input_cost_per_token": 3e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 943718,
"max_tokens": 943718,
"mode": "chat",
"output_cost_per_token": 1.0758e-05,
"output_cost_per_token": 1.5e-05,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -68442,14 +68445,14 @@
"supports_web_search": true
},
"openrouter/~z-ai/glm-flash-latest": {
"cache_read_input_token_cost": 1.5e-08,
"input_cost_per_token": 7.5e-08,
"cache_read_input_token_cost": 5e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1310720,
"max_output_tokens": 943718,
"max_tokens": 943718,
"mode": "chat",
"output_cost_per_token": 2.5e-07,
"output_cost_per_token": 5e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -68485,7 +68488,7 @@
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 8e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_input_tokens": 131072,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
@ -68505,7 +68508,7 @@
"cache_read_input_token_cost": 7.5e-07,
"input_cost_per_token": 3e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_input_tokens": 131072,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
@ -68525,7 +68528,47 @@
"cache_read_input_token_cost": 1.8e-07,
"input_cost_per_token": 7e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_input_tokens": 131072,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 1.4e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": false,
"supports_web_search": false
},
"openrouter/aion-labs/aion-3.5": {
"cache_read_input_token_cost": 7.5e-07,
"input_cost_per_token": 3e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 6e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": false,
"supports_web_search": false
},
"openrouter/aion-labs/aion-3.5-mini": {
"cache_read_input_token_cost": 1.8e-07,
"input_cost_per_token": 7e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
@ -71790,6 +71833,26 @@
"supports_vision": false,
"supports_web_search": false
},
"openrouter/upstage/solar-mini4": {
"cache_read_input_token_cost": 5e-09,
"input_cost_per_token": 5e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 524288,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 2e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": false,
"supports_web_search": false
},
"openrouter/writer/palmyra-x5": {
"input_cost_per_token": 6e-07,
"litellm_provider": "openrouter",
@ -72846,6 +72909,207 @@
"supports_tool_choice": true,
"supports_vision": true
},
"openrouter/anthropic/claude-opus-5.5:batch": {
"cache_creation_input_token_cost": 2.5e-06,
"cache_creation_input_token_cost_above_1hr": 4e-06,
"cache_read_input_token_cost": 1e-07,
"input_cost_per_token": 2e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1e-05,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"openrouter/cohere/command-a-plus": {
"cache_read_input_token_cost": 1.5e-07,
"input_cost_per_token": 3e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 192000,
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"output_cost_per_token": 1.5e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": false,
"supports_vision": true,
"supports_web_search": false
},
"openrouter/deepseek/deepseek-v4.1-flash:batch": {
"cache_read_input_token_cost": 3.36e-09,
"input_cost_per_token": 1.12e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 3.36e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
},
"openrouter/openai/gpt-6-luna-pro:batch": {
"cache_creation_input_token_cost": 6.25e-08,
"cache_creation_input_token_cost_above_272k_tokens": 1.25e-07,
"cache_read_input_token_cost": 5e-09,
"cache_read_input_token_cost_above_272k_tokens": 1e-08,
"input_cost_per_token": 5e-08,
"input_cost_per_token_above_272k_tokens": 1e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.5e-07,
"output_cost_per_token_above_272k_tokens": 3.75e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"openrouter/openai/gpt-6-luna:batch": {
"cache_creation_input_token_cost": 6.25e-08,
"cache_creation_input_token_cost_above_272k_tokens": 1.25e-07,
"cache_read_input_token_cost": 5e-09,
"cache_read_input_token_cost_above_272k_tokens": 1e-08,
"input_cost_per_token": 5e-08,
"input_cost_per_token_above_272k_tokens": 1e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.5e-07,
"output_cost_per_token_above_272k_tokens": 3.75e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"openrouter/openai/gpt-6-sol-pro:batch": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
"cache_read_input_token_cost": 1e-07,
"cache_read_input_token_cost_above_272k_tokens": 2e-07,
"input_cost_per_token": 1e-06,
"input_cost_per_token_above_272k_tokens": 2e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 5e-06,
"output_cost_per_token_above_272k_tokens": 7.5e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"openrouter/openai/gpt-6-sol:batch": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
"cache_read_input_token_cost": 1e-07,
"cache_read_input_token_cost_above_272k_tokens": 2e-07,
"input_cost_per_token": 1e-06,
"input_cost_per_token_above_272k_tokens": 2e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 5e-06,
"output_cost_per_token_above_272k_tokens": 7.5e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"openrouter/openai/gpt-oss-20b:batch": {
"input_cost_per_token": 2.4e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 131072,
"max_output_tokens": 117964,
"max_tokens": 117964,
"mode": "chat",
"output_cost_per_token": 1.12e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": false,
"supports_prompt_caching": false,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": false,
"supports_web_search": false
},
"openrouter/qwen/qwen3.8-omni-flash": {
"cache_read_input_token_cost": 1.6e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1000000,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 4.7e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_pdf_input": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
},
"vertex_ai/gemini-2.0-flash": {
"input_cost_per_audio_token": 1e-06,
"input_cost_per_audio_token_batches": 5e-07,

73
litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py
View file

@ -83,8 +83,14 @@ def _oauth_token_error(code: str, status: int = 400) -> JSONResponse:
def _user_id_from_session_cookie(request: Request) -> str | None:
"""Return user_id from the UI ``token`` cookie (HS256-signed with
``master_key``), or None if missing/invalid.
"""Return user_id from the UI ``token`` cookie, or None if missing/invalid."""
user_id, _ = _session_identity_from_cookie(request)
return user_id
def _session_identity_from_cookie(request: Request) -> tuple[str | None, str | None]:
"""Return ``(user_id, session_key)`` from the UI ``token`` cookie
(HS256-signed with ``master_key``), or ``(None, None)`` if missing/invalid.
The /token endpoint in this file ALSO issues master-key-signed JWTs
(type="byok_session") for MCP-client-side use. They must not be
@ -98,10 +104,10 @@ def _user_id_from_session_cookie(request: Request) -> str | None:
from litellm.proxy.proxy_server import master_key
if not master_key:
return None
return None, None
token: Final = request.cookies.get("token")
if not token:
return None
return None, None
try:
payload: Final = jwt.decode(
token,
@ -113,21 +119,68 @@ def _user_id_from_session_cookie(request: Request) -> str | None:
options={"require": ["exp"]},
)
except jwt.InvalidTokenError:
return None
return None, None
if payload.get("type") == "byok_session":
return None
return None, None
if payload.get("login_method") not in ("sso", "username_password"):
return None
return None, None
user_id: Final = payload.get("user_id")
return user_id if isinstance(user_id, str) and user_id else None
if not isinstance(user_id, str) or not user_id:
return None, None
session_key: Final = payload.get("key")
return user_id, session_key if isinstance(session_key, str) and session_key else None
async def _session_key_is_live(session_key: str | None) -> bool:
"""Whether the session key embedded in the UI cookie still resolves.
The cookie JWT stays signature-valid until ``exp``; the DB-backed session
key inside it is what ``POST /session/logout`` and password-change
revocation actually kill. Trusting the signature alone would let a
logged-out cookie keep authorizing BYOK credential writes, so re-resolve
the key here.
EXPERIMENTAL_UI_LOGIN blob tokens (non-``sk-``) have no DB row and are
unrevocable by construction (scoped out of revocation); they pass through
on their bounded 10-minute lifetime, as before.
"""
from litellm.proxy._types import hash_token
from litellm.proxy.auth.auth_checks import get_key_object
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if session_key is None:
# Older cookies predating the ``key`` claim: nothing to resolve.
return True
if not session_key.startswith("sk-"):
return True
if prisma_client is None:
return True
try:
await get_key_object(
hashed_token=hash_token(session_key),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
except Exception:
return False
return True
async def _byok_session_auth(request: Request) -> UserAPIKeyAuth:
"""Require the UI session cookie. Programmatic BYOK management uses
"""Require the UI session cookie, with the embedded session key
re-resolved against the DB so a revoked (logged-out) session cannot
authorize BYOK writes. Programmatic BYOK management uses
``POST /v1/mcp/server/{id}/user-credential`` instead."""
user_id: Final = _user_id_from_session_cookie(request)
user_id, session_key = _session_identity_from_cookie(request)
if not user_id:
raise HTTPException(status_code=401, detail="login_required")
if not await _session_key_is_live(session_key):
raise HTTPException(status_code=401, detail="login_required")
return UserAPIKeyAuth(api_key="byok_session_cookie", user_id=user_id)

5
litellm/proxy/_types.py
View file

@ -913,6 +913,7 @@ class LiteLLMRoutes(enum.Enum):
"/user/list", # org admins checked in endpoint; non-admins get 403
"/management/v1/users/bulk_delete", # proxy admins delete anyone, org admins only their orgs' users; others 403
"/user/password/change", # endpoint only ever writes the caller's own row
"/session/logout", # endpoint only ever revokes the caller's own session key
"/model/{model_id}/update",
"/prompt/list",
"/prompt/info",
@ -1948,6 +1949,10 @@ class ChangePasswordResponse(LiteLLMPydanticObjectBase):
message: str
class SessionLogoutResponse(LiteLLMPydanticObjectBase):
message: str
class DeleteUserRequest(LiteLLMPydanticObjectBase):
user_ids: list[str] # required

2
litellm/proxy/auth/login_utils.py
View file

@ -57,7 +57,7 @@ if TYPE_CHECKING:
from prisma import types as prisma_types
BREACH_RECHECK_INTERVAL: Final = timedelta(hours=24)
PASSWORD_RESET_ALLOWED_ROUTES: Final = ("/user/password/change",)
PASSWORD_RESET_ALLOWED_ROUTES: Final = ("/user/password/change", "/session/logout")
PASSWORD_SESSION_METADATA: Final = MappingProxyType({"login_method": "username_password"})

4
litellm/proxy/auth/route_checks.py
View file

@ -915,6 +915,10 @@ class RouteChecks:
if route == "/user/password/change":
return
# Self-service logout; the endpoint only revokes the caller's own session key.
if route == "/session/logout":
return
# Hard-block known write routes regardless of HTTP method (defensive
# — these are POSTs in practice, but pinning them here protects
# against future GET-shaped writes).

4
litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py
View file

@ -1,6 +1,6 @@
# litellm/proxy/guardrails/guardrail_hooks/pangea.py
import os
from typing import TYPE_CHECKING, Any, Final, cast
from typing import TYPE_CHECKING, Final
from fastapi import HTTPException
@ -230,7 +230,7 @@ class PangeaHandler(CustomGuardrail):
messages: Final = data.get("messages")
if messages is None:
return # No messages to check
input_messages = cast(list[dict[Any, Any]], messages)
input_messages = messages
else:
return

17
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -1583,6 +1583,23 @@ async def _update_single_user_helper(
response = inserted_user_row # pyright: ignore[reportAssignmentType] # insert_data returns a prisma row
if response is not None:
if "password" in non_default_values:
# An admin set this user's password, which implies the old one may be
# compromised; kill every existing UI session for the target. Revoke-all
# (no keep) — the caller is the admin, not the target, so the caller's
# own session is not among these.
from litellm.proxy.management_endpoints.session_endpoints import (
revoke_ui_session_keys,
)
target_user_id: Final = non_default_values.get("user_id")
if isinstance(target_user_id, str):
await revoke_ui_session_keys(
user_id=target_user_id,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
await _schedule_user_update_audit_log(
response=response,
existing_user_row=existing_user_row,

31
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -340,8 +340,21 @@ def _raise_on_strategy_router_write_violation(
)
def _stored_credential_name(existing_litellm_params: GenericLiteLLMParams | None) -> str | None:
if existing_litellm_params is None or existing_litellm_params.litellm_credential_name is None:
return None
return decrypt_value_helper(
value=existing_litellm_params.litellm_credential_name,
key="litellm_credential_name",
exception_type="debug",
return_original_value=True,
)
async def _raise_on_invalid_credential_name(
litellm_params: updateLiteLLMParams | None, prisma_client: PrismaClient
litellm_params: updateLiteLLMParams | None,
existing_litellm_params: GenericLiteLLMParams | None,
prisma_client: PrismaClient,
) -> None:
if litellm_params is None or "litellm_credential_name" not in litellm_params.model_fields_set:
return
@ -355,6 +368,8 @@ async def _raise_on_invalid_credential_name(
code=status.HTTP_400_BAD_REQUEST,
param="litellm_credential_name",
)
if credential_name == _stored_credential_name(existing_litellm_params):
return
if CredentialAccessor.find_credential(credential_name) is not None:
return
stored_credential: Final = await CredentialsRepository(WriterPinnedClient(prisma_client.db)).find_by_name(
@ -1192,7 +1207,7 @@ async def patch_model(
existing_litellm_params=db_model.litellm_params,
null_detaches=True,
)
await _raise_on_invalid_credential_name(patch_data.litellm_params, prisma_client)
await _raise_on_invalid_credential_name(patch_data.litellm_params, db_model.litellm_params, prisma_client)
ModelManagementAuthChecks.can_user_set_aws_session_tags(
litellm_params=patch_data.litellm_params,
@ -2012,18 +2027,8 @@ class ModelManagementAuthChecks:
return True
if litellm_params.litellm_credential_name is None and not null_detaches:
return True
existing_credential_name: Final = (
decrypt_value_helper(
value=existing_litellm_params.litellm_credential_name,
key="litellm_credential_name",
exception_type="debug",
return_original_value=True,
)
if existing_litellm_params is not None and existing_litellm_params.litellm_credential_name is not None
else None
)
requested_credential_name: Final = litellm_params.litellm_credential_name
if requested_credential_name == existing_credential_name:
if requested_credential_name == _stored_credential_name(existing_litellm_params):
return True
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
return True

10
litellm/proxy/management_endpoints/password_endpoints.py
View file

@ -26,6 +26,7 @@ from litellm.proxy._types import (
from litellm.proxy.auth.login_utils import PASSWORD_SESSION_METADATA
from litellm.proxy.auth.password_policy import validate_password_not_breached, validate_password_policy
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_endpoints.session_endpoints import revoke_ui_session_keys
from litellm.proxy.management_helpers.audit_logs import create_object_audit_log
from litellm.proxy.utils import hash_password, verify_password
from litellm.repositories.prisma_protocols import TableActions
@ -141,6 +142,15 @@ async def change_password(
}
await _user_table(prisma_client).update(where=find_user, data=password_update)
# The old password may have been compromised; revoke every other UI session
# so a holder of a stolen session token is cut off. The caller's own session
# is kept — they just proved they hold the current password.
await revoke_ui_session_keys(
user_id=user_id,
user_api_key_dict=user_api_key_dict,
keep_hashed_token=user_api_key_dict.token,
)
verbose_proxy_logger.info("Password changed via /user/password/change for user_id=%s", user_id)
await create_object_audit_log(
object_id=user_id,

175
litellm/proxy/management_endpoints/session_endpoints.py Normal file
View file

@ -0,0 +1,175 @@
"""
UI session revocation.
POST /session/logout — revoke the UI session key this request authenticated with.
revoke_ui_session_keys — revoke every UI session key a user holds (password writes).
Logging out of the dashboard was purely client-side (cookies cleared, redirect);
the DB-backed virtual key minted at login stayed valid until
LITELLM_UI_SESSION_DURATION elapsed, so a captured token kept working access
after logout, and changing a password did not invalidate existing sessions.
Deliberately NOT reusing /key/delete: its `can_modify_verification_token`
ownership checks can reject low-privilege roles, and a self-revoke endpoint
that takes no body cannot be aimed at other keys.
"""
from typing import TYPE_CHECKING, Annotated, Final, cast
from fastapi import APIRouter, Depends, HTTPException, Response
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.proxy._types import (
CommonProxyErrors,
HTTPExceptionErrorDetail,
LiteLLM_VerificationToken,
SessionLogoutResponse,
UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_checks import delete_cache_key_objects
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_endpoints.key_management_endpoints import (
_persist_deleted_verification_tokens,
)
from litellm.repositories.verification_token_repository import (
VerificationTokenRepository,
)
if TYPE_CHECKING:
from prisma import types as prisma_types
router: Final = APIRouter()
_TOKEN_LIST: Final = TypeAdapter(list[str])
def _error_detail(message: str) -> HTTPExceptionErrorDetail:
detail: Final[HTTPExceptionErrorDetail] = {"error": message}
return detail
async def revoke_ui_session_keys(
user_id: str,
user_api_key_dict: UserAPIKeyAuth,
*,
keep_hashed_token: str | None = None,
litellm_changed_by: str | None = None,
) -> int:
"""Revoke every UI session key belonging to ``user_id``, except
``keep_hashed_token`` (the caller's own session on a self-service password
change; the other password-write paths revoke all).
Best-effort: the password write this runs after has already committed, so a
revocation failure is logged loudly rather than failing the request — the
unrevoked keys still expire at LITELLM_UI_SESSION_DURATION.
Returns the number of sessions revoked.
"""
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
if prisma_client is None:
return 0
try:
where_user_sessions: Final[prisma_types.LiteLLM_VerificationTokenWhereInput] = {
"user_id": user_id,
"team_id": UI_SESSION_TOKEN_TEAM_ID,
}
rows: Final = cast( # cast-ok: find_many returns prisma rows shaped like the pydantic model
"tuple[LiteLLM_VerificationToken, ...]",
tuple(await VerificationTokenRepository(prisma_client).table.find_many(where=where_user_sessions)),
)
revoked_rows: Final = tuple(row for row in rows if row.token is not None and row.token != keep_hashed_token)
if not revoked_rows:
return 0
revoked_tokens: Final = _TOKEN_LIST.validate_python(tuple(row.token for row in revoked_rows))
await _persist_deleted_verification_tokens(
keys=revoked_rows,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
where_revoked: Final[prisma_types.LiteLLM_VerificationTokenWhereInput] = {"token": {"in": revoked_tokens}}
await VerificationTokenRepository(prisma_client).table.delete_many(where=where_revoked)
await delete_cache_key_objects(
hashed_tokens=revoked_tokens,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
verbose_proxy_logger.info(
"Revoked %s UI session key(s) for user_id=%s after password change",
len(revoked_tokens),
user_id,
)
return len(revoked_tokens)
except Exception: # noqa: BLE001 # the password write committed; revocation must not undo that
verbose_proxy_logger.exception(
"Failed to revoke UI session keys for user_id=%s; existing sessions remain valid until they expire",
user_id,
)
return 0
@router.post(
"/session/logout",
tags=("UI Session",),
)
async def session_logout(
response: Response,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> SessionLogoutResponse:
"""
Revoke the UI session key this request authenticated with.
Only accepts UI session keys (minted by dashboard login); any other
credential is refused, so this can never be used to delete arbitrary keys.
Revokes only the presented session, not the user's other sessions.
Idempotent: logging out an already-revoked session succeeds.
"""
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
if prisma_client is None:
raise HTTPException(
status_code=500,
detail=_error_detail(CommonProxyErrors.db_not_connected_error.value),
)
if user_api_key_dict.team_id != UI_SESSION_TOKEN_TEAM_ID:
raise HTTPException(
status_code=403,
detail=_error_detail("Only UI session tokens can be revoked through this endpoint."),
)
hashed_token: Final = user_api_key_dict.token
revoked = False
if hashed_token is not None:
where_token: Final[prisma_types.LiteLLM_VerificationTokenWhereUniqueInput] = {"token": hashed_token}
row: Final = await VerificationTokenRepository(prisma_client).table.find_unique(where=where_token)
# A missing row means the session is already revoked (or an
# EXPERIMENTAL_UI_LOGIN blob token); logout is idempotent either way.
if row is not None:
caller_row: Final = cast( # cast-ok: find_unique returns a prisma row shaped like the pydantic model
"LiteLLM_VerificationToken", row
)
await _persist_deleted_verification_tokens(
keys=(caller_row,),
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
)
await VerificationTokenRepository(prisma_client).table.delete_many(where=where_token)
revoked = True
await delete_cache_key_objects(
hashed_tokens=(hashed_token,),
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
# The server set this cookie at login (set_session_token_cookie); clear it
# here too so logout works even if the client-side clear is skipped.
response.delete_cookie("token")
return SessionLogoutResponse(
message="Session revoked." if revoked else "Session already revoked.",
)

17
litellm/proxy/proxy_server.py
View file

@ -630,6 +630,9 @@ from litellm.proxy.management_endpoints.prompt_caching_requests import (
from litellm.proxy.management_endpoints.router_settings_endpoints import (
router as router_settings_router,
)
from litellm.proxy.management_endpoints.session_endpoints import (
router as session_management_router,
)
from litellm.proxy.management_endpoints.tag_management_endpoints import (
router as tag_management_router,
)
@ -17012,6 +17015,19 @@ async def claim_onboarding_link(data: InvitationClaim, request: Request):
if user_obj and hasattr(user_obj, "__dict__"):
user_obj.__dict__.pop("password", None)
# The password just changed via an invitation/reset link; any UI session
# minted under the old password may be in hostile hands. Revoke them all —
# the caller holds only the short-lived onboarding JWT, and the fresh
# session key is minted below, after this sweep.
from litellm.proxy.management_endpoints.session_endpoints import (
revoke_ui_session_keys,
)
await revoke_ui_session_keys(
user_id=invite_obj.user_id,
user_api_key_dict=UserAPIKeyAuth(user_id=invite_obj.user_id),
)
try:
jwt_token: Final = await _generate_onboarding_ui_session_token(user_obj=user_obj)
except Exception as e:
@ -19431,6 +19447,7 @@ app.include_router(health_router)
app.include_router(key_management_router)
app.include_router(internal_user_router)
app.include_router(password_management_router)
app.include_router(session_management_router)
app.include_router(team_router)
app.include_router(ui_sso_router)
app.include_router(organization_router)

4
litellm/proxy/types_utils/utils.py
View file

@ -52,7 +52,7 @@ def get_instance_fn(value: str, config_file_path: str | None = None) -> Any:
module = importlib.import_module(module_name)
# Get the instance from the module
instance: Final = getattr(module, instance_name)
instance: Final[object] = getattr(module, instance_name)
return instance
except ImportError as e:
@ -167,7 +167,7 @@ def _load_instance_from_remote_storage(remote_url: str, config_file_path: str |
spec.loader.exec_module(module)
# Get the instance
instance: Final = getattr(module, instance_name)
instance: Final[object] = getattr(module, instance_name)
# Clean up the temporary file
try:

316
model_prices_and_context_window.json
View file

@ -3894,6 +3894,7 @@
"azure_ai/gpt-image-2": {
"cache_read_input_image_token_cost": 2e-06,
"cache_read_input_token_cost": 1.25e-06,
"deprecation_date": "2027-10-21",
"input_cost_per_image_token": 8e-06,
"input_cost_per_token": 5e-06,
"litellm_provider": "azure_ai",
@ -34171,6 +34172,7 @@
"supports_vision": true
},
"groq/llama-guard-3-8b": {
"deprecation_date": "2025-06-06",
"input_cost_per_token": 2e-07,
"litellm_provider": "groq",
"max_input_tokens": 8192,
@ -39858,21 +39860,21 @@
"supports_web_search": false
},
"openrouter/deepseek/deepseek-v4-pro": {
"input_cost_per_token": 8.92272e-07,
"input_cost_per_token": 8.8044e-07,
"input_cost_per_token_cache_hit": 4.4e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"max_tokens": 384000,
"mode": "chat",
"output_cost_per_token": 1.784544e-06,
"output_cost_per_token": 1.76088e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 7.4356e-08,
"cache_read_input_token_cost": 7.337e-08,
"supports_audio_input": false,
"supports_pdf_input": false,
"supports_vision": false,
@ -39884,10 +39886,10 @@
"cache_read_input_token_cost": 6e-09,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"max_tokens": 384000,
"max_output_tokens": 943718,
"max_tokens": 943718,
"mode": "chat",
"off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":1.5e-7,"output_cost_per_token":6e-7,"cache_read_input_token_cost":3e-9},
"off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":3e-7,"output_cost_per_token":0.0000012,"cache_read_input_token_cost":6e-9},
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -39915,7 +39917,7 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 4.4e-08,
"off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8},
"off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":0.00000132,"output_cost_per_token":0.00000396,"cache_read_input_token_cost":4.4e-8},
"supports_audio_input": false,
"supports_pdf_input": false,
"supports_vision": false,
@ -68170,14 +68172,14 @@
"supports_web_search": true
},
"openrouter/~deepseek/deepseek-flash-latest": {
"cache_read_input_token_cost": 3.6e-09,
"input_cost_per_token": 1.2e-07,
"cache_read_input_token_cost": 6e-09,
"input_cost_per_token": 3e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 943718,
"max_tokens": 943718,
"mode": "chat",
"output_cost_per_token": 4.8e-07,
"output_cost_per_token": 1.2e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -68190,14 +68192,15 @@
"supports_web_search": false
},
"openrouter/~deepseek/deepseek-pro-latest": {
"cache_read_input_token_cost": 1.2726e-08,
"input_cost_per_token": 3.9996e-07,
"cache_read_input_token_cost": 4.4e-08,
"input_cost_per_token": 1.32e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 393216,
"max_tokens": 393216,
"max_output_tokens": 384000,
"max_tokens": 384000,
"mode": "chat",
"output_cost_per_token": 1.19988e-06,
"off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8},
"output_cost_per_token": 3.96e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -68210,14 +68213,14 @@
"supports_web_search": false
},
"openrouter/~deepseek/deepseek-v4-flash-latest": {
"cache_read_input_token_cost": 8e-09,
"input_cost_per_token": 3e-08,
"cache_read_input_token_cost": 1.6e-08,
"input_cost_per_token": 4e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 1310720,
"max_output_tokens": 943718,
"max_tokens": 943718,
"mode": "chat",
"output_cost_per_token": 8e-07,
"output_cost_per_token": 6.4e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -68280,13 +68283,13 @@
},
"openrouter/~moonshotai/kimi-latest": {
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 1.4989e-06,
"input_cost_per_token": 3e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 943718,
"max_tokens": 943718,
"mode": "chat",
"output_cost_per_token": 1.0758e-05,
"output_cost_per_token": 1.5e-05,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -68442,14 +68445,14 @@
"supports_web_search": true
},
"openrouter/~z-ai/glm-flash-latest": {
"cache_read_input_token_cost": 1.5e-08,
"input_cost_per_token": 7.5e-08,
"cache_read_input_token_cost": 5e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1310720,
"max_output_tokens": 943718,
"max_tokens": 943718,
"mode": "chat",
"output_cost_per_token": 2.5e-07,
"output_cost_per_token": 5e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -68485,7 +68488,7 @@
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 8e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_input_tokens": 131072,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
@ -68505,7 +68508,7 @@
"cache_read_input_token_cost": 7.5e-07,
"input_cost_per_token": 3e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_input_tokens": 131072,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
@ -68525,7 +68528,47 @@
"cache_read_input_token_cost": 1.8e-07,
"input_cost_per_token": 7e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_input_tokens": 131072,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 1.4e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": false,
"supports_web_search": false
},
"openrouter/aion-labs/aion-3.5": {
"cache_read_input_token_cost": 7.5e-07,
"input_cost_per_token": 3e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 6e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": false,
"supports_web_search": false
},
"openrouter/aion-labs/aion-3.5-mini": {
"cache_read_input_token_cost": 1.8e-07,
"input_cost_per_token": 7e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
@ -71790,6 +71833,26 @@
"supports_vision": false,
"supports_web_search": false
},
"openrouter/upstage/solar-mini4": {
"cache_read_input_token_cost": 5e-09,
"input_cost_per_token": 5e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 524288,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 2e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": false,
"supports_web_search": false
},
"openrouter/writer/palmyra-x5": {
"input_cost_per_token": 6e-07,
"litellm_provider": "openrouter",
@ -72846,6 +72909,207 @@
"supports_tool_choice": true,
"supports_vision": true
},
"openrouter/anthropic/claude-opus-5.5:batch": {
"cache_creation_input_token_cost": 2.5e-06,
"cache_creation_input_token_cost_above_1hr": 4e-06,
"cache_read_input_token_cost": 1e-07,
"input_cost_per_token": 2e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1e-05,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"openrouter/cohere/command-a-plus": {
"cache_read_input_token_cost": 1.5e-07,
"input_cost_per_token": 3e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 192000,
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"output_cost_per_token": 1.5e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": false,
"supports_vision": true,
"supports_web_search": false
},
"openrouter/deepseek/deepseek-v4.1-flash:batch": {
"cache_read_input_token_cost": 3.36e-09,
"input_cost_per_token": 1.12e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 3.36e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
},
"openrouter/openai/gpt-6-luna-pro:batch": {
"cache_creation_input_token_cost": 6.25e-08,
"cache_creation_input_token_cost_above_272k_tokens": 1.25e-07,
"cache_read_input_token_cost": 5e-09,
"cache_read_input_token_cost_above_272k_tokens": 1e-08,
"input_cost_per_token": 5e-08,
"input_cost_per_token_above_272k_tokens": 1e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.5e-07,
"output_cost_per_token_above_272k_tokens": 3.75e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"openrouter/openai/gpt-6-luna:batch": {
"cache_creation_input_token_cost": 6.25e-08,
"cache_creation_input_token_cost_above_272k_tokens": 1.25e-07,
"cache_read_input_token_cost": 5e-09,
"cache_read_input_token_cost_above_272k_tokens": 1e-08,
"input_cost_per_token": 5e-08,
"input_cost_per_token_above_272k_tokens": 1e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.5e-07,
"output_cost_per_token_above_272k_tokens": 3.75e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"openrouter/openai/gpt-6-sol-pro:batch": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
"cache_read_input_token_cost": 1e-07,
"cache_read_input_token_cost_above_272k_tokens": 2e-07,
"input_cost_per_token": 1e-06,
"input_cost_per_token_above_272k_tokens": 2e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 5e-06,
"output_cost_per_token_above_272k_tokens": 7.5e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"openrouter/openai/gpt-6-sol:batch": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
"cache_read_input_token_cost": 1e-07,
"cache_read_input_token_cost_above_272k_tokens": 2e-07,
"input_cost_per_token": 1e-06,
"input_cost_per_token_above_272k_tokens": 2e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 5e-06,
"output_cost_per_token_above_272k_tokens": 7.5e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"openrouter/openai/gpt-oss-20b:batch": {
"input_cost_per_token": 2.4e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 131072,
"max_output_tokens": 117964,
"max_tokens": 117964,
"mode": "chat",
"output_cost_per_token": 1.12e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_pdf_input": false,
"supports_prompt_caching": false,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": false,
"supports_web_search": false
},
"openrouter/qwen/qwen3.8-omni-flash": {
"cache_read_input_token_cost": 1.6e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1000000,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 4.7e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_pdf_input": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
},
"vertex_ai/gemini-2.0-flash": {
"input_cost_per_audio_token": 1e-06,
"input_cost_per_audio_token_batches": 5e-07,

3
tests/e2e/ui/tests/integrationCritical/expected.json Normal file
View file

@ -0,0 +1,3 @@
[
"tests/e2e/ui/tests/integrationCritical/projectDetachment.spec.ts::project creation and explicit detachment preserve saved scope and restore serving"
]

6
tests/integration/AGENTS.md
View file

@ -21,5 +21,7 @@ function in a full stack is not
## Where it goes
By the domain a user would name: `pricing`, `spend`, `routing`. Add the node and its `covers` ids to
`contracts.json` or collection fails. Needs no proxy, DB or Redis: `tests/unit`
By the domain a user would name: `pricing`, `spend`, `routing`, `mcp`. A file only needs to live in a
directory that a `GROUPS` entry in `run.py` selects; there is no manifest and no `covers` marker on new
tests. A product bug the test exposes is `pytest.skip("BUG: <symptom>")` at the top of the body, not a
fix in the test and not a deletion. Needs no proxy, DB or Redis: `tests/unit`

12
tests/integration/README.md
View file

@ -2,9 +2,9 @@
These tests exercise a running gateway, PostgreSQL and Redis with an owned local upstream. CircleCI owns this suite. Tests are grouped by behavior, with no automatic test retries or fallback to paid provider calls
The `cost` group is driven by `cost_tracking_cases.json`, which contains the cost map, literal requests, literal provider responses and expected accounting values. Each case has a name, contract ID, cost-map model, optional deployment overrides, request body, tagged response and exact or recount expectations. Request bodies use `$MODEL` for the registered proxy model, while responses use `$REQUEST_ID` for the per-run scenario ID. To add a case, add a cost-map entry when the model is new, add the request body and exact provider response data, add hand-computed expected values and register the node ID in `contracts.json`. The upstream serves each stored response for any path under `/<scenario_id>`, while the test-owned cost map is served over loopback through `LITELLM_MODEL_COST_MAP_URL`
The `cost` group is driven by `cost_tracking_cases.json`, which contains the cost map, literal requests, literal provider responses and expected accounting values. Each case has a name, contract ID, cost-map model, optional deployment overrides, request body, tagged response and exact or recount expectations. Request bodies use `$MODEL` for the registered proxy model, while responses use `$REQUEST_ID` for the per-run scenario ID. To add a case, add a cost-map entry when the model is new, add the request body and exact provider response data, and add hand-computed expected values. The upstream serves each stored response for any path under `/<scenario_id>`, while the test-owned cost map is served over loopback through `LITELLM_MODEL_COST_MAP_URL`
Use `tests/integration/run.py management`, `accounting`, `database`, `providers`, `extensions`, `sdk` or `cost` to run a selected group. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate
Use `tests/integration/run.py management`, `accounting`, `database`, `providers`, `extensions`, `mcp`, `sdk` or `cost` to run a selected group. The group to directory mapping is the `GROUPS` literal at the top of `run.py`; a new directory needs a `GROUPS` entry and an `OWNED_DIRECTORIES` entry in `_support/manifest.py`. Set `INTEGRATION_WORKERS` above 1 to run a group under pytest-xdist; the `mcp` job does this in CI, so MCP tests must own their resources per scenario. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate
Management also requires `INTEGRATION_PEER_URL`, `REDIS_HOST` and `REDIS_PORT`. CircleCI starts two directly addressed proxy processes sharing only that job's stores. The test-only CLI wrapper supplies enterprise route entitlement, following the existing behavior suite's convention. It does not qualify license validation; run it with one worker and no reload
@ -14,7 +14,7 @@ Reuse the existing canned provider handlers through `_support/upstream.py`. It r
The CircleCI workflow starts its own database and Redis, restricts test-phase egress to its owned services and writes JUnit plus an executed-node manifest. Missing setup, failed cleanup or a selected test with neither a passed call nor a skip fail qualification. Skipped nodes are listed under `skipped` in `execution.json`, so the skip reasons double as the open bug list. Existing GitHub Actions jobs do not own these tests
Define integration contract IDs and their canonical test nodes in `contracts.json`. Every node must declare the same IDs with `covers`. The runner checks exact collected and passed-or-skipped selections against that mapping. These IDs belong to this CircleCI suite and must not be added to the separate E2E coverage registry. A manifest declaration alone does not mean a test passed
There is no per-node manifest. The runner fails only when pytest fails, when collection errors, or when a selected file collects zero tests. Older tests still carry `@pytest.mark.covers(...)` decorators; the marker stays registered so they collect, but the IDs are not checked against anything and new tests should not use it. The GitHub Actions coverage census reads the `GROUPS` literal in `run.py` and treats every `tests/integration/<directory>/test_*.py` file in a scheduled group as owned by CircleCI
Provider sentinels currently use the controlled server, not live recordings. The provider shard also runs the existing strict replay controls for changed requests, exhausted interactions, leftover interactions and no provider connection. Future recorded scenarios must use that replay-only implementation; missing recordings cannot fall back to a real provider. The observation endpoint is destructive and the current selection runs serially against one owned upstream
@ -30,6 +30,8 @@ Streaming checks send real HTTP transfer chunks, including one-byte partitions,
The sdk shard exercises the SDK's own HTTP clients against local protocol peers with no gateway in the path, so a case here fails only when the client library or its wire behavior changes. The HTTP/2 case runs a hypercorn TLS peer offering h2 and http/1.1 over ALPN, drives the sync and async httpx handlers at it with `LITELLM_HTTP2` off and on, and asserts the version both the client and the peer observed on the wire. Put a test here only when it needs no proxy, database or Redis; a case that reaches the gateway belongs in one of the other shards
The extensions shard reuses the existing MCP arithmetic functions with a real SDK server, and uses the built-in generic callback and guardrail transports. It checks actual tool calls after saved edits, discovery preservation, malformed/error responses, callback correlation and credential exclusion, guardrail rewriting and denial, retained OpenAI consumers, persisted toolsets and A2A wire versions
The extensions shard uses the built-in generic callback and guardrail transports. It checks callback correlation and credential exclusion, guardrail rewriting and denial, retained OpenAI consumers and A2A wire versions
Browser contracts live in `tests/e2e/ui/tests/integrationCritical` and run only through `tests/e2e/ui/integration.config.ts`. The CircleCI browser shard builds the checked-out dashboard, starts the owned proxy with that build, and verifies one exact browser result without retries or skips. The default Playwright selection excludes this directory. The focused project flow asserts the submitted create and clear values, fresh SQL state and actual blocked/restored serving while preserving model restrictions
The mcp shard runs the MCP gateway against SDK peers owned by each test (`_support/mcp.py`): streamable HTTP, SSE and stdio peers, an OpenAPI-spec app, and an OAuth 2.1 authorization-server double. Every peer records the requests it receives so a test can assert what reached the peer, not only what the proxy answered. The shard runs with `INTEGRATION_WORKERS` set and with `INTEGRATION_COVERAGE=1`, which starts the proxy under `coverage run --parallel-mode` limited to the MCP modules and stores `coverage.txt` plus an HTML report with the job artifacts. A test that fails because the product is wrong is skipped with `pytest.skip("BUG: <symptom>")` so the skip list in `execution.json` is the open MCP bug list
Browser contracts live in `tests/e2e/ui/tests/integrationCritical` and run only through `tests/e2e/ui/integration.config.ts`. The expected browser results are listed in `expected.json` in that directory and checked by `.circleci/scripts/verify_integration_browser.py`. The CircleCI browser shard builds the checked-out dashboard, starts the owned proxy with that build, and verifies one exact browser result without retries or skips. The default Playwright selection excludes this directory. The focused project flow asserts the submitted create and clear values, fresh SQL state and actual blocked/restored serving while preserving model restrictions

8
tests/integration/_support/asgi.py
View file

@ -4,8 +4,8 @@ import queue
import socket
import threading
import time
from collections.abc import Callable, Iterator
from concurrent.futures import Future
from collections.abc import Iterator
from contextlib import contextmanager
from typing import Final
@ -14,7 +14,7 @@ from starlette.types import ASGIApp
@contextmanager
def asgi_server(app: ASGIApp) -> Iterator[str]:
def asgi_server(app: ASGIApp, *, before_stop: Callable[[], None] | None = None) -> Iterator[str]:
with socket.socket() as listener:
listener.bind(("127.0.0.1", 0))
port: Final = listener.getsockname()[1]
@ -47,7 +47,7 @@ def asgi_server(app: ASGIApp) -> Iterator[str]:
class Capture(logging.Handler):
def emit(self, record: logging.LogRecord) -> None:
if record.thread == worker.ident and record.levelno >= logging.ERROR:
errors.put(record.getMessage())
errors.put(self.format(record))
handler: Final = Capture()
logger: Final = logging.getLogger("uvicorn.error")
@ -60,6 +60,8 @@ def asgi_server(app: ASGIApp) -> Iterator[str]:
time.sleep(0.01)
yield f"http://127.0.0.1:{port}"
finally:
if before_stop is not None:
before_stop()
server.should_exit = True
worker.join(timeout=8)
forced: Final = worker.is_alive()

13
tests/integration/_support/manifest.py
View file

@ -1,10 +1,5 @@
import json
from pathlib import Path
from typing import Final
from pydantic import TypeAdapter
MAPPING: Final = TypeAdapter(dict[str, tuple[str, ...]])
OWNED_DIRECTORIES: Final = frozenset(
{
"management",
@ -23,11 +18,3 @@ OWNED_DIRECTORIES: Final = frozenset(
"cost_calculation",
}
)
def contracts() -> dict[str, tuple[str, ...]]:
document: Final = json.loads((Path(__file__).resolve().parents[1] / "contracts.json").read_bytes())
result: Final = MAPPING.validate_python(document["tests"])
if not result or any(not values or any(not value.strip() for value in values) for values in result.values()):
raise ValueError("Integration manifest must contain nodes with contract IDs")
return result

569
tests/integration/_support/mcp.py
View file

@ -1,33 +1,70 @@
import asyncio
import json
import os
import queue
from collections.abc import Iterator
import sys
import time
from collections.abc import Callable, Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Final
from dataclasses import dataclass, field
from pathlib import Path
from typing import Final, Literal
import httpx
from integration._support.asgi import asgi_server
from integration._support.client import Gateway, Scenario
from integration._support.database import read_rows
from mcp.server.mcpserver import MCPServer
from integration._support.wire import Reply, Request, wire_server
from mcp import ClientSession
from mcp.client.sse import sse_client
from mcp.client.streamable_http import streamable_http_client
from mcp.server.mcpserver import Context, MCPServer
from mcp.server.transport_security import TransportSecuritySettings
from mcp.types import SamplingMessage, TextContent
from mcp_tests.mcp_e2e_upstream_server import add, multiply
from starlette.requests import Request
from pydantic import BaseModel
from sse_starlette.sse import AppStatus
from starlette.requests import Request as StarletteRequest
from starlette.responses import Response
from starlette.types import Message, Receive, Scope, Send
Transport = Literal["http", "sse", "stdio"]
STDIO_PEER: Final = Path(__file__).with_name("mcp_stdio_peer.py")
@dataclass(frozen=True, slots=True)
class McpPeer:
url: str
calls: queue.Queue[dict[str, object]]
transport: Transport = "http"
command: str | None = None
args: tuple[str, ...] = ()
record: Path | None = None
spec_path: Path | None = None
consumed: list[int] = field(default_factory=lambda: [0])
def drain(self) -> tuple[dict[str, object], ...]:
if self.record is not None:
lines: Final = self.record.read_text().splitlines() if self.record.exists() else []
fresh: Final = tuple(json.loads(line) for line in lines[self.consumed[0] :])
self.consumed[0] = len(lines)
return fresh
return tuple(self.calls.get_nowait() for _ in range(self.calls.qsize()))
def registration(self) -> dict[str, object]:
if self.transport == "stdio":
return {"transport": "stdio", "command": self.command, "args": list(self.args)}
if self.spec_path is not None:
return {"transport": "http", "url": self.url, "spec_path": str(self.spec_path)}
return {"transport": self.transport, "url": self.url}
@contextmanager
def mcp_peer() -> Iterator[McpPeer]:
service: Final = MCPServer("integration-math")
class Confirmation(BaseModel):
confirmed: bool
def math_service(name: str = "integration-math", *, rich: bool = False) -> MCPServer:
service: Final = MCPServer(name)
service.add_tool(add)
service.add_tool(multiply)
@ -35,22 +72,61 @@ def mcp_peer() -> Iterator[McpPeer]:
def fail() -> str:
raise ValueError("synthetic tool failure")
app: Final = service.streamable_http_app(
stateless_http=True,
json_response=True,
transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
)
observed: Final[queue.Queue[dict[str, object]]] = queue.Queue()
if not rich:
return service
@service.tool()
async def slow(seconds: float) -> str:
await asyncio.sleep(seconds)
return "slept"
@service.tool()
async def progress(steps: int, ctx: Context) -> str:
for step in range(steps):
await ctx.report_progress(step + 1, steps, f"step {step + 1}")
return f"{steps} steps"
@service.tool()
async def sample(prompt: str, ctx: Context) -> str:
result: Final = await ctx.session.create_message(
messages=[SamplingMessage(role="user", content=TextContent(type="text", text=prompt))],
max_tokens=32,
)
return "sampled:" + (result.content.text if isinstance(result.content, TextContent) else "")
@service.tool()
async def elicit(question: str, ctx: Context) -> str:
result: Final = await ctx.elicit(message=question, schema=Confirmation)
return f"elicited:{result.action}"
@service.prompt()
def greeting(name: str) -> str:
return f"Hello, {name}"
@service.resource("status://ready")
def status() -> str:
return "ready"
@service.resource("greeting://{name}")
def greeting_resource(name: str) -> str:
return f"Hello, {name}"
return service
def _capturing(app: Callable[[Scope, Receive, Send], object], observed: queue.Queue[dict[str, object]]):
async def capture(scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] != "http":
await app(scope, receive, send)
return
body: Final = await Request(scope, receive).body()
if scope["method"] == "GET" and scope["path"].endswith("/mcp"):
await Response(status_code=405, headers={"Allow": "POST, DELETE"})(scope, receive, send)
return
body: Final = await StarletteRequest(scope, receive).body()
assert len(body) <= 65536
if body:
observed.put({"body": json.loads(body), "headers": dict(scope["headers"])})
message: Final[Message] = {"type": "http.request", "body": body, "more_body": False}
observed.put({"body": json.loads(body), "headers": dict(scope["headers"]), "path": scope["path"]})
message: Final = {"type": "http.request", "body": body, "more_body": False}
pending: Final = iter((message,))
async def replay() -> Message:
@ -61,35 +137,285 @@ def mcp_peer() -> Iterator[McpPeer]:
await app(scope, replay, send)
with asgi_server(capture) as url:
yield McpPeer(url + "/mcp", observed)
return capture
def _drain_sse_streams() -> None:
AppStatus.should_exit = True
def _draining_sse_watcher(app: Callable[[Scope, Receive, Send], object]):
"""sse_starlette parks a per-loop watcher that only stops once AppStatus.should_exit flips."""
async def lifespan(scope: Scope, receive: Receive, send: Send) -> None:
while True:
message: Final = await receive()
if message["type"] == "lifespan.startup":
AppStatus.should_exit = False
await send({"type": "lifespan.startup.complete"})
elif message["type"] == "lifespan.shutdown":
_drain_sse_streams()
watchers: Final = tuple(
task for task in asyncio.all_tasks() if "_shutdown_watcher" in repr(task.get_coro())
)
await asyncio.gather(*watchers)
await send({"type": "lifespan.shutdown.complete"})
return
async def wrapped(scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] == "lifespan":
await lifespan(scope, receive, send)
return
starts: Final = [0]
async def send_once(message: Message) -> None:
if message["type"] == "http.response.start":
starts[0] += 1
if starts[0] == 2:
await send({"type": "http.response.body", "body": b"", "more_body": False})
if starts[0] > 1:
return
await send(message)
await app(scope, receive, send_once)
return wrapped
@contextmanager
def mcp_peer(transport: Literal["http", "sse"] = "http", *, rich: bool = False) -> Iterator[McpPeer]:
service: Final = math_service(rich=rich)
security: Final = TransportSecuritySettings(enable_dns_rebinding_protection=False)
app: Final = (
_draining_sse_watcher(service.sse_app(transport_security=security))
if transport == "sse"
else service.streamable_http_app(stateless_http=True, json_response=True, transport_security=security)
)
observed: Final[queue.Queue[dict[str, object]]] = queue.Queue()
with asgi_server(_capturing(app, observed), before_stop=_drain_sse_streams if transport == "sse" else None) as url:
yield McpPeer(url + ("/sse" if transport == "sse" else "/mcp"), observed, transport)
@contextmanager
def stdio_peer(directory: Path, *, rich: bool = False) -> Iterator[McpPeer]:
record: Final = directory / f"stdio-{os.getpid()}-{time.monotonic_ns()}.jsonl"
yield McpPeer(
"",
queue.Queue(),
"stdio",
sys.executable,
(str(STDIO_PEER), str(record), "rich" if rich else "plain"),
record,
)
JsonRpc = Mapping[str, object]
@dataclass(frozen=True, slots=True)
class ScriptedTool:
name: str
respond: Callable[[JsonRpc], Reply | JsonRpc]
def jsonrpc_reply(identity: object, result: JsonRpc) -> Reply:
return Reply(body=json.dumps({"jsonrpc": "2.0", "id": identity, "result": result}).encode())
def jsonrpc_error(identity: object, code: int, message: str) -> Reply:
return Reply(
body=json.dumps({"jsonrpc": "2.0", "id": identity, "error": {"code": code, "message": message}}).encode()
)
@contextmanager
def scripted_peer(*tools: ScriptedTool) -> Iterator[McpPeer]:
"""Raw JSON-RPC peer for shapes the SDK server cannot produce: half-written bodies, stalls, wire errors."""
observed: Final[queue.Queue[dict[str, object]]] = queue.Queue()
by_name: Final = {tool.name: tool for tool in tools}
def provider(request: Request) -> Reply:
if request.method != "POST":
return Reply(status=405)
body: Final = json.loads(request.body)
observed.put({"body": body, "headers": dict(request.headers), "path": request.target})
if "id" not in body:
return Reply(status=202)
identity: Final = body["id"]
method: Final = body["method"]
if method == "initialize":
return jsonrpc_reply(
identity,
{
"protocolVersion": body["params"]["protocolVersion"],
"capabilities": {"tools": {}},
"serverInfo": {"name": "integration-scripted-peer", "version": "1"},
},
)
if method == "tools/list":
return jsonrpc_reply(
identity, {"tools": [{"name": name, "inputSchema": {"type": "object"}} for name in by_name]}
)
if method != "tools/call":
return jsonrpc_error(identity, -32601, f"unsupported method {method}")
tool: Final = by_name.get(body["params"]["name"])
if tool is None:
return jsonrpc_error(identity, -32602, "unknown tool")
produced: Final = tool.respond(body["params"])
return produced if isinstance(produced, Reply) else jsonrpc_reply(identity, produced)
with wire_server(provider) as wire:
yield McpPeer(wire.url + "/mcp", observed)
def text_result(text: str) -> JsonRpc:
return {"content": [{"type": "text", "text": text}], "isError": False}
def slow_tool(name: str, seconds: float) -> ScriptedTool:
def respond(params: JsonRpc) -> JsonRpc:
time.sleep(seconds)
return text_result("slept")
return ScriptedTool(name, respond)
def disconnecting_tool(name: str) -> ScriptedTool:
return ScriptedTool(name, lambda params: Reply(chunks=(b'{"jsonrpc":"2.0",', b'"id":1}'), abort_after=1))
def echo_tool(name: str) -> ScriptedTool:
return ScriptedTool(name, lambda params: text_result(json.dumps(params.get("arguments", {}), sort_keys=True)))
@contextmanager
def openapi_peer() -> Iterator[McpPeer]:
"""OpenAPI-described HTTP service plus the spec file the proxy turns into MCP tools."""
observed: Final[queue.Queue[dict[str, object]]] = queue.Queue()
def provider(request: Request) -> Reply:
observed.put(
{
"body": json.loads(request.body) if request.body else None,
"headers": dict(request.headers),
"path": request.target,
"method": request.method,
}
)
if request.target.startswith("/pets/") and request.method == "GET":
return Reply(body=json.dumps({"id": request.target.rsplit("/", 1)[1], "name": "integration-pet"}).encode())
if request.target == "/pets" and request.method == "POST":
return Reply(status=201, body=json.dumps({"created": json.loads(request.body)}).encode())
return Reply(status=404, body=b'{"error":"synthetic not found"}')
with wire_server(provider) as wire:
spec: Final = {
"openapi": "3.0.0",
"info": {"title": "integration pets", "version": "1"},
"servers": [{"url": wire.url}],
"paths": {
"/pets/{petId}": {
"get": {
"operationId": "getPet",
"summary": "Fetch one pet",
"parameters": [{"name": "petId", "in": "path", "required": True, "schema": {"type": "string"}}],
"responses": {"200": {"description": "pet"}},
}
},
"/pets": {
"post": {
"operationId": "createPet",
"summary": "Create a pet",
"requestBody": {
"required": True,
"content": {
"application/json": {
"schema": {
"type": "object",
"properties": {"name": {"type": "string"}},
"required": ["name"],
}
}
},
},
"responses": {"201": {"description": "created"}},
}
},
},
}
yield McpPeer(wire.url, observed, spec_path=_spec_file(spec))
def scratch_directory() -> Path:
path: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", "/tmp")) / "mcp-peers"
path.mkdir(parents=True, exist_ok=True)
return path
def _spec_file(spec: JsonRpc) -> Path:
path: Final = scratch_directory() / f"openapi-{time.monotonic_ns()}.json"
path.write_text(json.dumps(spec))
return path
PeerKind = Literal["http", "sse", "stdio", "openapi"]
PEER_KINDS: Final[tuple[PeerKind, ...]] = ("http", "sse", "stdio", "openapi")
@contextmanager
def peer_of(kind: PeerKind, *, rich: bool = False) -> Iterator[McpPeer]:
if kind == "openapi":
with openapi_peer() as candidate:
yield candidate
elif kind == "stdio":
with stdio_peer(scratch_directory(), rich=rich) as candidate:
yield candidate
else:
with mcp_peer(kind, rich=rich) as candidate:
yield candidate
def register_mcp(scenario: Scenario, peer: McpPeer, alias: str, **fields: object) -> str:
response: Final = scenario.gateway.request(
"POST", "/v1/mcp/server", {"server_name": alias, "alias": alias, "url": peer.url, "transport": "http", **fields}
"POST", "/v1/mcp/server", {"server_name": alias, "alias": alias, **peer.registration(), **fields}
)
identity: Final = response.json()["server_id"]
scenario.cleanups.callback(delete_mcp, scenario.gateway, identity)
scenario.cleanups.callback(forget_mcp, scenario.gateway, identity)
assert response.status_code == 201, response.text
return identity
def forget_mcp(gateway: Gateway, identity: str) -> None:
response: Final = gateway.request("DELETE", f"/v1/mcp/server/{identity}")
assert response.status_code in (202, 404), response.text
def delete_mcp(gateway: Gateway, identity: str) -> None:
response: Final = gateway.request("DELETE", f"/v1/mcp/server/{identity}")
assert response.status_code == 202, response.text
assert read_rows('SELECT server_id FROM "LiteLLM_MCPServerTable" WHERE server_id = %s', (identity,)) == []
def tool_names(gateway: Gateway, key: str, identity: str) -> dict[str, str]:
response: Final = gateway.client.get("/mcp-rest/tools/list", headers={"x-litellm-api-key": key})
def listed_tools(gateway: Gateway, key: str, identity: str | None = None) -> dict[str, dict[str, object]]:
response: Final = gateway.client.get(
"/mcp-rest/tools/list",
headers={"x-litellm-api-key": key},
params={"server_id": identity} if identity else None,
)
assert response.status_code == 200, response.text
return {
name: tool["name"]
tool["name"]: tool
for tool in response.json()["tools"]
if tool.get("mcp_info", {}).get("server_id") == identity
for name in ("add", "multiply", "fail")
if tool["name"].endswith(name)
if identity is None or tool.get("mcp_info", {}).get("server_id") == identity
}
def tool_names(gateway: Gateway, key: str, identity: str) -> dict[str, str]:
return {
name: full
for full in listed_tools(gateway, key, identity)
for name in ("add", "multiply", "fail", "slow", "progress", "sample", "elicit")
if full.endswith(name)
}
@ -99,3 +425,192 @@ def call_tool(gateway: Gateway, key: str, identity: str, name: str, arguments: d
headers={"x-litellm-api-key": key},
json={"server_id": identity, "name": name, "arguments": arguments},
)
EntryPoint = Literal["mcp", "server_mcp", "root", "sse", "rest"]
ENTRY_POINTS: Final[tuple[EntryPoint, ...]] = ("mcp", "server_mcp", "root", "sse", "rest")
INITIALIZE: Final = {
"protocolVersion": "2025-06-18",
"capabilities": {},
"clientInfo": {"name": "integration", "version": "1"},
}
@dataclass(frozen=True, slots=True)
class Outcome:
"""What a caller saw from one MCP operation, normalised across entry points."""
status: int
error: str | None
tools: tuple[str, ...] = ()
text: str | None = None
raw: str = ""
@property
def ok(self) -> bool:
return self.status == 200 and self.error is None
def _parse_rpc_body(response: httpx.Response) -> Mapping[str, object] | None:
if response.headers.get("content-type", "").startswith("text/event-stream"):
data: Final = tuple(line[5:].strip() for line in response.text.splitlines() if line.startswith("data:"))
return json.loads(data[-1]) if data else None
try:
return json.loads(response.text)
except ValueError:
return None
def _outcome_from_rpc(response: httpx.Response) -> Outcome:
body: Final = _parse_rpc_body(response)
if response.status_code != 200 or body is None:
return Outcome(response.status_code, response.text or f"HTTP {response.status_code}", raw=response.text)
if "error" in body:
return Outcome(response.status_code, json.dumps(body["error"]), raw=response.text)
result: Final = body.get("result", {})
assert isinstance(result, dict)
if "tools" in result:
return Outcome(200, None, tuple(tool["name"] for tool in result["tools"]), raw=response.text)
content: Final = result.get("content", [])
text: Final = content[0].get("text") if content else None
if result.get("isError"):
return Outcome(200, text or "isError", text=text, raw=response.text)
return Outcome(200, None, text=text, raw=response.text)
def _outcome_from_rest(response: httpx.Response) -> Outcome:
if response.status_code != 200:
return Outcome(response.status_code, response.text, raw=response.text)
body: Final = response.json()
if "tools" in body:
return Outcome(200, None, tuple(tool["name"] for tool in body["tools"]), raw=response.text)
content: Final = body.get("content", [])
text: Final = content[0].get("text") if content else None
if body.get("isError"):
return Outcome(200, text or "isError", text=text, raw=response.text)
return Outcome(200, None, text=text, raw=response.text)
@dataclass(frozen=True, slots=True)
class McpCaller:
"""One caller's view of the gateway through a specific entry point."""
gateway: Gateway
key: str | None
entry: EntryPoint
alias: str | None = None
headers: Mapping[str, str] = field(default_factory=dict)
def _path(self) -> str:
if self.entry == "server_mcp":
assert self.alias is not None
return f"/{self.alias}/mcp"
return {"mcp": "/mcp", "root": "/mcp/", "sse": "/mcp/sse", "rest": "/mcp-rest"}[self.entry]
def _headers(self) -> dict[str, str]:
return {
**({"x-litellm-api-key": self.key} if self.key is not None else {}),
"Accept": "application/json, text/event-stream",
**self.headers,
}
def rpc(self, method: str, params: JsonRpc | None = None) -> httpx.Response:
if self.entry == "sse":
return _legacy_sse_rpc(self.gateway, self._headers(), method, params)
return self.gateway.client.post(
self._path(),
json={"jsonrpc": "2.0", "id": 1, "method": method, "params": dict(params or {})},
headers=self._headers(),
)
def initialize(self) -> Outcome:
if self.entry == "rest":
return Outcome(200, None)
return _outcome_from_rpc(self.rpc("initialize", INITIALIZE))
def list_tools(self, server_id: str | None = None) -> Outcome:
if self.entry == "rest":
return _outcome_from_rest(
self.gateway.client.get(
"/mcp-rest/tools/list",
headers=self._headers(),
params={"server_id": server_id} if server_id else None,
)
)
return _outcome_from_rpc(self.rpc("tools/list"))
def call(self, name: str, arguments: JsonRpc, server_id: str | None = None) -> Outcome:
if self.entry == "rest":
return _outcome_from_rest(
self.gateway.client.post(
"/mcp-rest/tools/call",
headers=self._headers(),
json={
"name": name,
"arguments": dict(arguments),
**({"server_id": server_id} if server_id else {}),
},
)
)
return _outcome_from_rpc(self.rpc("tools/call", {"name": name, "arguments": dict(arguments)}))
def _legacy_sse_rpc(
gateway: Gateway, headers: Mapping[str, str], method: str, params: JsonRpc | None
) -> httpx.Response:
"""Drive the legacy GET /mcp/sse + POST /mcp/sse/messages pair for one request and synthesise a JSON response."""
with gateway.client.stream("GET", "/mcp/sse", headers=headers, timeout=15) as stream:
if stream.status_code != 200:
stream.read()
return httpx.Response(stream.status_code, text=stream.text)
lines: Final = stream.iter_lines()
endpoint: Final = next(line[5:].strip() for line in lines if line.startswith("data:"))
init: Final = gateway.client.post(
endpoint,
json={"jsonrpc": "2.0", "id": 0, "method": "initialize", "params": INITIALIZE},
headers=headers,
)
assert init.status_code in (200, 202), init.text
gateway.client.post(endpoint, json={"jsonrpc": "2.0", "method": "notifications/initialized"}, headers=headers)
posted: Final = gateway.client.post(
endpoint, json={"jsonrpc": "2.0", "id": 1, "method": method, "params": dict(params or {})}, headers=headers
)
if posted.status_code not in (200, 202):
return httpx.Response(posted.status_code, text=posted.text)
for line in lines:
if line.startswith("data:") and '"id": 1' in line.replace('"id":1', '"id": 1'):
return httpx.Response(200, text=line[5:].strip(), headers={"content-type": "application/json"})
return httpx.Response(599, text="legacy SSE stream ended without a reply")
def official_client_outcomes(
gateway: Gateway, key: str, path: str, name: str, arguments: JsonRpc, *, legacy_sse: bool = False
) -> tuple[Outcome, Outcome]:
"""List then call through the official MCP client session, returning both outcomes."""
url: Final = str(gateway.client.base_url).rstrip("/") + path
headers: Final = {"x-litellm-api-key": key}
async def run() -> tuple[Outcome, Outcome]:
transport: Final = (
sse_client(url, headers=headers)
if legacy_sse
else streamable_http_client(url, http_client=httpx.AsyncClient(headers=headers, timeout=30))
)
async with transport as streams, ClientSession(streams[0], streams[1]) as session:
await session.initialize()
listed: Final = await session.list_tools()
result: Final = await session.call_tool(name, dict(arguments))
content: Final = result.content[0] if result.content else None
text: Final = content.text if isinstance(content, TextContent) else None
return (
Outcome(200, None, tuple(tool.name for tool in listed.tools)),
Outcome(200, (text or "isError") if result.is_error else None, text=text),
)
return asyncio.run(run())
def tool_calls(observed: tuple[dict[str, object], ...]) -> tuple[dict[str, object], ...]:
return tuple(
item for item in observed if isinstance(item.get("body"), dict) and item["body"].get("method") == "tools/call"
)

151
tests/integration/_support/mcp_grants.py Normal file
View file

@ -0,0 +1,151 @@
import uuid
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Final, Literal
from integration._support.client import Gateway, Scenario, string_value
Subject = Literal["key", "team", "org", "user", "end_user", "agent", "access_group", "toolset", "allowed_tools"]
SUBJECTS: Final[tuple[Subject, ...]] = (
"key",
"team",
"org",
"user",
"end_user",
"agent",
"access_group",
"toolset",
"allowed_tools",
)
@dataclass(frozen=True, slots=True)
class Caller:
"""A key plus the request headers that make the proxy resolve the granted subject."""
key: str
headers: Mapping[str, str]
def _mcp_permission(server_ids: tuple[str, ...]) -> dict[str, list[str]]:
return {"mcp_servers": list(server_ids)}
def delete_organization(gateway: Gateway, identity: str) -> None:
response: Final = gateway.request("DELETE", "/organization/delete", {"organization_ids": [identity]})
assert response.status_code == 200, response.text
def delete_end_user(gateway: Gateway, identity: str) -> None:
response: Final = gateway.request("POST", "/end_user/delete", {"user_ids": [identity]})
assert response.status_code == 200, response.text
def delete_agent(gateway: Gateway, identity: str) -> None:
response: Final = gateway.request("DELETE", f"/v1/agents/{identity}")
assert response.status_code == 200, response.text
def delete_toolset(gateway: Gateway, identity: str) -> None:
response: Final = gateway.request("DELETE", f"/v1/mcp/toolset/{identity}")
assert response.status_code in (200, 202, 204), response.text
def create_toolset(scenario: Scenario, tools: tuple[tuple[str, str], ...]) -> str:
response: Final = scenario.gateway.request(
"POST",
"/v1/mcp/toolset",
{
"toolset_name": f"integration-{uuid.uuid4().hex[:10]}",
"tools": [{"server_id": server_id, "tool_name": tool} for server_id, tool in tools],
},
)
assert response.status_code == 201, response.text
identity: Final = string_value(response.json()["toolset_id"])
scenario.cleanups.callback(delete_toolset, scenario.gateway, identity)
return identity
def grant(
scenario: Scenario,
subject: Subject,
granted: tuple[str, ...],
ceiling: tuple[str, ...],
*,
access_group: str | None = None,
allowed_tools: Mapping[str, tuple[str, ...]] | None = None,
) -> Caller:
"""Build a caller whose ``subject`` level grants exactly ``granted`` out of ``ceiling``.
``ceiling`` is what the key itself can reach before the subject narrows it; the key subject grants
``granted`` directly. Access groups take the group name that the granted servers were registered with,
and ``allowed_tools`` maps server id to the tools the key may call on it."""
gateway: Final = scenario.gateway
match subject:
case "key":
return Caller(scenario.key(object_permission=_mcp_permission(granted)), {})
case "team":
team: Final = scenario.team(object_permission=_mcp_permission(granted))
return Caller(scenario.key(team_id=team), {})
case "org":
created: Final = gateway.post(
"/organization/new",
{
"organization_alias": f"integration-{uuid.uuid4().hex[:10]}",
"object_permission": _mcp_permission(granted),
},
)
org: Final = string_value(created["organization_id"])
scenario.cleanups.callback(delete_organization, gateway, org)
org_team: Final = scenario.team(organization_id=org, object_permission=_mcp_permission(ceiling))
return Caller(scenario.key(team_id=org_team), {})
case "user":
user: Final = scenario.user(object_permission=_mcp_permission(granted))
return Caller(scenario.key(user_id=user, object_permission=_mcp_permission(ceiling)), {})
case "end_user":
end_user: Final = f"integration-{uuid.uuid4().hex[:10]}"
response: Final = gateway.request(
"POST", "/end_user/new", {"user_id": end_user, "object_permission": _mcp_permission(granted)}
)
assert response.status_code == 200, response.text
scenario.cleanups.callback(delete_end_user, gateway, end_user)
return Caller(scenario.key(object_permission=_mcp_permission(ceiling)), {"x-litellm-end-user-id": end_user})
case "agent":
agent: Final = gateway.post(
"/v1/agents",
{
"agent_name": f"integration-{uuid.uuid4().hex[:10]}",
"agent_card_params": {
"protocolVersion": "0.3.0",
"name": "integration",
"description": "integration agent",
"url": "http://127.0.0.1:1/agent",
"version": "1",
"capabilities": {},
"defaultInputModes": ["text"],
"defaultOutputModes": ["text"],
"skills": [],
},
"object_permission": _mcp_permission(granted),
},
)
agent_id: Final = string_value(agent["agent_id"])
scenario.cleanups.callback(delete_agent, gateway, agent_id)
return Caller(scenario.key(agent_id=agent_id, object_permission=_mcp_permission(ceiling)), {})
case "access_group":
assert access_group is not None
return Caller(scenario.key(object_permission={"mcp_access_groups": [access_group]}), {})
case "toolset":
toolset: Final = create_toolset(scenario, tuple((server, "add") for server in granted))
return Caller(scenario.key(object_permission={"mcp_toolsets": [toolset]}), {})
case "allowed_tools":
assert allowed_tools is not None
return Caller(
scenario.key(
object_permission={
"mcp_servers": list(granted),
"mcp_tool_permissions": {server: list(tools) for server, tools in allowed_tools.items()},
}
),
{},
)

49
tests/integration/_support/mcp_stdio_peer.py Normal file
View file

@ -0,0 +1,49 @@
"""Stdio MCP peer the proxy spawns; every inbound JSON-RPC line is appended to the record file."""
import sys
from pathlib import Path
sys.path[0] = str(Path(__file__).resolve().parents[2])
import asyncio # noqa: E402 # the script directory holds mcp.py, which would shadow the mcp package
import json # noqa: E402
import os # noqa: E402
from typing import Final # noqa: E402
import anyio # noqa: E402
from integration._support.mcp import math_service # noqa: E402
from mcp.server.stdio import stdio_server # noqa: E402
class Recording:
def __init__(self, source: anyio.AsyncFile[str], record: Path) -> None:
self.source = source
self.record = record
def __aiter__(self) -> "Recording":
return self
async def __anext__(self) -> str:
line: Final = await self.source.readline()
if not line:
raise StopAsyncIteration
with self.record.open("a") as sink:
passed: Final = {name: value for name, value in os.environ.items() if name.startswith("PEER_")}
sink.write(json.dumps({"body": json.loads(line), "env": passed}) + "\n")
return line
async def readline(self) -> str:
return await self.__anext__()
async def main() -> None:
record: Final = Path(sys.argv[1])
service: Final = math_service("integration-stdio", rich=sys.argv[2] == "rich")
stdin: Final = anyio.wrap_file(sys.stdin)
async with stdio_server(stdin=Recording(stdin, record)) as (read_stream, write_stream):
lowlevel: Final = service._lowlevel_server
await lowlevel.run(read_stream, write_stream, lowlevel.create_initialization_options())
if __name__ == "__main__":
asyncio.run(main())

198
tests/integration/_support/oauth_server.py Normal file
View file

@ -0,0 +1,198 @@
"""OAuth 2.1 authorization-server double: metadata, DCR, PKCE authorization code, refresh, client credentials,
token exchange and revocation, every request recorded."""
import base64
import hashlib
import json
import secrets
import threading
import uuid
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Final
from urllib.parse import parse_qs, urlencode, urlsplit
from integration._support.wire import Reply, Request, Wire, wire_server
TOKEN_EXCHANGE: Final = "urn:ietf:params:oauth:grant-type:token-exchange"
@dataclass(slots=True)
class AuthorizationServer:
wire: Wire
clients: dict[str, str] = field(default_factory=dict)
codes: dict[str, dict[str, str]] = field(default_factory=dict)
access_tokens: dict[str, dict[str, str]] = field(default_factory=dict)
refresh_tokens: dict[str, dict[str, str]] = field(default_factory=dict)
revoked: set[str] = field(default_factory=set)
lock: threading.Lock = field(default_factory=threading.Lock)
@property
def issuer(self) -> str:
return self.wire.url
def drain(self) -> tuple[Request, ...]:
return self.wire.drain()
def token_requests(self) -> tuple[dict[str, str], ...]:
return tuple(
{name: values[0] for name, values in parse_qs(item.body.decode()).items()}
for item in self.drain()
if item.target.startswith("/token")
)
def is_live(self, token: str) -> bool:
with self.lock:
return token in self.access_tokens and token not in self.revoked
def issue(self, grant: str, client_id: str, subject: str, scope: str) -> dict[str, object]:
access: Final = f"at-{grant}-{secrets.token_urlsafe(8)}"
refresh: Final = f"rt-{secrets.token_urlsafe(8)}"
with self.lock:
self.access_tokens[access] = {"client_id": client_id, "subject": subject, "scope": scope, "grant": grant}
self.refresh_tokens[refresh] = {"client_id": client_id, "subject": subject, "scope": scope}
return {
"access_token": access,
"token_type": "Bearer",
"expires_in": 3600,
"refresh_token": refresh,
"scope": scope,
}
def _pkce_matches(challenge: str, verifier: str) -> bool:
digest: Final = hashlib.sha256(verifier.encode()).digest()
return base64.urlsafe_b64encode(digest).rstrip(b"=").decode() == challenge
def _json(status: int, body: dict[str, object]) -> Reply:
return Reply(status=status, body=json.dumps(body).encode())
def _client_credentials(request: Request, form: dict[str, str]) -> tuple[str, str | None]:
header: Final = request.headers.get("authorization", "")
if header.lower().startswith("basic "):
decoded: Final = base64.b64decode(header.split(" ", 1)[1]).decode()
client_id, _, secret = decoded.partition(":")
return client_id, secret
return form.get("client_id", ""), form.get("client_secret")
@contextmanager
def oauth_server(*, scopes: tuple[str, ...] = ("tools.read", "tools.call")) -> Iterator[AuthorizationServer]:
holder: list[AuthorizationServer] = []
def respond(request: Request) -> Reply:
server: Final = holder[0]
path: Final = urlsplit(request.target).path
query: Final = {name: values[0] for name, values in parse_qs(urlsplit(request.target).query).items()}
form: Final = {name: values[0] for name, values in parse_qs(request.body.decode()).items()}
if path.startswith("/.well-known/oauth-authorization-server") or path == "/.well-known/openid-configuration":
return _json(
200,
{
"issuer": server.issuer,
"authorization_endpoint": server.issuer + "/authorize",
"token_endpoint": server.issuer + "/token",
"registration_endpoint": server.issuer + "/register",
"revocation_endpoint": server.issuer + "/revoke",
"introspection_endpoint": server.issuer + "/introspect",
"scopes_supported": list(scopes),
"response_types_supported": ["code"],
"grant_types_supported": [
"authorization_code",
"refresh_token",
"client_credentials",
TOKEN_EXCHANGE,
],
"code_challenge_methods_supported": ["S256"],
"token_endpoint_auth_methods_supported": ["client_secret_post", "client_secret_basic", "none"],
},
)
if path == "/register" and request.method == "POST":
metadata: Final = json.loads(request.body or b"{}")
client_id: Final = f"dcr-{uuid.uuid4().hex[:12]}"
secret: Final = f"secret-{secrets.token_urlsafe(8)}"
with server.lock:
server.clients[client_id] = secret
return _json(
201,
{
"client_id": client_id,
"client_secret": secret,
"client_id_issued_at": 0,
"redirect_uris": metadata.get("redirect_uris", []),
"grant_types": metadata.get("grant_types", ["authorization_code"]),
"token_endpoint_auth_method": metadata.get("token_endpoint_auth_method", "client_secret_post"),
},
)
if path == "/authorize" and request.method == "GET":
missing: Final = tuple(
name for name in ("client_id", "redirect_uri", "code_challenge", "state") if name not in query
)
if missing or query.get("code_challenge_method", "S256") != "S256" or query.get("response_type") != "code":
return _json(400, {"error": "invalid_request", "missing": list(missing), "received": query})
code: Final = f"code-{secrets.token_urlsafe(8)}"
with server.lock:
server.codes[code] = {
"client_id": query["client_id"],
"redirect_uri": query["redirect_uri"],
"code_challenge": query["code_challenge"],
"scope": query.get("scope", " ".join(scopes)),
}
location: Final = (
query["redirect_uri"]
+ ("&" if "?" in query["redirect_uri"] else "?")
+ urlencode({"code": code, "state": query["state"]})
)
return Reply(status=302, body=b"", headers={"location": location})
if path == "/token" and request.method == "POST":
grant: Final = form.get("grant_type", "")
client_id, client_secret = _client_credentials(request, form)
if grant == "authorization_code":
with server.lock:
issued: Final = server.codes.pop(form.get("code", ""), None)
if issued is None:
return _json(400, {"error": "invalid_grant", "error_description": "unknown or reused code"})
if issued["client_id"] != client_id:
return _json(400, {"error": "invalid_client", "error_description": "code issued to another client"})
if not _pkce_matches(issued["code_challenge"], form.get("code_verifier", "")):
return _json(400, {"error": "invalid_grant", "error_description": "pkce verifier mismatch"})
return _json(200, server.issue("authorization_code", client_id, "integration-user", issued["scope"]))
if grant == "refresh_token":
with server.lock:
known: Final = server.refresh_tokens.pop(form.get("refresh_token", ""), None)
if known is None:
return _json(400, {"error": "invalid_grant", "error_description": "unknown refresh token"})
return _json(200, server.issue("refresh_token", known["client_id"], known["subject"], known["scope"]))
if grant == "client_credentials":
with server.lock:
expected: Final = server.clients.get(client_id)
if not client_id or (expected is not None and expected != client_secret) or not client_secret:
return _json(401, {"error": "invalid_client"})
return _json(200, server.issue("client_credentials", client_id, client_id, form.get("scope", "")))
if grant == TOKEN_EXCHANGE:
subject: Final = form.get("subject_token", "")
if not subject:
return _json(400, {"error": "invalid_request", "error_description": "subject_token required"})
if not client_id:
return _json(401, {"error": "invalid_client"})
token: Final = server.issue("token_exchange", client_id, f"exchanged:{subject}", form.get("scope", ""))
return _json(200, {**token, "issued_token_type": "urn:ietf:params:oauth:token-type:access_token"})
return _json(400, {"error": "unsupported_grant_type", "grant_type": grant})
if path == "/revoke" and request.method == "POST":
with server.lock:
server.revoked.add(form.get("token", ""))
return Reply(status=200, body=b"{}")
if path == "/introspect" and request.method == "POST":
token: Final = form.get("token", "")
with server.lock:
info: Final = server.access_tokens.get(token)
active: Final = info is not None and token not in server.revoked
return _json(200, {"active": active, **(info or {})})
return _json(404, {"error": "not_found", "path": path, "method": request.method})
with wire_server(respond) as wire:
holder.append(AuthorizationServer(wire))
yield holder[0]

8
tests/integration/_support/proxy.py
View file

@ -1,11 +1,19 @@
"""Run the normal single-process CLI with the existing behavior-suite test entitlement."""
import signal
import sys
from types import FrameType
from unittest.mock import patch
from litellm import run_server
def _exit_on_reraised_term(signum: int, frame: FrameType | None) -> None:
sys.exit(0)
def main() -> None:
signal.signal(signal.SIGTERM, _exit_on_reraised_term)
with patch( # test-quality-ok: route entitlement only; license validation is outside these HTTP/DB contracts
"litellm.proxy.auth.litellm_license.LicenseCheck.is_premium", return_value=True
):

43
tests/integration/conftest.py
View file

@ -14,7 +14,7 @@ from redis import Redis
from tests.integration._support.client import Gateway, eventually, gateway_from_environment
from tests.integration._support.generation import LIFECYCLE_SETTINGS
from tests.integration._support.manifest import OWNED_DIRECTORIES, contracts
from tests.integration._support.manifest import OWNED_DIRECTORIES
COLLECTED: Final = pytest.StashKey[tuple[str, ...]]()
REPORTS: Final = pytest.StashKey[list[pytest.TestReport]]()
@ -26,7 +26,7 @@ def pytest_addoption(parser: pytest.Parser) -> None:
def pytest_configure(config: pytest.Config) -> None:
config.addinivalue_line("markers", "integration: owned real-service integration contracts")
config.addinivalue_line("markers", "covers(*ids): independently asserted behavior contracts")
config.addinivalue_line("markers", "covers(*ids): legacy contract IDs kept for existing tests, not enforced")
config.stash[REPORTS] = []
config.pluginmanager.register(IntegrationReportPlugin(config))
@ -53,7 +53,6 @@ def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item
if order_seed:
# rebind-ok: pytest requires this hook to reorder its shared collection list in place.
items.sort(key=lambda item: hashlib.sha256(f"{order_seed}:{item.nodeid}".encode()).digest())
manifest: Final = contracts()
root: Final = Path(__file__).parent
owned: Final = tuple(
item
@ -63,12 +62,7 @@ def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item
if owned and os.environ.get("GITHUB_ACTIONS") == "true":
raise pytest.UsageError("Integration contracts are owned by CircleCI")
for item in owned:
if item.nodeid not in manifest:
raise pytest.UsageError(f"Integration node missing from manifest: {item.nodeid}")
item.add_marker(pytest.mark.integration)
declared: Final = tuple(value for mark in item.iter_markers("covers") for value in mark.args)
if set(declared) != set(manifest[item.nodeid]):
raise pytest.UsageError(f"Contract mapping differs for {item.nodeid}")
config.stash[COLLECTED] = tuple(item.nodeid for item in owned)
@ -91,22 +85,25 @@ def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None:
output: Final = Path(destination)
output.mkdir(parents=True, exist_ok=True)
(output / "execution.json").write_text(
json.dumps({
"collected": collected,
"passed": passed,
"skipped": skipped,
"complete": complete,
"exitstatus": exitstatus,
"hypothesis_version": version("hypothesis"),
"hypothesis_seed": session.config.getoption("hypothesis_seed"),
"order_seed": session.config.getoption("integration_order_seed"),
"generation": {
"max_examples": LIFECYCLE_SETTINGS.max_examples,
"stateful_step_count": LIFECYCLE_SETTINGS.stateful_step_count,
"database": str(LIFECYCLE_SETTINGS.database),
"phases": [phase.name for phase in LIFECYCLE_SETTINGS.phases],
json.dumps(
{
"collected": collected,
"passed": passed,
"skipped": skipped,
"complete": complete,
"exitstatus": exitstatus,
"hypothesis_version": version("hypothesis"),
"hypothesis_seed": session.config.getoption("hypothesis_seed"),
"order_seed": session.config.getoption("integration_order_seed"),
"generation": {
"max_examples": LIFECYCLE_SETTINGS.max_examples,
"stateful_step_count": LIFECYCLE_SETTINGS.stateful_step_count,
"database": str(LIFECYCLE_SETTINGS.database),
"phases": [phase.name for phase in LIFECYCLE_SETTINGS.phases],
},
},
}, indent=2)
indent=2,
)
+ "\n"
)
if not complete and exitstatus == 0:

2103
tests/integration/contracts.json
View file

File diff suppressed because it is too large Load diff

136
tests/integration/management/test_model_credential_name_updates.py Normal file
View file

@ -0,0 +1,136 @@
import uuid
from typing import Final
import httpx
import pytest
from pydantic import JsonValue
from tests.integration._support.client import JSON_OBJECT, Gateway, Scenario, object_value, string_value
from tests.integration._support.database import read_rows
def _dangling_credential(gateway: Gateway, scenario: Scenario) -> str:
name: Final = f"credential-{uuid.uuid4().hex}"
gateway.post(
"/credentials",
{"credential_name": name, "credential_values": {"api_key": "synthetic-credential"}, "credential_info": {}},
)
scenario.cleanups.callback(_delete_credential_if_present, gateway, name)
return name
def _delete_credential_if_present(gateway: Gateway, name: str) -> None:
response: Final = gateway.request("DELETE", f"/credentials/{name}")
assert response.status_code in (200, 404), response.text
def _delete_credential(gateway: Gateway, name: str) -> None:
response: Final = gateway.request("DELETE", f"/credentials/{name}")
assert response.status_code == 200, response.text
assert read_rows('SELECT credential_name FROM "LiteLLM_CredentialsTable" WHERE credential_name = %s', (name,)) == []
def _model_with_credential(gateway: Gateway, scenario: Scenario, credential: str, **model_info: JsonValue) -> str:
created: Final = gateway.post(
"/model/new",
{
"model_name": f"integration-{uuid.uuid4().hex}",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_base": f"{gateway.upstream_url}/v1",
"litellm_credential_name": credential,
"rpm": 5,
},
"model_info": dict(model_info),
},
)
identity: Final = string_value(object_value(created["model_info"])["id"])
scenario.cleanups.callback(scenario.delete_model, identity)
return identity
def _stored_params(gateway: Gateway, identity: str) -> dict[str, JsonValue]:
entries: Final = gateway.get("/model/info", {"litellm_model_id": identity})["data"]
assert isinstance(entries, list) and len(entries) == 1, entries
return object_value(object_value(entries[0])["litellm_params"])
def _error(response: httpx.Response) -> dict[str, JsonValue]:
return object_value(JSON_OBJECT.validate_json(response.content)["error"])
@pytest.mark.covers("mgmt.model.update.unchanged_credential_name_is_not_revalidated")
def test_unrelated_patch_succeeds_when_resent_credential_name_is_dangling(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
credential: Final = _dangling_credential(gateway, scenario)
identity: Final = _model_with_credential(gateway, scenario, credential)
_delete_credential(gateway, credential)
before: Final = _stored_params(gateway, identity)
assert before["litellm_credential_name"] == credential
assert before["rpm"] == 5
patched: Final = gateway.request(
"PATCH",
f"/model/{identity}/update",
{"litellm_params": {"litellm_credential_name": before["litellm_credential_name"], "rpm": 7}},
)
assert patched.status_code == 200, patched.text
after: Final = _stored_params(gateway, identity)
assert after == {**before, "rpm": 7}
@pytest.mark.covers(
"mgmt.model.update.non_admin_detach_is_rejected",
"mgmt.model.update.empty_credential_name_is_rejected",
)
def test_non_admin_detach_and_empty_credential_name_still_rejected(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
credential: Final = _dangling_credential(gateway, scenario)
user: Final = scenario.user(user_role="internal_user")
team: Final = scenario.team(members_with_roles=[{"user_id": user, "role": "admin"}])
team_admin: Final = scenario.key(user_id=user, team_id=team)
identity: Final = _model_with_credential(gateway, scenario, credential, team_id=team)
before: Final = _stored_params(gateway, identity)
detached: Final = gateway.request(
"PATCH", f"/model/{identity}/update", {"litellm_params": {"litellm_credential_name": None}}, key=team_admin
)
assert detached.status_code == 403, detached.text
assert _error(detached) == {
"message": "Only a proxy admin can detach a stored credential (litellm_credential_name) on a model. "
"Your role=internal_user.",
"type": "auth_error",
"param": "litellm_credential_name",
"code": "403",
}
emptied: Final = gateway.request(
"PATCH", f"/model/{identity}/update", {"litellm_params": {"litellm_credential_name": ""}}
)
assert emptied.status_code == 400, emptied.text
assert _error(emptied) == {
"message": "litellm_credential_name cannot be an empty string. Send null to detach the stored credential "
"or omit the field to leave it unchanged.",
"type": "validation_error",
"param": "litellm_credential_name",
"code": "400",
}
assert _stored_params(gateway, identity) == before
@pytest.mark.covers("mgmt.model.update.changed_missing_credential_name_is_rejected")
def test_changing_credential_name_to_missing_credential_is_rejected(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
credential: Final = _dangling_credential(gateway, scenario)
identity: Final = _model_with_credential(gateway, scenario, credential)
_delete_credential(gateway, credential)
before: Final = _stored_params(gateway, identity)
missing: Final = f"credential-{uuid.uuid4().hex}"
rejected: Final = gateway.request(
"PATCH", f"/model/{identity}/update", {"litellm_params": {"litellm_credential_name": missing, "rpm": 7}}
)
assert rejected.status_code == 400, rejected.text
assert _error(rejected) == {
"message": f"Credential '{missing}' not found. Create it via /credentials before attaching it to a model.",
"type": "validation_error",
"param": "litellm_credential_name",
"code": "400",
}
assert _stored_params(gateway, identity) == before

124
tests/integration/mcp/test_mcp_access_matrix.py Normal file
View file

@ -0,0 +1,124 @@
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.mcp import (
ENTRY_POINTS,
EntryPoint,
McpCaller,
McpPeer,
PeerKind,
peer_of,
register_mcp,
tool_calls,
)
from integration._support.mcp_grants import SUBJECTS, Subject, grant
CALLABLE: Final = {"add": {"a": 1, "b": 2}, "multiply": {"a": 2, "b": 3}}
RESULTS: Final = {"add": "3", "multiply": "6"}
def _server_scoped(entry: EntryPoint, identity: str) -> str | None:
return identity if entry == "rest" else None
def _name(entry: EntryPoint, alias: str, tool: str) -> str:
return tool if entry == "rest" else f"{alias}-{tool}"
def _assert_denied(caller: McpCaller, peer: McpPeer, name: str, identity: str, entry: EntryPoint) -> None:
peer.drain()
outcome: Final = caller.call(name, CALLABLE["add"], _server_scoped(entry, identity))
assert outcome.error is not None, f"denied call succeeded: {outcome.raw}"
assert outcome.text not in RESULTS.values(), outcome.raw
assert tool_calls(peer.drain()) == (), "denied call reached the peer"
@pytest.mark.parametrize("entry", ENTRY_POINTS)
@pytest.mark.parametrize("subject", SUBJECTS)
@pytest.mark.parametrize("peer_kind", ("http", "sse"))
def test_subject_grant_lists_only_reachable_tools_and_denies_the_rest(
gateway: Gateway, peer_kind: PeerKind, subject: Subject, entry: EntryPoint
) -> None:
with peer_of(peer_kind) as granted_peer, peer_of(peer_kind) as denied_peer, gateway.scenario() as scenario:
group: Final = "grp" + uuid.uuid4().hex[:8]
granted_alias: Final = "yes" + uuid.uuid4().hex[:8]
denied_alias: Final = "no" + uuid.uuid4().hex[:8]
granted: Final = register_mcp(scenario, granted_peer, granted_alias, mcp_access_groups=[group])
denied: Final = register_mcp(scenario, denied_peer, denied_alias)
caller: Final = grant(
scenario, subject, (granted,), (granted, denied), access_group=group, allowed_tools={granted: ("add",)}
)
reach: Final = McpCaller(gateway, caller.key, entry, granted_alias, caller.headers)
listed: Final = reach.list_tools(_server_scoped(entry, granted))
assert listed.ok, listed.raw
expected: Final = (
{_name(entry, granted_alias, "add")}
if subject in ("toolset", "allowed_tools")
else {_name(entry, granted_alias, tool) for tool in ("add", "multiply", "fail")}
)
assert set(listed.tools) == expected, listed.tools
for tool, arguments in CALLABLE.items():
name: Final = _name(entry, granted_alias, tool)
if name not in listed.tools:
continue
granted_peer.drain()
outcome: Final = reach.call(name, arguments, _server_scoped(entry, granted))
assert outcome.ok and outcome.text == RESULTS[tool], outcome.raw
assert [call["body"]["params"]["name"] for call in tool_calls(granted_peer.drain())] == [tool]
if subject in ("toolset", "allowed_tools"):
_assert_denied(reach, granted_peer, _name(entry, granted_alias, "multiply"), granted, entry)
blocked: Final = McpCaller(gateway, caller.key, entry, denied_alias, caller.headers)
_assert_denied(blocked, denied_peer, _name(entry, denied_alias, "add"), denied, entry)
denied_listed: Final = blocked.list_tools(_server_scoped(entry, denied))
if entry == "rest":
assert denied_listed.status == 403 and "access_denied" in denied_listed.raw, denied_listed.raw
assert denied_listed.tools == ()
else:
assert not any(name.startswith(denied_alias) for name in denied_listed.tools), denied_listed.tools
@pytest.mark.parametrize("entry", ("mcp", "server_mcp", "rest"))
def test_key_without_any_grant_sees_no_scoped_server(gateway: Gateway, entry: EntryPoint) -> None:
with peer_of("http") as peer, gateway.scenario() as scenario:
alias: Final = "none" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias)
other: Final = scenario.key(object_permission={"mcp_servers": ["no-mcp-servers"]})
caller: Final = McpCaller(gateway, other, entry, alias)
_assert_denied(caller, peer, _name(entry, alias, "add"), identity, entry)
listed: Final = caller.list_tools(_server_scoped(entry, identity))
assert not any(name.startswith(alias) for name in listed.tools), listed.tools
@pytest.mark.parametrize("entry", ("mcp", "server_mcp", "rest", "root", "sse"))
def test_missing_or_wrong_key_is_rejected_before_the_peer(gateway: Gateway, entry: EntryPoint) -> None:
with peer_of("http") as peer, gateway.scenario() as scenario:
alias: Final = "anon" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias)
for key in (None, "sk-integration-wrong-" + uuid.uuid4().hex):
caller: Final = McpCaller(gateway, key, entry, alias)
peer.drain()
outcome: Final = caller.call(_name(entry, alias, "add"), CALLABLE["add"], _server_scoped(entry, identity))
assert outcome.status in (401, 403) or outcome.error is not None, outcome.raw
assert outcome.text not in RESULTS.values(), outcome.raw
assert tool_calls(peer.drain()) == ()
def test_same_tool_name_on_two_servers_routes_by_prefix(gateway: Gateway) -> None:
with peer_of("http") as first, peer_of("sse") as second, gateway.scenario() as scenario:
first_alias: Final = "one" + uuid.uuid4().hex[:8]
second_alias: Final = "two" + uuid.uuid4().hex[:8]
first_id: Final = register_mcp(scenario, first, first_alias)
second_id: Final = register_mcp(scenario, second, second_alias)
key: Final = scenario.key(object_permission={"mcp_servers": [first_id, second_id]})
caller: Final = McpCaller(gateway, key, "mcp", None)
listed: Final = caller.list_tools()
assert listed.ok and len(listed.tools) == len(set(listed.tools)) == 6, listed.tools
assert {f"{first_alias}-add", f"{second_alias}-add"} <= set(listed.tools)
first.drain()
second.drain()
outcome: Final = caller.call(f"{second_alias}-add", {"a": 5, "b": 5})
assert outcome.ok and outcome.text == "10", outcome.raw
assert tool_calls(first.drain()) == ()
assert [call["body"]["params"]["name"] for call in tool_calls(second.drain())] == ["add"]

197
tests/integration/mcp/test_mcp_accounting_guardrails.py Normal file
View file

@ -0,0 +1,197 @@
import uuid
from collections.abc import Iterator
from contextlib import contextmanager
from hashlib import sha256
from typing import Final
import pytest
from integration._support.client import Gateway, JsonValue, Scenario, eventually
from integration._support.database import read_rows
from integration._support.mcp import (
ENTRY_POINTS,
EntryPoint,
McpCaller,
McpPeer,
Outcome,
mcp_peer,
register_mcp,
tool_calls,
)
DEFAULT_COST: Final = 0.25
ADD_COST: Final = 0.5
FORBIDDEN: Final = "forbidden-integration-word"
SPEND_ROWS: Final = (
'SELECT call_type, model, spend, status, metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s ORDER BY "startTime"'
)
def _digest(key: str) -> str:
return sha256(key.encode()).hexdigest()
def _rows(key: str, count: int) -> list[dict[str, JsonValue]]:
return eventually(lambda: read_rows(SPEND_ROWS, (_digest(key),)), lambda rows: len(rows) >= count, seconds=70)
def _priced_server(scenario: Scenario, peer: McpPeer, alias: str) -> str:
return register_mcp(
scenario,
peer,
alias,
mcp_info={
"server_name": alias,
"mcp_server_cost_info": {
"default_cost_per_query": DEFAULT_COST,
"tool_name_to_cost_per_query": {"add": ADD_COST},
},
},
)
def _tool_metadata(row: dict[str, JsonValue]) -> dict[str, JsonValue]:
metadata: Final = row["metadata"]
assert isinstance(metadata, dict), row
tool: Final = metadata.get("mcp_tool_call_metadata")
assert isinstance(tool, dict), metadata
return tool
def _call(caller: McpCaller, name: str, arguments: dict[str, object], entry: EntryPoint, identity: str) -> Outcome:
return caller.call(name, arguments, identity if entry == "rest" else None)
@pytest.mark.parametrize("entry", ENTRY_POINTS)
def test_each_tool_call_writes_one_spend_row_with_server_tool_and_configured_cost(
gateway: Gateway, entry: EntryPoint
) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "acct" + uuid.uuid4().hex[:8]
identity: Final = _priced_server(scenario, peer, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
caller: Final = McpCaller(gateway, key, entry, alias)
peer.drain()
assert _call(caller, f"{alias}-add", {"a": 2, "b": 3}, entry, identity).text == "5"
assert _call(caller, f"{alias}-multiply", {"a": 2, "b": 3}, entry, identity).text == "6"
assert len(tool_calls(peer.drain())) == 2
rows: Final = [row for row in _rows(key, 2) if row["call_type"] == "call_mcp_tool"]
assert len(rows) == 2, rows
by_tool: Final = {_tool_metadata(row)["name"]: row for row in rows}
assert set(by_tool) == {"add", "multiply"}, rows
assert float(str(by_tool["add"]["spend"])) == pytest.approx(ADD_COST)
assert float(str(by_tool["multiply"]["spend"])) == pytest.approx(DEFAULT_COST)
for row in rows:
assert _tool_metadata(row)["mcp_server_name"] == alias, row
assert row["model"] == f"MCP: {alias}-{_tool_metadata(row)['name']}", row
later: Final = read_rows(SPEND_ROWS, (_digest(key),))
assert len([row for row in later if row["call_type"] == "call_mcp_tool"]) == 2, later
def test_key_spend_and_key_max_budget_count_mcp_tool_calls(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "budget" + uuid.uuid4().hex[:8]
identity: Final = _priced_server(scenario, peer, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]}, max_budget=ADD_COST / 2)
caller: Final = McpCaller(gateway, key, "mcp", alias)
assert caller.call(f"{alias}-add", {"a": 2, "b": 3}).text == "5"
info: Final = eventually(
lambda: gateway.client.get("/key/info", params={"key": key}, headers={"x-litellm-api-key": gateway.key}),
lambda response: response.status_code == 200 and float(response.json()["info"]["spend"]) > 0,
seconds=70,
)
assert float(info.json()["info"]["spend"]) == pytest.approx(ADD_COST)
eventually(
lambda: caller.call(f"{alias}-add", {"a": 2, "b": 3}),
lambda outcome: outcome.error is not None,
seconds=70,
)
peer.drain()
denied: Final = caller.call(f"{alias}-add", {"a": 2, "b": 3})
assert denied.error is not None and "budget" in str(denied.raw).lower(), denied.raw
assert tool_calls(peer.drain()) == (), "over-budget call reached the peer"
@contextmanager
def _content_filter(gateway: Gateway, mode: str) -> Iterator[str]:
name: Final = "filter" + uuid.uuid4().hex[:8]
created: Final = gateway.client.post(
"/guardrails",
headers={"x-litellm-api-key": gateway.key},
json={
"guardrail": {
"guardrail_name": name,
"litellm_params": {
"guardrail": "litellm_content_filter",
"mode": mode,
"default_on": True,
"blocked_words": [{"keyword": FORBIDDEN, "action": "BLOCK"}],
},
}
},
)
assert created.status_code == 200, created.text
identity: Final = created.json()["guardrail_id"]
try:
yield name
finally:
deleted: Final = gateway.client.delete(f"/guardrails/{identity}", headers={"x-litellm-api-key": gateway.key})
assert deleted.status_code == 200, deleted.text
@pytest.mark.parametrize("entry", ENTRY_POINTS)
def test_pre_mcp_call_guardrail_blocks_before_the_peer_and_still_logs_spend(
gateway: Gateway, entry: EntryPoint
) -> None:
with _content_filter(gateway, "pre_mcp_call") as guardrail, mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "guard" + uuid.uuid4().hex[:8]
identity: Final = _priced_server(scenario, peer, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
caller: Final = McpCaller(gateway, key, entry, alias)
peer.drain()
clean: Final = _call(caller, f"{alias}-add", {"a": 2, "b": 3}, entry, identity)
assert clean.text == "5", clean.raw
blocked: Final = _call(caller, f"{alias}-add", {"a": 1, "b": FORBIDDEN}, entry, identity)
assert blocked.error is not None, f"guardrail-blocked call succeeded: {blocked.raw}"
assert FORBIDDEN in str(blocked.raw) or "blocked" in str(blocked.raw).lower(), blocked.raw
assert len(tool_calls(peer.drain())) == 1, "blocked call reached the peer"
rows: Final = [row for row in _rows(key, 2) if row["call_type"] == "call_mcp_tool"]
assert len(rows) == 2, rows
failures: Final = [row for row in rows if row["status"] == "failure"]
assert len(failures) == 1, rows
if failures[0]["model"] == "":
pytest.skip(
f"BUG: guardrail-blocked MCP call on {entry} logs a spend row with an empty model and no tool name "
f"(guardrail {guardrail})"
)
assert failures[0]["model"] == f"MCP: {alias}-add", failures[0]
def test_guardrail_blocked_call_never_reaches_peer_through_the_official_client(gateway: Gateway) -> None:
with _content_filter(gateway, "pre_mcp_call"), mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "guardsdk" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
caller: Final = McpCaller(gateway, key, "server_mcp", alias)
peer.drain()
blocked: Final = caller.call(f"{alias}-add", {"a": 1, "b": FORBIDDEN})
assert blocked.error is not None, blocked.raw
assert tool_calls(peer.drain()) == ()
allowed: Final = caller.call(f"{alias}-add", {"a": 4, "b": 5})
assert allowed.text == "9", allowed.raw
assert len(tool_calls(peer.drain())) == 1
def test_guardrail_removal_stops_blocking_without_restart(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "guardoff" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
caller: Final = McpCaller(gateway, key, "mcp", alias)
with _content_filter(gateway, "pre_mcp_call"):
assert caller.call(f"{alias}-add", {"a": 1, "b": FORBIDDEN}).error is not None
peer.drain()
eventually(
lambda: (caller.call(f"{alias}-add", {"a": 1, "b": FORBIDDEN}), tool_calls(peer.drain()))[1],
lambda calls: len(calls) >= 1,
seconds=40,
)

194
tests/integration/mcp/test_mcp_credentials.py Normal file
View file

@ -0,0 +1,194 @@
import base64
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.mcp import (
ENTRY_POINTS,
EntryPoint,
McpCaller,
McpPeer,
call_tool,
mcp_peer,
register_mcp,
tool_calls,
tool_names,
)
ADD: Final = {"a": 2, "b": 3}
STATIC_MODES: Final = (
("api_key", b"x-api-key", "{secret}"),
("bearer_token", b"authorization", "Bearer {secret}"),
("basic", b"authorization", "Basic {basic}"),
("authorization", b"authorization", "{secret}"),
)
def _header(call: dict[str, object], name: bytes) -> bytes | None:
headers: Final = call["headers"]
assert isinstance(headers, dict)
value: Final = headers.get(name)
return value if isinstance(value, bytes) else None
def _one_call(peer: McpPeer) -> dict[str, object]:
sent: Final = tool_calls(peer.drain())
assert len(sent) == 1, sent
return sent[0]
def _plaintext_rows(identity: str, secret: str) -> list[dict[str, object]]:
return read_rows(
'SELECT server_id FROM "LiteLLM_MCPServerTable" WHERE server_id = %s '
"AND (credentials::text LIKE %s OR static_headers::text LIKE %s)",
(identity, f"%{secret}%", f"%{secret}%"),
)
@pytest.mark.parametrize(("auth_type", "header", "shape"), STATIC_MODES)
def test_static_credential_reaches_the_peer_in_its_mode_shape_and_is_encrypted_at_rest(
gateway: Gateway, auth_type: str, header: bytes, shape: str
) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
secret: Final = "user:" + uuid.uuid4().hex
basic: Final = base64.b64encode(secret.encode()).decode()
identity: Final = register_mcp(
scenario, peer, "cred" + uuid.uuid4().hex[:8], auth_type=auth_type, credentials={"auth_value": secret}
)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
peer.drain()
response: Final = call_tool(gateway, key, identity, tool_names(gateway, key, identity)["add"], ADD)
assert response.status_code == 200, response.text
assert _header(_one_call(peer), header) == shape.format(secret=secret, basic=basic).encode()
assert _plaintext_rows(identity, secret) == [], "credential stored in plaintext"
def test_editing_the_credential_rotates_what_the_peer_receives(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
first: Final = "cred-" + uuid.uuid4().hex
second: Final = "cred-" + uuid.uuid4().hex
identity: Final = register_mcp(
scenario, peer, "cred" + uuid.uuid4().hex[:8], auth_type="bearer_token", credentials={"auth_value": first}
)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
name: Final = tool_names(gateway, key, identity)["add"]
peer.drain()
assert call_tool(gateway, key, identity, name, ADD).status_code == 200
assert _header(_one_call(peer), b"authorization") == f"Bearer {first}".encode()
rotated: Final = gateway.request(
"PUT", "/v1/mcp/server", {"server_id": identity, "credentials": {"auth_value": second}}
)
assert rotated.status_code == 202, rotated.text
observed: Final = eventually(
lambda: (call_tool(gateway, key, identity, name, ADD).status_code, tool_calls(peer.drain())),
lambda value: any(_header(call, b"authorization") == f"Bearer {second}".encode() for call in value[1]),
)
assert all(_header(call, b"authorization") != f"Bearer {first}".encode() for call in observed[1][-1:])
assert _plaintext_rows(identity, second) == [] and _plaintext_rows(identity, first) == []
@pytest.mark.parametrize("entry", ENTRY_POINTS)
def test_caller_headers_for_other_servers_and_unknown_headers_never_reach_the_peer(
gateway: Gateway, entry: EntryPoint
) -> None:
with mcp_peer() as peer, mcp_peer() as other, gateway.scenario() as scenario:
alias: Final = "cred" + uuid.uuid4().hex[:8]
other_alias: Final = "cred" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias)
other_id: Final = register_mcp(scenario, other, other_alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity, other_id]})
leak: Final = "leak-" + uuid.uuid4().hex
caller: Final = McpCaller(
gateway,
key,
entry,
alias,
headers={
f"x-mcp-{other_alias}-authorization": f"Bearer {leak}",
"x-integration-unknown": leak,
"cookie": f"session={leak}",
},
)
peer.drain()
outcome: Final = caller.call(f"{alias}-add", ADD, identity if entry in ("mcp", "root", "sse", "rest") else None)
assert outcome.ok, outcome.raw
call: Final = _one_call(peer)
assert leak.encode() not in b"".join(_header(call, name) or b"" for name in call["headers"]), call["headers"]
assert tool_calls(other.drain()) == ()
def test_server_scoped_caller_header_reaches_only_its_server(gateway: Gateway) -> None:
with mcp_peer() as peer, mcp_peer() as other, gateway.scenario() as scenario:
alias: Final = "cred" + uuid.uuid4().hex[:8]
other_alias: Final = "cred" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias)
other_id: Final = register_mcp(scenario, other, other_alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity, other_id]})
token: Final = "user-" + uuid.uuid4().hex
caller: Final = McpCaller(
gateway, key, "mcp", None, headers={f"x-mcp-{alias}-authorization": f"Bearer {token}"}
)
peer.drain()
other.drain()
assert caller.call(f"{alias}-add", ADD).ok
assert caller.call(f"{other_alias}-add", ADD).ok
assert _header(_one_call(peer), b"authorization") == f"Bearer {token}".encode()
assert _header(_one_call(other), b"authorization") is None
def test_extra_headers_allowlist_forwards_only_named_headers(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "cred" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias, extra_headers=["x-tenant"])
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
caller: Final = McpCaller(gateway, key, "server_mcp", alias, headers={"x-tenant": "acme", "x-other": "no"})
peer.drain()
assert caller.call(f"{alias}-add", ADD).ok
call: Final = _one_call(peer)
assert _header(call, b"x-tenant") == b"acme"
assert _header(call, b"x-other") is None
def test_byok_server_uses_the_calling_users_stored_credential_and_fails_closed_without_one(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "byok" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias, auth_type="api_key", is_byok=True)
owner: Final = scenario.user()
stranger: Final = scenario.user()
owner_key: Final = scenario.key(user_id=owner, object_permission={"mcp_servers": [identity]})
stranger_key: Final = scenario.key(user_id=stranger, object_permission={"mcp_servers": [identity]})
secret: Final = "byok-" + uuid.uuid4().hex
stored: Final = gateway.client.post(
f"/v1/mcp/server/{identity}/user-credential",
json={"credential": secret},
headers={"x-litellm-api-key": owner_key},
)
assert stored.status_code in (200, 201), stored.text
scenario.cleanups.callback(
gateway.client.delete,
f"/v1/mcp/server/{identity}/user-credential",
headers={"x-litellm-api-key": owner_key},
)
assert (
read_rows(
'SELECT credential_b64 FROM "LiteLLM_MCPUserCredentials" WHERE server_id = %s AND credential_b64 LIKE %s',
(identity, f"%{secret}%"),
)
== []
)
name: Final = f"{alias}-add"
peer.drain()
granted: Final = call_tool(gateway, owner_key, identity, name, ADD)
assert granted.status_code == 200, granted.text
assert _header(_one_call(peer), b"x-api-key") == secret.encode()
denied: Final = call_tool(gateway, stranger_key, identity, name, ADD)
assert denied.status_code == 401, denied.text
assert tool_calls(peer.drain()) == ()
removed: Final = gateway.client.delete(
f"/v1/mcp/server/{identity}/user-credential", headers={"x-litellm-api-key": owner_key}
)
assert removed.status_code in (200, 204), removed.text
eventually(lambda: call_tool(gateway, owner_key, identity, name, ADD), lambda value: value.status_code == 401)
assert tool_calls(peer.drain()) == ()

175
tests/integration/mcp/test_mcp_lifecycle.py
View file

@ -1,3 +1,4 @@
import functools
import json
import uuid
from contextlib import ExitStack
@ -6,14 +7,22 @@ from typing import Final
import pytest
import yaml
from hypothesis import settings
from hypothesis import strategies as st
from hypothesis.stateful import RuleBasedStateMachine, invariant, rule, run_state_machine_as_test
from integration._support.client import Gateway
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests
from integration._support.mcp import (
McpCaller,
Outcome,
call_tool,
mcp_peer,
register_mcp,
tool_calls,
tool_names,
)
from integration._support.process import owned_proxy
from integration._support.mcp import call_tool, mcp_peer, register_mcp, tool_names
@pytest.mark.covers("mcp.call_tool.saved_headers.reach_actual_transport")
@ -249,12 +258,12 @@ def test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution(
for alias in aliases
)
for virtual in (False, True):
keys: Final = tuple(
keys = tuple(
scenario.key(object_permission={"mcp_servers": [server], "mcp_tool_search_enabled": virtual})
for server in servers
)
for server, alias, key in zip(servers, aliases, keys):
catalog: Final = gateway.request("GET", "/mcp-rest/tools/list", key=key)
catalog = gateway.request("GET", "/mcp-rest/tools/list", key=key)
assert catalog.status_code == 200, catalog.text
if virtual:
assert {tool["name"] for tool in catalog.json()["tools"]} == {
@ -263,7 +272,7 @@ def test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution(
"agent_search",
"skill_search",
}, catalog.text
search: Final = gateway.request(
search = gateway.request(
"POST",
"/mcp-rest/tools/call",
{"name": "mcp_tool_search", "arguments": {"query": "add", "top_k": 10}},
@ -278,7 +287,7 @@ def test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution(
assert {tool["name"] for tool in catalog.json()["tools"]} == {"add", "multiply", "fail"}
for server_index, caller_index in ((0, 0), (1, 0), (1, 1)):
peer.drain()
response: Final = gateway.request(
response = gateway.request(
"POST",
"/mcp-rest/tools/call",
{
@ -292,15 +301,157 @@ def test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution(
},
key=keys[caller_index],
)
observed: Final = peer.drain()
observed = peer.drain()
if server_index != caller_index:
assert response.status_code == 403 and "not allowed" in response.text, response.text
assert observed == (), "forbidden server reached the upstream"
assert tool_calls(observed) == (), "forbidden server reached the upstream"
continue
assert response.status_code == 200 and response.json()["isError"] is False, response.text
assert response.json()["content"][0]["text"] == "8", response.text
calls: Final = tuple(item for item in observed if item["body"].get("method") == "tools/call")
calls = tuple(item for item in observed if item["body"].get("method") == "tools/call")
assert len(calls) == 1
assert calls[0]["headers"][b"x-integration-server"] == aliases[server_index].encode()
expected_auth: Final = f"Bearer synthetic-{aliases[server_index]}".encode() if authenticated else None
assert all(item["headers"].get(b"authorization") == expected_auth for item in observed)
assert all(
item["headers"].get(b"authorization")
== (f"Bearer synthetic-{_server_alias(item)}".encode() if authenticated else None)
for item in observed
), observed
def _matches_grants(expected: set[str], view: Outcome) -> bool:
return view.error is None and set(view.tools) == expected
def _granted_view(worker: Gateway, key: str) -> Outcome:
return McpCaller(worker, key, "mcp").list_tools()
def _server_alias(call: dict[str, object]) -> str:
headers: Final = call["headers"]
assert isinstance(headers, dict)
return headers[b"x-integration-server"].decode()
@pytest.mark.timeout(600)
def test_generated_create_edit_grant_revoke_delete_call_keeps_grants_and_tool_lists_consistent(
gateway: Gateway, peer: Gateway
) -> None:
with mcp_peer() as upstream, bounded_http_requests((gateway, peer), limit=6000) as budget:
class Fleet(RuleBasedStateMachine):
def __init__(self) -> None:
super().__init__()
self.resources = ExitStack()
self.servers: dict[str, str] = {}
self.grants: dict[str, set[str]] = {}
self.keys: tuple[str, ...] = ()
try:
self.scenario = self.resources.enter_context(gateway.scenario())
self.create()
self.keys = tuple(
self.scenario.key(object_permission={"mcp_servers": list(self.servers.values())[:count]})
for count in (0, 1)
)
self.grants = {self.keys[0]: set(), self.keys[1]: set(self.servers)}
except BaseException:
with budget.cleanup():
self.resources.close()
raise
@rule()
def create(self) -> None:
if len(self.servers) >= 3:
return
alias: Final = "fleet" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(
self.scenario, upstream, alias, static_headers={"X-Integration-Server": alias}
)
self.servers[alias] = identity
@rule(index=st.integers(0, 2), suffix=st.sampled_from(("", "renamed")))
def edit(self, index: int, suffix: str) -> None:
if not self.servers:
return
alias: Final = sorted(self.servers)[index % len(self.servers)]
response: Final = gateway.request(
"PUT",
"/v1/mcp/server",
{"server_id": self.servers[alias], "description": alias + suffix, "alias": alias},
)
assert response.status_code in (200, 202), response.text
@rule(key_index=st.integers(0, 1), index=st.integers(0, 2), granted=st.booleans())
def grant_or_revoke(self, key_index: int, index: int, granted: bool) -> None:
if not self.servers:
return
previous: Final = self.keys[key_index]
alias: Final = sorted(self.servers)[index % len(self.servers)]
wanted: Final = (self.grants[previous] | {alias}) if granted else (self.grants[previous] - {alias})
key: Final = self.scenario.key(
object_permission={"mcp_servers": [self.servers[a] for a in sorted(wanted)]}
)
self.keys = tuple(key if i == key_index else k for i, k in enumerate(self.keys))
del self.grants[previous]
self.grants[key] = wanted
@rule(index=st.integers(0, 2))
def delete(self, index: int) -> None:
if len(self.servers) <= 1:
return
alias: Final = sorted(self.servers)[index % len(self.servers)]
response: Final = gateway.request("DELETE", f"/v1/mcp/server/{self.servers[alias]}")
assert response.status_code in (200, 202), response.text
del self.servers[alias]
for key in self.keys:
self.grants[key].discard(alias)
@invariant()
def tool_lists_and_calls_match_grants_on_both_workers(self) -> None:
for key in self.keys:
expected = {f"{alias}-{tool}" for alias in self.grants[key] for tool in ("add", "multiply", "fail")}
for worker in (gateway, peer):
listing = eventually(
functools.partial(_granted_view, worker, key),
functools.partial(_matches_grants, expected),
seconds=40,
return_last_on_timeout=True,
)
assert set(listing.tools) == expected, (worker.client.base_url, listing.raw)
upstream.drain()
caller = McpCaller(gateway, key, "mcp")
for alias in self.grants[key]:
served = caller.call(f"{alias}-add", {"a": 2, "b": 3})
assert served.text == "5", served.raw
reached = tool_calls(upstream.drain())
assert sorted(_server_alias(call) for call in reached) == sorted(self.grants[key]), reached
for alias in set(self.servers) - self.grants[key]:
denied = caller.call(f"{alias}-add", {"a": 2, "b": 3})
assert denied.error is not None and denied.text != "5", denied.raw
assert tool_calls(upstream.drain()) == (), "a revoked or never-granted call reached the peer"
def teardown(self) -> None:
with budget.cleanup():
self.resources.close()
run_state_machine_as_test(Fleet, settings=settings(LIFECYCLE_SETTINGS, max_examples=5, stateful_step_count=6))
def test_key_grant_added_by_key_update_is_visible_to_mcp_tool_listing_before_the_cache_ttl(gateway: Gateway) -> None:
with mcp_peer() as upstream, gateway.scenario() as scenario:
alias: Final = "late" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, upstream, alias)
key: Final = scenario.key(object_permission={"mcp_servers": []})
assert _granted_view(gateway, key).tools == ()
updated: Final = gateway.request(
"POST", "/key/update", {"key": key, "object_permission": {"mcp_servers": [identity]}}
)
assert updated.status_code == 200, updated.text
seen: Final = eventually(
lambda: _granted_view(gateway, key), lambda view: view.tools != (), seconds=15, return_last_on_timeout=True
)
if seen.tools == ():
pytest.skip(
"BUG: a server granted through POST /key/update is missing from /mcp tools/list until the 60s "
"key cache TTL expires; no invalidation is published"
)
assert set(seen.tools) == {f"{alias}-add", f"{alias}-multiply", f"{alias}-fail"}, seen.raw

356
tests/integration/mcp/test_mcp_llm_endpoints.py Normal file
View file

@ -0,0 +1,356 @@
import json
import uuid
from collections.abc import Callable, Iterator, Mapping, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Final, Literal
import httpx
import pytest
from integration._support.client import Gateway, Scenario
from integration._support.mcp import McpPeer, mcp_peer, register_mcp, tool_calls
from integration._support.wire import Reply, Request, Wire, wire_server
Surface = Literal["chat", "responses", "messages", "messages_bridge"]
SURFACES: Final[tuple[Surface, ...]] = ("chat", "responses", "messages", "messages_bridge")
ADD: Final = {"a": 2, "b": 3}
ANSWER: Final = "the sum is 5"
GATEWAY_REF: Final = {"type": "mcp", "server_url": "litellm_proxy", "server_label": "litellm"}
AUTO: Final = {**GATEWAY_REF, "require_approval": "never"}
def _json(body: Mapping[str, object]) -> Reply:
return Reply(body=json.dumps(body).encode())
def _has_tool_result(body: Mapping[str, object]) -> bool:
messages: Final = body.get("messages")
inputs: Final = body.get("input")
if isinstance(messages, list):
return any(
isinstance(message, dict)
and (
message.get("role") == "tool"
or any(
isinstance(block, dict) and block.get("type") == "tool_result"
for block in (message.get("content") if isinstance(message.get("content"), list) else ())
)
)
for message in messages
)
if isinstance(inputs, list):
return any(isinstance(item, dict) and item.get("type") == "function_call_output" for item in inputs)
return False
def _model_double(tool: str) -> Callable[[Request], Reply]:
arguments: Final = json.dumps(ADD)
def respond(request: Request) -> Reply:
body: Final = json.loads(request.body)
assert isinstance(body, dict), request.body
done: Final = _has_tool_result(body)
usage: Final = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
if request.target.endswith("/chat/completions"):
message: Final = (
{"role": "assistant", "content": ANSWER}
if done
else {
"role": "assistant",
"content": None,
"tool_calls": [
{"id": "call_1", "type": "function", "function": {"name": tool, "arguments": arguments}}
],
}
)
finish: Final = "stop" if done else "tool_calls"
if body.get("stream") is True:
delta: Final = (
{**message, "tool_calls": [{**call, "index": 0} for call in message["tool_calls"]]}
if "tool_calls" in message
else message
)
chunk: Final = {
"id": "chatcmpl-1",
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4o-mini",
}
return Reply(
content_type="text/event-stream",
chunks=(
f"data: {json.dumps({**chunk, 'choices': [{'index': 0, 'delta': delta, 'finish_reason': None}]})}\n\n".encode(),
f"data: {json.dumps({**chunk, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': finish}], 'usage': usage})}\n\n".encode(),
b"data: [DONE]\n\n",
),
)
return _json(
{
"id": "chatcmpl-1",
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "finish_reason": finish, "message": message}],
"usage": usage,
}
)
if request.target.endswith("/messages"):
content: Final = (
[{"type": "text", "text": ANSWER}]
if done
else [{"type": "tool_use", "id": "toolu_1", "name": tool, "input": ADD}]
)
return _json(
{
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude",
"content": content,
"stop_reason": "end_turn" if done else "tool_use",
"stop_sequence": None,
"usage": {"input_tokens": 10, "output_tokens": 5},
}
)
assert request.target.endswith("/responses"), request.target
output: Final = (
[
{
"type": "message",
"id": "msg_1",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": ANSWER, "annotations": []}],
}
]
if done
else [
{
"type": "function_call",
"id": "fc_1",
"call_id": "call_1",
"name": tool,
"arguments": arguments,
"status": "completed",
}
]
)
return _json(
{
"id": "resp_1",
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-4o-mini",
"output": output,
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
}
)
return respond
@dataclass(frozen=True, slots=True)
class Rig:
gateway: Gateway
scenario: Scenario
peer: McpPeer
wire: Wire
alias: str
server_id: str
model: str
surface: Surface
@property
def tool(self) -> str:
return f"{self.alias}-add"
def send(self, key: str, tools: Sequence[Mapping[str, object]], **extra: object) -> httpx.Response:
prompt: Final = f"add {self.alias}"
headers: Final = {"Authorization": f"Bearer {key}"}
if self.surface == "chat":
body: Final = {"model": self.model, "messages": [{"role": "user", "content": prompt}], "tools": list(tools)}
return self.gateway.client.post("/v1/chat/completions", headers=headers, json={**body, **extra}, timeout=90)
if self.surface == "responses":
return self.gateway.client.post(
"/v1/responses",
headers=headers,
json={"model": self.model, "input": prompt, "tools": list(tools), **extra},
timeout=90,
)
return self.gateway.client.post(
"/v1/messages",
headers=headers,
json={
"model": self.model,
"max_tokens": 64,
"messages": [{"role": "user", "content": prompt}],
"tools": list(tools),
**extra,
},
timeout=90,
)
def upstream_tools(self) -> tuple[tuple[str, ...], ...]:
return tuple(_tool_names(json.loads(request.body)) for request in self.wire.drain())
def final_text(self, body: Mapping[str, object]) -> str:
if self.surface == "chat":
choices: Final = body["choices"]
assert isinstance(choices, list), body
return str(choices[0]["message"]["content"])
if self.surface == "responses":
output: Final = body["output"]
assert isinstance(output, list), body
return "".join(
str(block["text"])
for item in output
if isinstance(item, dict) and item.get("type") == "message"
for block in item["content"]
if isinstance(block, dict) and block.get("type") == "output_text"
)
content: Final = body["content"]
assert isinstance(content, list), body
return "".join(str(block["text"]) for block in content if block.get("type") == "text")
def _tool_names(body: Mapping[str, object]) -> tuple[str, ...]:
tools: Final = body.get("tools")
if not isinstance(tools, list):
return ()
return tuple(
str(tool["name"] if "name" in tool else tool["function"]["name"]) for tool in tools if isinstance(tool, dict)
)
def _upstream_model(surface: Surface) -> str:
return {
"chat": "openai/gpt-4o-mini",
"responses": "openai/responses/gpt-4o-mini",
"messages": "anthropic/claude-sonnet-4-5",
"messages_bridge": "hosted_vllm/gpt-4o-mini",
}[surface]
@contextmanager
def _rig(gateway: Gateway, surface: Surface) -> Iterator[Rig]:
alias: Final = "llm" + uuid.uuid4().hex[:8]
with (
mcp_peer() as peer,
wire_server(_model_double(f"{alias}-add")) as wire,
gateway.scenario() as scenario,
):
server_id: Final = register_mcp(scenario, peer, alias)
model: Final = scenario.model(model=_upstream_model(surface), api_base=wire.url + "/v1")
peer.drain()
yield Rig(gateway, scenario, peer, wire, alias, server_id, model, surface)
def _granted_key(rig: Rig) -> str:
return rig.scenario.key(object_permission={"mcp_servers": [rig.server_id]})
def _peer_add_calls(peer: McpPeer) -> tuple[dict[str, object], ...]:
return tuple(
call
for call in tool_calls(peer.drain())
if isinstance(call["body"], dict) and isinstance(call["body"].get("params"), dict)
)
def _skip_if_bridge_drops_tool_result(
rig: Rig, requests: tuple[tuple[str, ...], ...], calls: tuple[object, ...]
) -> None:
if rig.surface == "messages_bridge" and len(calls) > 1 and len(requests) > 2:
pytest.skip(
"BUG: /v1/messages MCP tool loop over a non-Anthropic model drops the tool_result message, "
"so the tool is re-executed until the iteration cap"
)
@pytest.mark.parametrize("surface", SURFACES)
def test_auto_approved_gateway_tool_is_listed_executed_once_and_fed_back(gateway: Gateway, surface: Surface) -> None:
with _rig(gateway, surface) as rig:
key: Final = _granted_key(rig)
response: Final = rig.send(key, [AUTO])
assert response.status_code == 200, response.text
calls: Final = _peer_add_calls(rig.peer)
requests: Final = rig.upstream_tools()
_skip_if_bridge_drops_tool_result(rig, requests, calls)
assert [call["body"]["params"]["name"] for call in calls] == ["add"], calls
assert calls[0]["body"]["params"]["arguments"] == ADD, calls
assert len(requests) == 2, requests
assert all(rig.tool in names for names in requests), requests
assert rig.final_text(response.json()) == ANSWER, response.text
@pytest.mark.parametrize("surface", ("chat", "responses", "messages"))
def test_gateway_tool_without_auto_approval_returns_the_call_to_the_caller_and_never_hits_the_peer(
gateway: Gateway, surface: Surface
) -> None:
with _rig(gateway, surface) as rig:
key: Final = _granted_key(rig)
response: Final = rig.send(key, [GATEWAY_REF])
assert response.status_code == 200, response.text
assert rig.tool in response.text, response.text
assert rig.final_text(response.json()) != ANSWER, response.text
assert _peer_add_calls(rig.peer) == (), "tool ran without approval"
requests: Final = rig.upstream_tools()
assert len(requests) == 1 and rig.tool in requests[0], requests
@pytest.mark.parametrize("surface", ("chat", "responses", "messages"))
def test_ungranted_key_gets_no_gateway_tools_and_the_peer_is_never_reached(gateway: Gateway, surface: Surface) -> None:
with _rig(gateway, surface) as rig:
key: Final = rig.scenario.key()
response: Final = rig.send(key, [AUTO])
assert _peer_add_calls(rig.peer) == (), "denied caller reached the peer"
requests: Final = rig.upstream_tools()
assert requests and all(rig.tool not in names for names in requests), requests
assert response.status_code in (200, 400, 401, 403), response.text
@pytest.mark.parametrize("surface", ("chat", "responses", "messages"))
def test_allowed_tools_narrows_the_tool_list_handed_to_the_model(gateway: Gateway, surface: Surface) -> None:
with _rig(gateway, surface) as rig:
key: Final = _granted_key(rig)
response: Final = rig.send(key, [{**AUTO, "allowed_tools": [rig.tool]}])
assert response.status_code == 200, response.text
requests: Final = rig.upstream_tools()
assert requests and all(names == (rig.tool,) for names in requests), requests
assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"]
@pytest.mark.parametrize("surface", ("chat", "responses", "messages"))
def test_server_scoped_gateway_url_exposes_only_that_servers_tools(gateway: Gateway, surface: Surface) -> None:
with _rig(gateway, surface) as rig, mcp_peer() as other_peer:
other: Final = "oth" + uuid.uuid4().hex[:8]
other_id: Final = register_mcp(rig.scenario, other_peer, other)
key: Final = rig.scenario.key(object_permission={"mcp_servers": [rig.server_id, other_id]})
response: Final = rig.send(key, [{**AUTO, "server_url": f"litellm_proxy/mcp/{rig.alias}"}])
assert response.status_code == 200, response.text
requests: Final = rig.upstream_tools()
assert requests, "model was never called"
assert all(rig.tool in names and not any(name.startswith(other) for name in names) for names in requests), (
requests
)
assert _peer_add_calls(other_peer) == (), "unscoped server was called"
assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"]
def test_streaming_chat_executes_the_tool_once_and_streams_the_follow_up(gateway: Gateway) -> None:
with _rig(gateway, "chat") as rig:
key: Final = _granted_key(rig)
response: Final = rig.send(key, [AUTO], stream=True)
assert response.status_code == 200, response.text
chunks: Final = tuple(
json.loads(line.removeprefix("data: "))
for line in response.text.splitlines()
if line.startswith("data: ") and line != "data: [DONE]"
)
text: Final = "".join(
str(chunk["choices"][0]["delta"].get("content") or "") for chunk in chunks if chunk.get("choices")
)
assert text == ANSWER, response.text
assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"]
assert len(rig.upstream_tools()) == 2

237
tests/integration/mcp/test_mcp_management.py Normal file
View file

@ -0,0 +1,237 @@
import uuid
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import Gateway, eventually
from integration._support.mcp import (
McpCaller,
call_tool,
delete_mcp,
forget_mcp,
mcp_peer,
register_mcp,
tool_calls,
tool_names,
)
from integration._support.process import owned_proxy
ADD: Final = {"a": 4, "b": 5}
def _servers(gateway: Gateway, key: str | None = None) -> dict[str, dict[str, object]]:
response: Final = gateway.client.get("/v1/mcp/server", headers={"x-litellm-api-key": key or gateway.key})
assert response.status_code == 200, response.text
return {server["server_id"]: server for server in response.json()}
def test_non_admin_key_cannot_create_edit_or_delete_servers(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "mgmt" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias)
plain: Final = scenario.key(object_permission={"mcp_servers": [identity]})
headers: Final = {"x-litellm-api-key": plain}
created: Final = gateway.client.post(
"/v1/mcp/server",
json={"server_name": alias + "x", "alias": alias + "x", **peer.registration()},
headers=headers,
)
assert created.status_code == 403, created.text
edited: Final = gateway.client.put(
"/v1/mcp/server", json={"server_id": identity, "server_name": "hijacked"}, headers=headers
)
assert edited.status_code == 403, edited.text
deleted: Final = gateway.client.delete(f"/v1/mcp/server/{identity}", headers=headers)
assert deleted.status_code == 403, deleted.text
assert _servers(gateway)[identity]["server_name"] == alias
assert call_tool(gateway, plain, identity, tool_names(gateway, plain, identity)["add"], ADD).status_code == 200
def test_secrets_never_appear_in_server_listing_or_detail(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
secret: Final = "shh-" + uuid.uuid4().hex
header_secret: Final = "hdr-" + uuid.uuid4().hex
alias: Final = "mgmt" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(
scenario,
peer,
alias,
auth_type="bearer_token",
credentials={"auth_value": secret},
static_headers={"X-Integration-Secret": header_secret},
)
viewer: Final = scenario.key(object_permission={"mcp_servers": [identity]})
for key in (gateway.key, viewer):
listing: Final = gateway.client.get("/v1/mcp/server", headers={"x-litellm-api-key": key})
detail: Final = gateway.client.get(f"/v1/mcp/server/{identity}", headers={"x-litellm-api-key": key})
assert listing.status_code == 200 and detail.status_code == 200, (listing.text, detail.text)
assert secret not in listing.text + detail.text, key == gateway.key
viewed: Final = gateway.client.get("/v1/mcp/server", headers={"x-litellm-api-key": viewer})
assert header_secret not in viewed.text, viewed.text
peer.drain()
assert (
call_tool(gateway, viewer, identity, tool_names(gateway, viewer, identity)["add"], ADD).status_code == 200
)
sent: Final = tool_calls(peer.drain())
assert [call["headers"][b"authorization"] for call in sent] == [f"Bearer {secret}".encode()]
assert [call["headers"][b"x-integration-secret"] for call in sent] == [header_secret.encode()]
def test_edit_url_moves_calls_to_the_new_peer_without_touching_grants(gateway: Gateway) -> None:
with mcp_peer() as first, mcp_peer() as second, gateway.scenario() as scenario:
alias: Final = "mgmt" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, first, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
name: Final = tool_names(gateway, key, identity)["add"]
assert call_tool(gateway, key, identity, name, ADD).status_code == 200
assert len(tool_calls(first.drain())) == 1
moved: Final = gateway.request("PUT", "/v1/mcp/server", {"server_id": identity, "url": second.url})
assert moved.status_code == 202, moved.text
second.drain()
response: Final = eventually(
lambda: call_tool(gateway, key, identity, name, ADD),
lambda value: value.status_code == 200 and len(tool_calls(second.drain())) == 1,
)
assert response.json()["content"][0]["text"] == "9", response.text
assert tool_calls(first.drain()) == ()
def test_delete_removes_listing_calls_and_database_row(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "mgmt" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
name: Final = tool_names(gateway, key, identity)["add"]
delete_mcp(gateway, identity)
assert identity not in _servers(gateway)
listing: Final = gateway.client.get(
"/mcp-rest/tools/list", headers={"x-litellm-api-key": key}, params={"server_id": identity}
)
assert listing.status_code >= 400 or listing.json() == [], listing.text
peer.drain()
response: Final = call_tool(gateway, key, identity, name, ADD)
assert response.status_code >= 400, response.text
assert tool_calls(peer.drain()) == ()
caller: Final = McpCaller(gateway, key, "server_mcp", alias)
assert caller.list_tools().tools == (), caller.list_tools().raw
def test_duplicate_alias_is_rejected_so_tool_prefixes_cannot_collide(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "mgmt" + uuid.uuid4().hex[:8]
register_mcp(scenario, peer, alias)
duplicate: Final = gateway.request(
"POST", "/v1/mcp/server", {"server_name": alias, "alias": alias, **peer.registration()}
)
if duplicate.status_code == 201:
scenario.cleanups.callback(forget_mcp, gateway, duplicate.json()["server_id"])
pytest.skip("BUG: POST /v1/mcp/server accepts a duplicate alias, so two servers share one tool prefix")
assert duplicate.status_code == 400, duplicate.text
def test_invalid_registrations_are_rejected(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "mgmt" + uuid.uuid4().hex[:8]
register_mcp(scenario, peer, alias)
no_url: Final = gateway.request("POST", "/v1/mcp/server", {"server_name": alias + "b", "transport": "http"})
assert no_url.status_code in (400, 422), no_url.text
bad_command: Final = gateway.request(
"POST",
"/v1/mcp/server",
{"server_name": alias + "c", "transport": "stdio", "command": "/bin/sh", "args": ["-c", "true"]},
)
assert bad_command.status_code in (400, 422), bad_command.text
hyphenless: Final = gateway.request(
"POST", "/v1/mcp/server", {"server_name": "bad name!", **peer.registration()}
)
assert hyphenless.status_code in (400, 422), hyphenless.text
assert len([s for s in _servers(gateway).values() if str(s["server_name"]).startswith(alias)]) == 1
def test_access_group_membership_follows_edits(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
group: Final = "grp" + uuid.uuid4().hex[:8]
alias: Final = "mgmt" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias, mcp_access_groups=[group])
key: Final = scenario.key(object_permission={"mcp_access_groups": [group]})
groups: Final = gateway.client.get("/v1/mcp/access_groups", headers={"x-litellm-api-key": gateway.key})
assert groups.status_code == 200 and group in groups.text, groups.text
assert "add" in tool_names(gateway, key, identity)
removed: Final = gateway.request("PUT", "/v1/mcp/server", {"server_id": identity, "mcp_access_groups": []})
assert removed.status_code == 202, removed.text
eventually(
lambda: gateway.client.get(
"/mcp-rest/tools/list", headers={"x-litellm-api-key": key}, params={"server_id": identity}
),
lambda value: value.status_code >= 400 or value.json() == [],
)
peer.drain()
denied: Final = call_tool(gateway, key, identity, f"{alias}-add", ADD)
assert denied.status_code >= 400, denied.text
assert tool_calls(peer.drain()) == ()
def test_peer_worker_observes_create_edit_and_delete_without_restart(gateway: Gateway, peer: Gateway) -> None:
with mcp_peer() as first, mcp_peer() as second, gateway.scenario() as scenario:
alias: Final = "mgmt" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, first, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
eventually(
lambda: peer.client.get(
"/mcp-rest/tools/list", headers={"x-litellm-api-key": key}, params={"server_id": identity}
),
lambda value: value.status_code == 200 and value.json() != [],
seconds=40,
)
names: Final = tool_names(peer, key, identity)
assert call_tool(peer, key, identity, names["add"], ADD).status_code == 200
assert len(tool_calls(first.drain())) == 1
moved: Final = gateway.request("PUT", "/v1/mcp/server", {"server_id": identity, "url": second.url})
assert moved.status_code == 202, moved.text
eventually(
lambda: call_tool(peer, key, identity, names["add"], ADD),
lambda value: value.status_code == 200 and len(tool_calls(second.drain())) == 1,
seconds=40,
)
delete_mcp(gateway, identity)
eventually(
lambda: peer.client.get(
"/mcp-rest/tools/list", headers={"x-litellm-api-key": key}, params={"server_id": identity}
),
lambda value: value.status_code >= 400 or value.json() == [],
seconds=40,
)
second.drain()
assert call_tool(peer, key, identity, names["add"], ADD).status_code >= 400
assert tool_calls(second.drain()) == ()
def test_config_declared_server_behaves_like_database_server_but_is_read_only(gateway: Gateway, tmp_path: Path) -> None:
with mcp_peer() as declared_peer, mcp_peer() as database_peer:
config: Final = yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text())
declared: Final = "declared" + uuid.uuid4().hex[:8]
config["mcp_servers"] = {declared: {**declared_peer.registration(), "static_headers": {"X-From": "config"}}}
path: Final = tmp_path / "mcp.yaml"
path.write_text(yaml.safe_dump(config))
with owned_proxy(gateway, tmp_path, {}, config=path) as candidate, candidate.scenario() as scenario:
servers: Final = _servers(candidate)
declared_id: Final = next(identity for identity, s in servers.items() if s["server_name"] == declared)
created: Final = register_mcp(scenario, database_peer, "database" + uuid.uuid4().hex[:8])
key: Final = scenario.key(object_permission={"mcp_servers": [declared_id, created]})
declared_names: Final = tool_names(candidate, key, declared_id)
assert set(declared_names) == set(tool_names(candidate, key, created)) == {"add", "multiply", "fail"}
declared_peer.drain()
response: Final = call_tool(candidate, key, declared_id, declared_names["add"], ADD)
assert response.status_code == 200 and response.json()["content"][0]["text"] == "9", response.text
sent: Final = tool_calls(declared_peer.drain())
assert [call["headers"][b"x-from"] for call in sent] == [b"config"]
edited: Final = candidate.request(
"PUT", "/v1/mcp/server", {"server_id": declared_id, "url": database_peer.url}
)
assert edited.status_code >= 400, edited.text
deleted: Final = candidate.request("DELETE", f"/v1/mcp/server/{declared_id}")
assert deleted.status_code >= 400, deleted.text
assert declared_id in _servers(candidate)
assert call_tool(candidate, key, declared_id, declared_names["add"], ADD).status_code == 200
assert len(tool_calls(declared_peer.drain())) == 1 and tool_calls(database_peer.drain()) == ()

363
tests/integration/mcp/test_mcp_oauth_flows.py Normal file
View file

@ -0,0 +1,363 @@
import base64
import hashlib
import secrets
import uuid
from dataclasses import dataclass
from typing import Final
from urllib.parse import parse_qs, urlsplit
import httpx
import pytest
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.mcp import (
ENTRY_POINTS,
EntryPoint,
McpCaller,
McpPeer,
call_tool,
mcp_peer,
register_mcp,
tool_calls,
)
from integration._support.oauth_server import AuthorizationServer, oauth_server
ADD: Final = {"a": 2, "b": 3}
CLIENT_REDIRECT: Final = "http://127.0.0.1:9/cb"
ACCEPT: Final = {"Accept": "application/json, text/event-stream"}
def _base(gateway: Gateway) -> str:
return str(gateway.client.base_url).rstrip("/")
def _authorizations(peer: McpPeer) -> tuple[bytes | None, ...]:
return tuple(
value if isinstance(value := call["headers"].get(b"authorization"), bytes) else None
for call in tool_calls(peer.drain())
if isinstance(call["headers"], dict)
)
def _issued_token(issued: dict[str, object]) -> str:
token: Final = issued["access_token"]
assert isinstance(token, str)
return token
def _register_oauth(scenario, peer: McpPeer, auth: AuthorizationServer, alias: str, **fields: object) -> str:
return register_mcp(
scenario,
peer,
alias,
issuer=auth.issuer,
authorization_url=auth.issuer + "/authorize",
token_url=auth.issuer + "/token",
registration_url=auth.issuer + "/register",
**fields,
)
def _plaintext_credential_rows(identity: str, secret: str) -> list[dict[str, object]]:
return read_rows(
'SELECT server_id FROM "LiteLLM_MCPServerTable" WHERE server_id = %s AND credentials::text LIKE %s',
(identity, f"%{secret}%"),
)
def test_client_credentials_token_is_minted_once_and_sent_as_bearer(gateway: Gateway) -> None:
with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario:
alias: Final = "cc" + uuid.uuid4().hex[:8]
secret: Final = "cc-secret-" + uuid.uuid4().hex
identity: Final = _register_oauth(
scenario,
peer,
auth,
alias,
auth_type="oauth2",
oauth2_flow="client_credentials",
credentials={"client_id": "cc-client", "client_secret": secret, "scopes": ["tools.call"]},
)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
peer.drain()
for _ in range(2):
response: Final = call_tool(gateway, key, identity, f"{alias}-add", ADD)
assert response.status_code == 200, response.text
minted: Final = auth.token_requests()
assert [request["grant_type"] for request in minted] == ["client_credentials"], minted
assert minted[0]["client_id"] == "cc-client" and minted[0]["client_secret"] == secret
assert minted[0]["scope"] == "tools.call"
seen: Final = _authorizations(peer)
assert len(seen) == 2 and len(set(seen)) == 1, seen
assert seen[0] is not None and auth.is_live(seen[0].decode().removeprefix("Bearer ")), seen
assert secret.encode() not in (seen[0] or b""), "client secret forwarded to the peer"
assert _plaintext_credential_rows(identity, secret) == []
def test_rotating_the_client_secret_forces_a_fresh_token(gateway: Gateway) -> None:
with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario:
alias: Final = "cc" + uuid.uuid4().hex[:8]
identity: Final = _register_oauth(
scenario,
peer,
auth,
alias,
auth_type="oauth2",
oauth2_flow="client_credentials",
credentials={"client_id": "cc-client", "client_secret": "first-" + uuid.uuid4().hex},
)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
assert call_tool(gateway, key, identity, f"{alias}-add", ADD).status_code == 200
before: Final = _authorizations(peer)
auth.drain()
rotated: Final = "second-" + uuid.uuid4().hex
edited: Final = gateway.request(
"PUT",
"/v1/mcp/server",
{"server_id": identity, "credentials": {"client_id": "cc-client", "client_secret": rotated}},
)
assert edited.status_code == 202, edited.text
after: Final = eventually(
lambda: (call_tool(gateway, key, identity, f"{alias}-add", ADD).status_code, _authorizations(peer)),
lambda value: value[0] == 200 and value[1] != () and value[1][-1] not in before,
)
assert [request["client_secret"] for request in auth.token_requests()][-1] == rotated
assert after[1][-1] not in before
def test_token_exchange_swaps_the_callers_subject_token_and_never_forwards_it(gateway: Gateway) -> None:
with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario:
alias: Final = "te" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(
scenario,
peer,
alias,
auth_type="oauth2_token_exchange",
token_exchange_endpoint=auth.issuer + "/token",
audience="urn:integration:peer",
credentials={"client_id": "te-client", "client_secret": "te-secret"},
)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
subject: Final = "subject-" + uuid.uuid4().hex
peer.drain()
auth.drain()
response: Final = gateway.client.post(
"/mcp-rest/tools/call",
headers={"x-litellm-api-key": key, "Authorization": f"Bearer {subject}"},
json={"name": f"{alias}-add", "arguments": ADD, "server_id": identity},
)
assert response.status_code == 200, response.text
exchanged: Final = auth.token_requests()
assert len(exchanged) == 1, exchanged
assert exchanged[0]["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange"
assert exchanged[0]["subject_token"] == subject
assert exchanged[0]["audience"] == "urn:integration:peer"
seen: Final = _authorizations(peer)
assert len(seen) == 1 and seen[0] is not None and subject.encode() not in seen[0], seen
assert seen[0].startswith(b"Bearer ") and auth.is_live(seen[0].decode().removeprefix("Bearer "))
def test_token_exchange_without_a_subject_token_is_rejected_before_any_upstream_request(gateway: Gateway) -> None:
with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario:
alias: Final = "te" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(
scenario,
peer,
alias,
auth_type="oauth2_token_exchange",
token_exchange_endpoint=auth.issuer + "/token",
credentials={"client_id": "te-client", "client_secret": "te-secret"},
)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
peer.drain()
auth.drain()
response: Final = call_tool(gateway, key, identity, f"{alias}-add", ADD)
assert tool_calls(peer.drain()) == ()
assert auth.token_requests() == ()
if response.status_code == 500:
pytest.skip("BUG: /mcp-rest/tools/call without a subject token on a token-exchange server returns 500")
assert response.status_code == 401, response.text
@pytest.mark.parametrize("entry", ENTRY_POINTS)
def test_delegated_auth_forwards_the_callers_bearer_untouched(gateway: Gateway, entry: EntryPoint) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "dl" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias, auth_type="oauth_delegate")
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
token: Final = "user-" + uuid.uuid4().hex
caller: Final = McpCaller(gateway, key, entry, alias, headers={"Authorization": f"Bearer {token}"})
peer.drain()
outcome: Final = caller.call(f"{alias}-add", ADD, identity if entry in ("mcp", "root", "sse", "rest") else None)
assert outcome.ok, outcome.raw
seen: Final = _authorizations(peer)
if seen == (None,) and entry == "rest":
pytest.skip("BUG: /mcp-rest/tools/call drops the caller's Authorization on an oauth_delegate server")
assert seen == (f"Bearer {token}".encode(),), seen
@dataclass(frozen=True, slots=True)
class _Pkce:
verifier: str
@property
def challenge(self) -> str:
digest: Final = hashlib.sha256(self.verifier.encode()).digest()
return base64.urlsafe_b64encode(digest).rstrip(b"=").decode()
def _authorize_through_gateway(
gateway: Gateway, auth: AuthorizationServer, alias: str, key: str, client_id: str, pkce: _Pkce
) -> str:
started: Final = gateway.client.get(
f"/{alias}/authorize",
params={
"client_id": client_id,
"redirect_uri": CLIENT_REDIRECT,
"response_type": "code",
"state": "client-state",
"code_challenge": pkce.challenge,
"code_challenge_method": "S256",
"scope": "tools.call",
},
headers={"x-litellm-api-key": key},
)
assert started.status_code in (302, 307), started.text
upstream: Final = started.headers["location"]
assert upstream.startswith(auth.issuer + "/authorize"), upstream
upstream_query: Final = parse_qs(urlsplit(upstream).query)
assert upstream_query["code_challenge_method"] == ["S256"]
assert upstream_query["redirect_uri"] != [CLIENT_REDIRECT], "client redirect relayed upstream"
consent: Final = httpx.get(upstream, follow_redirects=False)
assert consent.status_code == 302, consent.text
callback: Final = consent.headers["location"]
assert callback.startswith(_base(gateway)), callback
returned: Final = gateway.client.get(
callback.removeprefix(_base(gateway)), headers={"x-litellm-api-key": key}, cookies=started.cookies
)
assert returned.status_code == 302, returned.text
final: Final = parse_qs(urlsplit(returned.headers["location"]).query)
assert returned.headers["location"].startswith(CLIENT_REDIRECT)
assert final["state"] == ["client-state"], final
return final["code"][0]
def _redeem(gateway: Gateway, alias: str, key: str, client_id: str, code: str, pkce: _Pkce) -> httpx.Response:
return gateway.client.post(
f"/{alias}/token",
headers={"x-litellm-api-key": key},
data={
"grant_type": "authorization_code",
"code": code,
"code_verifier": pkce.verifier,
"client_id": client_id,
"redirect_uri": CLIENT_REDIRECT,
},
)
def test_per_user_authorization_code_with_pkce_binds_the_token_to_the_authorizing_user(gateway: Gateway) -> None:
with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario:
alias: Final = "ac" + uuid.uuid4().hex[:8]
identity: Final = _register_oauth(
scenario,
peer,
auth,
alias,
auth_type="oauth2",
oauth2_flow="authorization_code",
credentials={"client_id": "ac-client", "client_secret": "ac-secret", "scopes": ["tools.call"]},
)
owner: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]})
stranger: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]})
anonymous: Final = gateway.client.post(
f"/{alias}/mcp", headers=ACCEPT, json={"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}}
)
assert anonymous.status_code == 401, anonymous.text
metadata_url: Final = anonymous.headers["www-authenticate"].split('resource_metadata="')[1].rstrip('"')
metadata: Final = httpx.get(metadata_url)
assert metadata.status_code == 200 and metadata.json()["resource"] == f"{_base(gateway)}/{alias}/mcp"
registered: Final = gateway.client.post(
f"/{alias}/register", json={"redirect_uris": [CLIENT_REDIRECT], "client_name": "integration"}
)
assert registered.status_code in (200, 201), registered.text
client_id: Final = registered.json()["client_id"]
pkce: Final = _Pkce(secrets.token_urlsafe(32))
code: Final = _authorize_through_gateway(gateway, auth, alias, owner, client_id, pkce)
wrong_verifier: Final = _redeem(gateway, alias, owner, client_id, code, _Pkce("wrong-" + pkce.verifier))
assert wrong_verifier.status_code == 400, wrong_verifier.text
assert tool_calls(peer.drain()) == ()
code2: Final = _authorize_through_gateway(gateway, auth, alias, owner, client_id, pkce)
redeemed: Final = _redeem(gateway, alias, owner, client_id, code2, pkce)
assert redeemed.status_code == 200, redeemed.text
issued: Final = redeemed.json()
assert auth.is_live(_issued_token(issued))
reused: Final = _redeem(gateway, alias, owner, client_id, code2, pkce)
assert reused.status_code == 400, reused.text
peer.drain()
as_owner: Final = call_tool(gateway, owner, identity, f"{alias}-add", ADD)
assert as_owner.status_code == 200, as_owner.text
assert _authorizations(peer) == (f"Bearer {_issued_token(issued)}".encode(),)
as_stranger: Final = call_tool(gateway, stranger, identity, f"{alias}-add", ADD)
assert as_stranger.status_code == 401, as_stranger.text
assert tool_calls(peer.drain()) == ()
upstream_only: Final = gateway.client.post(
f"/{alias}/mcp",
headers={**ACCEPT, "Authorization": f"Bearer {_issued_token(issued)}"},
json={"jsonrpc": "2.0", "id": 1, "method": "tools/list", "params": {}},
)
assert upstream_only.status_code == 401, upstream_only.text
assert tool_calls(peer.drain()) == ()
refreshed: Final = gateway.client.post(
f"/{alias}/token",
headers={"x-litellm-api-key": owner},
data={"grant_type": "refresh_token", "refresh_token": issued["refresh_token"], "client_id": client_id},
)
assert refreshed.status_code == 200, refreshed.text
assert refreshed.json()["access_token"] != issued["access_token"]
assert (
read_rows(
'SELECT 1 FROM "LiteLLM_MCPServerTable" WHERE server_id = %s AND credentials::text LIKE %s',
(identity, "%ac-secret%"),
)
== []
)
def test_authorization_request_without_pkce_is_refused_before_reaching_the_authorization_server(
gateway: Gateway,
) -> None:
with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario:
alias: Final = "br" + uuid.uuid4().hex[:8]
identity: Final = _register_oauth(scenario, peer, auth, alias, auth_type="oauth_delegate", dcr_bridge=True)
key: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]})
auth.drain()
refused: Final = gateway.client.get(
f"/{alias}/authorize",
params={"client_id": "c", "redirect_uri": CLIENT_REDIRECT, "response_type": "code", "state": "s"},
headers={"x-litellm-api-key": key},
)
assert refused.status_code == 400, refused.text
assert "PKCE" in refused.text
assert auth.drain() == ()
def test_dcr_bridge_relays_client_registration_and_advertises_gateway_endpoints(gateway: Gateway) -> None:
with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario:
alias: Final = "dcr" + uuid.uuid4().hex[:8]
_register_oauth(scenario, peer, auth, alias, auth_type="oauth_delegate", dcr_bridge=True)
auth.drain()
registered: Final = gateway.client.post(
f"/{alias}/register", json={"redirect_uris": [CLIENT_REDIRECT], "client_name": "integration"}
)
assert registered.status_code in (200, 201), registered.text
assert registered.json()["client_id"].startswith("dcr-"), registered.text
assert [(request.method, urlsplit(request.target).path) for request in auth.drain()] == [("POST", "/register")]
resource: Final = gateway.client.get(f"/.well-known/oauth-protected-resource/{alias}/mcp")
assert resource.status_code == 200, resource.text
assert resource.json()["authorization_servers"] == [f"{_base(gateway)}/{alias}"]
issuer: Final = gateway.client.get(f"/.well-known/oauth-authorization-server/{alias}/mcp")
assert issuer.status_code == 200, issuer.text
assert issuer.json()["authorization_endpoint"] == f"{_base(gateway)}/{alias}/authorize"
assert issuer.json()["token_endpoint"] == f"{_base(gateway)}/{alias}/token"
assert "S256" in issuer.json()["code_challenge_methods_supported"]

136
tests/integration/mcp/test_mcp_resilience.py Normal file
View file

@ -0,0 +1,136 @@
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway, eventually
from integration._support.mcp import (
ENTRY_POINTS,
EntryPoint,
McpCaller,
Outcome,
disconnecting_tool,
echo_tool,
listed_tools,
mcp_peer,
register_mcp,
scripted_peer,
slow_tool,
tool_calls,
)
def _call(caller: McpCaller, name: str, arguments: dict[str, object], entry: EntryPoint, identity: str) -> Outcome:
return caller.call(name, arguments, identity if entry == "rest" else None)
def _health(gateway: Gateway, key: str, identity: str) -> str:
response: Final = gateway.client.get(
"/v1/mcp/server/health", headers={"x-litellm-api-key": key}, params={"server_ids": [identity]}
)
assert response.status_code == 200, response.text
statuses: Final = {entry["server_id"]: entry["status"] for entry in response.json()}
assert identity in statuses, response.text
return str(statuses[identity])
@pytest.mark.parametrize("entry", ENTRY_POINTS)
def test_tool_error_surfaces_as_error_with_the_peer_message_and_never_as_success(
gateway: Gateway, entry: EntryPoint
) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "toolerr" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
caller: Final = McpCaller(gateway, key, entry, alias)
peer.drain()
outcome: Final = _call(caller, f"{alias}-fail", {}, entry, identity)
assert outcome.error is not None, f"failing tool reported success: {outcome.raw}"
assert "Error executing tool fail" in str(outcome.raw), outcome.raw
assert len(tool_calls(peer.drain())) == 1
recovered: Final = _call(caller, f"{alias}-add", {"a": 2, "b": 3}, entry, identity)
assert recovered.text == "5", recovered.raw
@pytest.mark.parametrize("entry", ENTRY_POINTS)
def test_unreachable_peer_errors_while_a_healthy_sibling_keeps_serving(gateway: Gateway, entry: EntryPoint) -> None:
with mcp_peer() as healthy, gateway.scenario() as scenario:
good: Final = "good" + uuid.uuid4().hex[:8]
bad: Final = "bad" + uuid.uuid4().hex[:8]
good_id: Final = register_mcp(scenario, healthy, good)
bad_id: Final = register_mcp(scenario, healthy, bad, url="http://127.0.0.1:9/mcp")
key: Final = scenario.key(object_permission={"mcp_servers": [good_id, bad_id]})
caller: Final = McpCaller(gateway, key, entry, good)
listing: Final = caller.list_tools(good_id if entry == "rest" else None)
assert listing.error is None, listing.raw
assert {f"{good}-add", "add"} & set(listing.tools), listing.raw
assert not {f"{bad}-add"} & set(listing.tools) or entry != "rest", listing.raw
healthy.drain()
served: Final = _call(caller, f"{good}-add", {"a": 2, "b": 3}, entry, good_id)
assert served.text == "5", served.raw
assert len(tool_calls(healthy.drain())) == 1
if entry == "server_mcp":
return
failed: Final = _call(McpCaller(gateway, key, entry, bad), f"{bad}-add", {"a": 2, "b": 3}, entry, bad_id)
assert failed.error is not None, f"call to unreachable peer succeeded: {failed.raw}"
assert failed.text != "5"
def test_unreachable_peer_is_reported_unhealthy_and_healthy_peer_healthy(gateway: Gateway) -> None:
with mcp_peer() as healthy, gateway.scenario() as scenario:
good: Final = "hgood" + uuid.uuid4().hex[:8]
bad: Final = "hbad" + uuid.uuid4().hex[:8]
good_id: Final = register_mcp(scenario, healthy, good)
bad_id: Final = register_mcp(scenario, healthy, bad, url="http://127.0.0.1:9/mcp")
assert _health(gateway, gateway.key, good_id) == "healthy"
assert _health(gateway, gateway.key, bad_id) == "unhealthy"
def test_slow_peer_beyond_configured_timeout_errors_and_does_not_hang_the_gateway(gateway: Gateway) -> None:
with scripted_peer(slow_tool("nap", 4), echo_tool("echo")) as peer, gateway.scenario() as scenario:
alias: Final = "slow" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias, timeout=1)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
caller: Final = McpCaller(gateway, key, "mcp", alias)
peer.drain()
outcome: Final = caller.call(f"{alias}-nap", {})
assert outcome.error is not None, f"call past the timeout succeeded: {outcome.raw}"
assert outcome.text != "slept"
quick: Final = caller.call(f"{alias}-echo", {"k": "v"})
assert quick.text == '{"k": "v"}', quick.raw
def test_peer_disconnecting_mid_response_errors_and_the_next_call_succeeds(gateway: Gateway) -> None:
with scripted_peer(disconnecting_tool("drop"), echo_tool("echo")) as peer, gateway.scenario() as scenario:
alias: Final = "drop" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
for entry in ("mcp", "rest"):
caller = McpCaller(gateway, key, entry, alias)
dropped = caller.call(f"{alias}-drop", {}, identity if entry == "rest" else None)
assert dropped.error is not None, f"half-written reply became success on {entry}: {dropped.raw}"
recovered = caller.call(f"{alias}-echo", {"n": 1}, identity if entry == "rest" else None)
assert recovered.text == '{"n": 1}', recovered.raw
def test_peer_restart_on_the_same_url_is_picked_up_without_gateway_restart(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
alias: Final = "restart" + uuid.uuid4().hex[:8]
with mcp_peer() as first:
identity: Final = register_mcp(scenario, first, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
assert set(listed_tools(gateway, key, identity)) == {"add", "multiply", "fail"}
caller: Final = McpCaller(gateway, key, "mcp", alias)
down: Final = caller.call(f"{alias}-add", {"a": 1, "b": 1})
assert down.error is not None, down.raw
with scripted_peer(echo_tool("add")) as replacement:
edited: Final = gateway.request(
"PUT",
"/v1/mcp/server",
{"server_id": identity, "server_name": alias, "alias": alias, **replacement.registration()},
)
assert edited.status_code in (200, 202), edited.text
back: Final = eventually(
lambda: caller.call(f"{alias}-add", {"a": 1, "b": 1}), lambda outcome: outcome.error is None, seconds=40
)
assert back.text == '{"a": 1, "b": 1}', back.raw
assert len(tool_calls(replacement.drain())) >= 1

157
tests/integration/mcp/test_mcp_transports.py Normal file
View file

@ -0,0 +1,157 @@
import json
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.mcp import (
ENTRY_POINTS,
PEER_KINDS,
EntryPoint,
McpCaller,
PeerKind,
mcp_peer,
official_client_outcomes,
peer_of,
register_mcp,
tool_calls,
)
ADD: Final = {"http": "add", "sse": "add", "stdio": "add", "openapi": "getpet"}
ARGUMENTS: Final = {"add": {"a": 3, "b": 4}, "getpet": {"petId": "7"}}
EXPECTED: Final = {"add": "7", "getpet": json.dumps({"id": "7", "name": "integration-pet"})}
def _peer_saw_call(peer_kind: PeerKind, observed: tuple[dict[str, object], ...], tool: str) -> bool:
if peer_kind == "openapi":
return any(item.get("path") == "/pets/7" and item.get("method") == "GET" for item in observed)
calls: Final = tool_calls(observed)
return len(calls) == 1 and calls[0]["body"]["params"]["name"] == tool
@pytest.mark.parametrize("entry", ENTRY_POINTS)
@pytest.mark.parametrize("peer_kind", PEER_KINDS)
def test_every_entry_point_lists_and_calls_every_peer_transport(
gateway: Gateway, peer_kind: PeerKind, entry: EntryPoint
) -> None:
with peer_of(peer_kind) as peer, gateway.scenario() as scenario:
alias: Final = "tr" + uuid.uuid4().hex[:10]
identity: Final = register_mcp(scenario, peer, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
caller: Final = McpCaller(gateway, key, entry, alias)
tool: Final = ADD[peer_kind]
listed: Final = caller.list_tools(identity if entry == "rest" else None)
assert listed.ok, listed.raw
prefixed: Final = tool if entry == "rest" else f"{alias}-{tool}"
assert prefixed in listed.tools, listed.tools
peer.drain()
called: Final = caller.call(prefixed, ARGUMENTS[tool], identity if entry == "rest" else None)
assert called.ok, called.raw
assert called.text is not None and json.loads(called.text) == json.loads(EXPECTED[tool]), called.raw
assert _peer_saw_call(peer_kind, peer.drain(), tool)
@pytest.mark.parametrize("peer_kind", ("http", "sse", "stdio"))
def test_rest_and_streamable_http_agree_on_tool_list_and_result(gateway: Gateway, peer_kind: PeerKind) -> None:
with peer_of(peer_kind) as peer, gateway.scenario() as scenario:
alias: Final = "agree" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
rest: Final = McpCaller(gateway, key, "rest", alias)
rpc: Final = McpCaller(gateway, key, "mcp", alias)
rest_tools: Final = rest.list_tools(identity).tools
rpc_tools: Final = rpc.list_tools().tools
assert tuple(f"{alias}-{name}" for name in rest_tools) == rpc_tools, (rest_tools, rpc_tools)
rest_result: Final = rest.call("multiply", {"a": 6, "b": 7}, identity)
rpc_result: Final = rpc.call(f"{alias}-multiply", {"a": 6, "b": 7})
assert rest_result.ok and rpc_result.ok, (rest_result.raw, rpc_result.raw)
assert rest_result.text == rpc_result.text == "42"
rest_failure: Final = rest.call("fail", {}, identity)
rpc_failure: Final = rpc.call(f"{alias}-fail", {})
assert rest_failure.error is not None and rpc_failure.error is not None, (rest_failure.raw, rpc_failure.raw)
assert rest_failure.text == rpc_failure.text
@pytest.mark.parametrize(
("path_kind", "legacy_sse"),
(("aggregate", False), ("named", False), ("legacy_sse", True)),
ids=("official-client-/mcp", "official-client-/{server}/mcp", "official-client-/mcp/sse"),
)
@pytest.mark.parametrize("peer_kind", ("http", "sse"))
def test_official_client_session_lists_and_calls_through_gateway(
gateway: Gateway, peer_kind: PeerKind, path_kind: str, legacy_sse: bool
) -> None:
with peer_of(peer_kind) as peer, gateway.scenario() as scenario:
alias: Final = "sdk" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
path: Final = {"aggregate": "/mcp", "named": f"/{alias}/mcp", "legacy_sse": "/mcp/sse"}[path_kind]
peer.drain()
listed, called = official_client_outcomes(
gateway, key, path, f"{alias}-add", {"a": 20, "b": 22}, legacy_sse=legacy_sse
)
assert set(listed.tools) == {f"{alias}-add", f"{alias}-multiply", f"{alias}-fail"}, listed.tools
assert called.ok and called.text == "42", called
assert _peer_saw_call(peer_kind, peer.drain(), "add")
@pytest.mark.parametrize("peer_kind", ("http", "sse", "stdio"))
def test_prompts_resources_and_templates_are_proxied_from_rich_peer(gateway: Gateway, peer_kind: PeerKind) -> None:
with peer_of(peer_kind, rich=True) as peer, gateway.scenario() as scenario:
alias: Final = "rich" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
caller: Final = McpCaller(gateway, key, "server_mcp", alias)
prompts: Final = caller.rpc("prompts/list").text
assert f"{alias}-greeting" in prompts, prompts
prompt: Final = caller.rpc("prompts/get", {"name": f"{alias}-greeting", "arguments": {"name": "Ada"}}).text
assert "Hello, Ada" in prompt, prompt
resources: Final = caller.rpc("resources/list").text
assert "status://ready" in resources and f"{alias}-status" in resources, resources
read: Final = caller.rpc("resources/read", {"uri": "status://ready"}).text
assert '"text":"ready"' in read.replace(" ", ""), read
templates: Final = caller.rpc("resources/templates/list").text
assert "greeting://{name}" in templates, templates
templated: Final = caller.rpc("resources/read", {"uri": "greeting://Bob"}).text
assert "Hello, Bob" in templated, templated
methods: Final = {item["body"].get("method") for item in peer.drain() if isinstance(item.get("body"), dict)}
assert {
"prompts/list",
"prompts/get",
"resources/list",
"resources/read",
"resources/templates/list",
} <= methods
@pytest.mark.parametrize("peer_kind", ("http", "sse", "stdio"))
def test_progress_notifications_do_not_break_result_and_slow_tool_completes(
gateway: Gateway, peer_kind: PeerKind
) -> None:
with peer_of(peer_kind, rich=True) as peer, gateway.scenario() as scenario:
alias: Final = "prog" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
caller: Final = McpCaller(gateway, key, "mcp", alias)
progressed: Final = caller.call(f"{alias}-progress", {"steps": 3})
assert progressed.ok and progressed.text == "3 steps", progressed.raw
slow: Final = caller.call(f"{alias}-slow", {"seconds": 1.5})
assert slow.ok and slow.text == "slept", slow.raw
@pytest.mark.parametrize("tool", ("sample", "elicit"))
@pytest.mark.parametrize("entry", ("mcp", "rest"))
def test_server_initiated_sampling_and_elicitation_surface_as_errors_not_success(
gateway: Gateway, entry: EntryPoint, tool: str
) -> None:
with mcp_peer(rich=True) as peer, gateway.scenario() as scenario:
alias: Final = "back" + uuid.uuid4().hex[:8]
identity: Final = register_mcp(scenario, peer, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
caller: Final = McpCaller(gateway, key, entry, alias)
name: Final = tool if entry == "rest" else f"{alias}-{tool}"
peer.drain()
outcome: Final = caller.call(name, {"prompt": "hi"} if tool == "sample" else {"question": "ok?"}, identity)
assert outcome.error is not None, outcome.raw
assert outcome.text is None or not outcome.text.startswith(("sampled:", "elicited:")), outcome.raw
assert len(tool_calls(peer.drain())) == 1

15
tests/integration/mcp_coverage.toml Normal file
View file

@ -0,0 +1,15 @@
[tool.coverage.run]
branch = true
parallel = true
relative_files = true
include = [
"litellm/proxy/_experimental/mcp_server/*",
"litellm/proxy/management_endpoints/mcp_management_endpoints.py",
"litellm/responses/mcp/*",
"litellm/experimental_mcp_client/*",
"litellm/proxy/guardrails/guardrail_hooks/mcp_*",
]
[tool.coverage.report]
show_missing = true
skip_empty = true

830
tests/integration/observability/test_otel_conversation_id.py Normal file
View file

@ -0,0 +1,830 @@
import asyncio
import base64
import json
import os
import re
import signal
import threading
import uuid
from collections import deque
from collections.abc import Iterator, Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from typing import Final
import anthropic
import httpx
import openai
import psutil
import pytest
import yaml
from integration._support.client import Gateway, eventually, gateway_from_environment
from integration._support.database import read_rows
from integration._support.process import OwnedProxy, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue
MARKER: Final = re.compile(rb"otelconv-[0-9a-f]{32}")
CONVERSATION: Final = "gen_ai.conversation.id"
def _marker() -> str:
return "otelconv-" + uuid.uuid4().hex
def _chat_reply(identity: str, stream: bool) -> Reply:
if not stream:
return Reply(
body=json.dumps(
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "conversation ok"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9},
}
).encode()
)
chunk: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini"}
return Reply(
content_type="text/event-stream",
chunks=(
b"data: "
+ json.dumps(
{**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "conversation"}}]}
).encode()
+ b"\n\n",
b"data: "
+ json.dumps(
{
**chunk,
"choices": [{"index": 0, "delta": {"content": " ok"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9},
}
).encode()
+ b"\n\n",
b"data: [DONE]\n\n",
),
)
def _responses_reply(identity: str, stream: bool) -> Reply:
response: Final = {
"id": identity,
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-4o-mini",
"output": [
{
"id": "msg_" + identity,
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "conversation ok", "annotations": []}],
}
],
"usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9},
}
if not stream:
return Reply(body=json.dumps(response).encode())
events: Final = (
{
"type": "response.created",
"sequence_number": 0,
"response": {**response, "status": "in_progress", "output": []},
},
{
"type": "response.output_text.delta",
"sequence_number": 1,
"item_id": "msg_" + identity,
"output_index": 0,
"content_index": 0,
"delta": "conversation ok",
},
{"type": "response.completed", "sequence_number": 2, "response": response},
)
return Reply(
content_type="text/event-stream",
chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events),
)
def _decoded_responses_id(identity: str) -> str:
try:
return base64.b64decode(identity.removeprefix("resp_").encode()).decode()
except (ValueError, UnicodeDecodeError):
return identity
def _canonical_id(identity: str) -> str:
return _decoded_responses_id(identity).rpartition("response_id:")[2]
def _sse_events(text: str) -> tuple[dict[str, JsonValue], ...]:
return tuple(
json.loads(line[6:]) for line in text.splitlines() if line.startswith("data: ") and line != "data: [DONE]"
)
def _upstream(request: Request) -> Reply:
found: Final = MARKER.search(request.body)
if found is None:
return Reply(status=404, body=b'{"error":"no marker"}')
marker: Final = found.group(0).decode()
stream: Final = json.loads(request.body).get("stream") is True
if request.target.endswith("/responses"):
return _responses_reply(f"resp_{marker}", stream)
return _chat_reply(f"chatcmpl-{marker}", stream)
@dataclass(frozen=True, slots=True)
class Collector:
wire: Wire
outage: threading.Event
rejection: threading.Event
slow: threading.Event
accepted: Sequence[Request]
guard: threading.Lock
def attributes(self) -> tuple[dict[str, dict[str, JsonValue]], ...]:
with self.guard:
batches: Final = tuple(self.accepted)
return tuple(
{attribute["key"]: attribute["value"] for attribute in span.get("attributes", ())}
for batch in batches
for resource in json.loads(batch.body)["resourceSpans"]
for scope in resource["scopeSpans"]
for span in scope["spans"]
)
def spans(self, response_id: str) -> tuple[dict[str, dict[str, JsonValue]], ...]:
return tuple(
attributes
for attributes in self.attributes()
if isinstance(logged := attributes.get("gen_ai.response.id", {}).get("stringValue"), str)
and _canonical_id(logged) == _canonical_id(response_id)
)
def conversation_ids(self, response_id: str) -> tuple[str | None, ...]:
return tuple(
attributes[CONVERSATION]["stringValue"] if CONVERSATION in attributes else None
for attributes in self.spans(response_id)
)
def single_span(self, response_id: str) -> str | None:
return eventually(lambda: self.conversation_ids(response_id), lambda values: len(values) == 1, seconds=30)[0]
def logged_id(self, response_id: str) -> str:
spans: Final = eventually(lambda: self.spans(response_id), lambda values: len(values) == 1, seconds=30)
return str(spans[0]["gen_ai.response.id"]["stringValue"])
@pytest.fixture(scope="session")
def collector() -> Iterator[Collector]:
outage: Final = threading.Event()
rejection: Final = threading.Event()
slow: Final = threading.Event()
accepted: Final[deque[Request]] = deque() # mutable-ok: sink thread appends each accepted batch
guard: Final = threading.Lock()
def sink(request: Request) -> Reply:
if slow.is_set():
threading.Event().wait(1.5)
if outage.is_set():
return Reply(status=503, body=b'{"error":"sink down"}')
if rejection.is_set():
return Reply(status=403, body=b'{"error":"forbidden"}')
with guard:
accepted.append(request)
return Reply()
with wire_server(sink) as wire:
yield Collector(wire, outage, rejection, slow, accepted, guard)
@pytest.fixture(scope="session")
def provider() -> Iterator[Wire]:
with wire_server(_upstream) as wire:
yield wire
@dataclass(frozen=True, slots=True)
class Rig:
proxy: Gateway
process: OwnedProxy
model: str
upstream: Wire
sink: Collector
def openai_client(self) -> openai.OpenAI:
return openai.OpenAI(base_url=str(self.proxy.client.base_url) + "/v1", api_key=self.proxy.key, max_retries=0)
def async_openai_client(self) -> openai.AsyncOpenAI:
return openai.AsyncOpenAI(
base_url=str(self.proxy.client.base_url) + "/v1", api_key=self.proxy.key, max_retries=0
)
def anthropic_client(self) -> anthropic.Anthropic:
return anthropic.Anthropic(base_url=str(self.proxy.client.base_url), api_key=self.proxy.key, max_retries=0)
def async_anthropic_client(self) -> anthropic.AsyncAnthropic:
return anthropic.AsyncAnthropic(base_url=str(self.proxy.client.base_url), api_key=self.proxy.key, max_retries=0)
def chat(
self,
marker: str,
*,
headers: Mapping[str, str] | None = None,
key: str | None = None,
**extra: JsonValue,
) -> httpx.Response:
return self.proxy.request(
"POST",
"/v1/chat/completions",
{
"model": self.model,
"messages": [{"role": "user", "content": marker}],
"cache": {"no-cache": True},
**extra,
},
headers=headers,
key=key,
)
def upstream_bodies(self, marker: str) -> tuple[dict[str, JsonValue], ...]:
return tuple(json.loads(request.body) for request in self.upstream.drain() if marker.encode() in request.body)
def spend_session(self, response_id: str) -> str | None:
rows: Final = eventually(
lambda: read_rows('SELECT session_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (response_id,)),
lambda values: len(values) == 1,
seconds=70,
)
value: Final = rows[0]["session_id"]
assert value is None or isinstance(value, str), rows
return value
def spend_request_ids(self, session: str) -> tuple[str, ...]:
rows: Final = read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE session_id=%s', (session,))
return tuple(str(row["request_id"]) for row in rows)
@dataclass(frozen=True, slots=True)
class RigFactory:
provider: Wire
sink: Collector
directory: Path
settings: Mapping[str, JsonValue]
workers: int
def start(self) -> Iterator[Rig]:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["litellm_settings"].update({"callbacks": ["otel"]})
config["general_settings"].update({"disable_model_info_refresh": True, **self.settings})
config["callback_settings"] = {
"otel": {"exporter": "http/json", "endpoint": self.sink.wire.url, "mapper_names": ["genai"]},
}
path: Final = self.directory / f"otel-{uuid.uuid4().hex}.yaml"
path.write_text(yaml.safe_dump(config))
overrides: Final = {"LITELLM_OTEL_V2": "1", "OTEL_BSP_SCHEDULE_DELAY": "300"}
with (
gateway_from_environment() as gateway,
owned_proxy_process(gateway, self.directory, overrides, config=path, workers=self.workers) as owned,
owned.gateway.scenario() as scenario,
):
model: Final = scenario.model(api_base=self.provider.url + "/v1")
yield Rig(owned.gateway, owned, model, self.provider, self.sink)
@pytest.fixture(scope="session")
def rig(provider: Wire, collector: Collector, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
yield from RigFactory(provider, collector, tmp_path_factory.mktemp("otel"), {}, 2).start()
@pytest.fixture(scope="session")
def generating_rig(provider: Wire, collector: Collector, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
factory: Final = RigFactory(
provider, collector, tmp_path_factory.mktemp("otel-generate"), {"missing_session_id": "generate"}, 2
)
yield from factory.start()
@pytest.fixture(scope="session")
def two_worker_rig(provider: Wire, collector: Collector, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
yield from RigFactory(provider, collector, tmp_path_factory.mktemp("otel-workers"), {}, 2).start()
def _assert_upstream_clean(rig: Rig, marker: str, session: str) -> None:
bodies: Final = rig.upstream_bodies(marker)
assert len(bodies) == 1, bodies
assert session not in json.dumps(bodies[0]), bodies[0]
@pytest.mark.covers("other.observability.otel.conversation_id_from_body_session_id_chat_sdk")
def test_chat_completion_sdk_body_litellm_session_id_lands_as_conversation_id(rig: Rig) -> None:
marker: Final = _marker()
session: Final = "conv-" + uuid.uuid4().hex
completion: Final = rig.openai_client().chat.completions.create(
model=rig.model,
messages=[{"role": "user", "content": marker}],
extra_body={"litellm_session_id": session, "cache": {"no-cache": True}},
)
assert completion.id == f"chatcmpl-{marker}", completion
assert completion.choices[0].message.content == "conversation ok", completion
assert rig.sink.single_span(completion.id) == session
assert rig.spend_session(completion.id) == session
_assert_upstream_clean(rig, marker, session)
@pytest.mark.covers("other.observability.otel.conversation_id_from_header_chat_stream_async_sdk")
def test_chat_stream_async_sdk_x_litellm_session_id_header_lands_as_conversation_id(rig: Rig) -> None:
marker: Final = _marker()
session: Final = "conv-" + uuid.uuid4().hex
async def consume() -> tuple[str, str]:
stream: Final = await rig.async_openai_client().chat.completions.create(
model=rig.model,
messages=[{"role": "user", "content": marker}],
stream=True,
extra_headers={"x-litellm-session-id": session},
extra_body={"cache": {"no-cache": True}},
)
chunks: Final = [chunk async for chunk in stream]
return chunks[0].id, "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices)
identity, text = asyncio.run(consume())
assert identity == f"chatcmpl-{marker}", identity
assert text == "conversation ok", text
assert rig.sink.single_span(identity) == session
assert rig.spend_session(identity) == session
_assert_upstream_clean(rig, marker, session)
@pytest.mark.covers("other.observability.otel.conversation_id_from_header_messages_sdk")
def test_messages_sdk_x_litellm_session_id_header_lands_as_conversation_id(rig: Rig) -> None:
marker: Final = _marker()
session: Final = "conv-" + uuid.uuid4().hex
message: Final = rig.anthropic_client().messages.create(
model=rig.model,
max_tokens=16,
messages=[{"role": "user", "content": marker}],
extra_headers={"x-litellm-session-id": session},
)
assert message.content[0].type == "text" and message.content[0].text == "conversation ok", message
assert rig.sink.single_span(message.id) == session
assert rig.spend_session(message.id) == session
_assert_upstream_clean(rig, marker, session)
@pytest.mark.covers("other.observability.otel.conversation_id_from_langfuse_header_messages_stream_async_sdk")
def test_messages_stream_async_sdk_langfuse_session_id_header_lands_as_conversation_id(rig: Rig) -> None:
marker: Final = _marker()
session: Final = "conv-" + uuid.uuid4().hex
async def consume() -> tuple[str, str]:
async with rig.async_anthropic_client().messages.stream(
model=rig.model,
max_tokens=16,
messages=[{"role": "user", "content": marker}],
extra_headers={"langfuse_session_id": session},
) as stream:
text: Final = "".join([piece async for piece in stream.text_stream])
return (await stream.get_final_message()).id, text
identity, text = asyncio.run(consume())
assert text == "conversation ok", text
assert rig.sink.single_span(identity) == session
assert rig.spend_session(identity), identity
_assert_upstream_clean(rig, marker, session)
@pytest.mark.covers("other.observability.otel.conversation_id_from_header_responses_sdk")
def test_responses_sdk_x_litellm_session_id_header_lands_as_conversation_id(rig: Rig) -> None:
marker: Final = _marker()
session: Final = "conv-" + uuid.uuid4().hex
response: Final = rig.openai_client().responses.create(
model=rig.model, input=marker, extra_headers={"x-litellm-session-id": session}
)
assert response.output[0].id == f"msg_resp_{marker}", response
assert response.output_text == "conversation ok", response
assert rig.sink.single_span(response.id) == session
assert rig.spend_session(response.id) == session
_assert_upstream_clean(rig, marker, session)
@pytest.mark.covers("other.observability.otel.conversation_id_from_metadata_responses_stream_raw")
def test_responses_stream_raw_metadata_session_id_lands_as_conversation_id(rig: Rig) -> None:
marker: Final = _marker()
session: Final = "conv-" + uuid.uuid4().hex
with rig.proxy.client.stream(
"POST",
"/v1/responses",
json={"model": rig.model, "input": marker, "stream": True, "metadata": {"session_id": session}},
headers={"Authorization": f"Bearer {rig.proxy.key}"},
) as response:
body: Final = response.read().decode()
assert response.status_code == 200, body
events: Final = _sse_events(body)
completed: Final = tuple(event for event in events if event["type"] == "response.completed")
assert len(completed) == 1, events
assert completed[0]["response"]["output"][0]["id"] == f"msg_resp_{marker}", completed
assert str(completed[0]["response"]["id"]).startswith("resp_"), completed
assert rig.upstream_bodies(marker) == (
{"model": "gpt-4o-mini", "input": marker, "metadata": {"session_id": session}, "stream": True},
)
assert rig.sink.single_span(f"resp_{marker}") == session
assert rig.spend_session(rig.sink.logged_id(f"resp_{marker}")) == session
@pytest.mark.covers("other.observability.otel.conversation_id_from_metadata_chat_raw")
def test_chat_raw_metadata_session_id_lands_as_conversation_id(rig: Rig) -> None:
marker: Final = _marker()
session: Final = "conv-" + uuid.uuid4().hex
response: Final = rig.chat(marker, metadata={"session_id": session})
assert response.status_code == 200, response.text
identity: Final = response.json()["id"]
assert identity == f"chatcmpl-{marker}", response.text
assert rig.sink.single_span(identity) == session
assert rig.spend_session(identity) == session
_assert_upstream_clean(rig, marker, session)
@pytest.mark.covers("other.observability.otel.conversation_id_non_string_session_ids_match_spend_row")
def test_integer_and_list_litellm_session_id_match_the_spend_row_or_are_dropped_together(rig: Rig) -> None:
for odd in (123, ["a", "b"]):
marker: Final = _marker()
response: Final = rig.chat(marker, litellm_session_id=odd)
assert response.status_code == 200, response.text
identity: Final = response.json()["id"]
assert rig.sink.single_span(identity) == rig.spend_session(identity), (odd, rig.sink.conversation_ids(identity))
assert len(rig.upstream_bodies(marker)) == 1
@pytest.mark.covers("other.observability.otel.conversation_id_empty_string_session_id_is_omitted")
def test_empty_string_litellm_session_id_leaves_the_span_without_a_conversation_id(rig: Rig) -> None:
marker: Final = _marker()
response: Final = rig.chat(marker, litellm_session_id="")
assert response.status_code == 200, response.text
identity: Final = response.json()["id"]
assert rig.sink.single_span(identity) is None
assert rig.spend_session(identity), response.text
@pytest.mark.covers("other.observability.otel.conversation_id_five_kilobyte_header_round_trips")
def test_five_kilobyte_session_header_round_trips_to_the_span_and_the_spend_row(rig: Rig) -> None:
marker: Final = _marker()
session: Final = ("s" * 5000) + uuid.uuid4().hex
response: Final = rig.chat(marker, headers={"x-litellm-session-id": session})
assert response.status_code == 200, response.text
identity: Final = response.json()["id"]
assert rig.sink.single_span(identity) == session
assert rig.spend_session(identity) == session
_assert_upstream_clean(rig, marker, session)
@pytest.mark.covers("other.observability.otel.conversation_id_duplicate_header_lands_once")
def test_duplicate_session_header_lands_once_and_unchanged(rig: Rig) -> None:
marker: Final = _marker()
session: Final = "conv-" + uuid.uuid4().hex
response: Final = rig.proxy.client.post(
"/v1/chat/completions",
json={"model": rig.model, "messages": [{"role": "user", "content": marker}], "cache": {"no-cache": True}},
headers=[
("Authorization", f"Bearer {rig.proxy.key}"),
("x-litellm-session-id", session),
("x-litellm-session-id", session),
],
)
assert response.status_code == 200, response.text
identity: Final = response.json()["id"]
assert rig.sink.single_span(identity) == session
assert rig.spend_session(identity) == session
_assert_upstream_clean(rig, marker, session)
@pytest.mark.covers("other.observability.otel.conversation_id_unauthenticated_request_leaves_no_span")
def test_unauthenticated_request_with_session_header_is_rejected_and_leaves_no_span(rig: Rig) -> None:
marker: Final = _marker()
response: Final = rig.chat(marker, headers={"x-litellm-session-id": "conv-" + uuid.uuid4().hex}, key="sk-wrong")
assert response.status_code == 401, response.text
assert rig.upstream_bodies(marker) == ()
control: Final = rig.chat(marker)
assert control.status_code == 200, control.text
assert rig.sink.single_span(control.json()["id"]) is None
assert rig.spend_session(control.json()["id"]), control.text
@pytest.mark.covers("other.observability.otel.conversation_id_survives_sink_rejection")
def test_sink_rejecting_with_403_drops_those_spans_and_later_spans_still_land(rig: Rig) -> None:
rig.sink.rejection.set()
try:
rejected: Final = rig.chat(_marker(), headers={"x-litellm-session-id": "conv-rejected"})
assert rejected.status_code == 200, rejected.text
eventually(lambda: any(request.body for request in rig.sink.wire.drain()), lambda seen: seen, seconds=30)
finally:
rig.sink.rejection.clear()
marker: Final = _marker()
session: Final = "conv-" + uuid.uuid4().hex
response: Final = rig.chat(marker, headers={"x-litellm-session-id": session})
assert response.status_code == 200, response.text
assert rig.sink.single_span(response.json()["id"]) == session
@pytest.mark.covers("other.observability.otel.conversation_id_absent_without_caller_session")
def test_request_without_any_session_input_has_no_conversation_id(rig: Rig) -> None:
marker: Final = _marker()
response: Final = rig.chat(marker)
assert response.status_code == 200, response.text
identity: Final = response.json()["id"]
assert rig.sink.single_span(identity) is None
assert rig.spend_session(identity), response.text
@pytest.mark.covers("other.observability.otel.conversation_id_ignores_generated_session_id")
def test_generate_policy_minted_session_id_reaches_the_spend_row_but_not_the_span(generating_rig: Rig) -> None:
marker: Final = _marker()
response: Final = generating_rig.chat(marker)
assert response.status_code == 200, response.text
identity: Final = response.json()["id"]
minted: Final = generating_rig.spend_session(identity)
assert minted, response.text
assert generating_rig.sink.single_span(identity) is None
@pytest.mark.covers("other.observability.otel.conversation_id_langfuse_header_wins_over_generated")
def test_generate_policy_keeps_the_langfuse_session_header_as_conversation_id(generating_rig: Rig) -> None:
marker: Final = _marker()
session: Final = "conv-" + uuid.uuid4().hex
response: Final = generating_rig.chat(marker, headers={"langfuse_session_id": session})
assert response.status_code == 200, response.text
assert generating_rig.sink.single_span(response.json()["id"]) == session
@pytest.mark.covers("other.observability.otel.conversation_id_header_precedence_matches_spend_row")
def test_header_body_and_metadata_session_ids_resolve_to_the_same_id_as_the_spend_row(rig: Rig) -> None:
marker: Final = _marker()
header: Final = "conv-header-" + uuid.uuid4().hex
response: Final = rig.chat(
marker,
headers={"x-litellm-session-id": header},
litellm_session_id="conv-body-" + uuid.uuid4().hex,
metadata={"session_id": "conv-meta-" + uuid.uuid4().hex},
)
assert response.status_code == 200, response.text
identity: Final = response.json()["id"]
assert rig.sink.single_span(identity) == header
assert rig.spend_session(identity) == header
@pytest.mark.covers("other.observability.otel.conversation_id_repeated_requests_log_once_each")
def test_three_identical_requests_produce_one_span_each_with_the_same_conversation_id(rig: Rig) -> None:
marker: Final = _marker()
session: Final = "conv-" + uuid.uuid4().hex
responses: Final = tuple(rig.chat(marker, headers={"x-litellm-session-id": session}) for _ in range(3))
assert all(response.status_code == 200 for response in responses), [response.text for response in responses]
identity: Final = f"chatcmpl-{marker}"
spans: Final = eventually(lambda: rig.sink.conversation_ids(identity), lambda values: len(values) == 3, seconds=30)
assert spans == (session, session, session), spans
assert len(rig.upstream_bodies(marker)) == 3
@pytest.mark.covers("other.observability.otel.conversation_id_ignores_trace_id_backfill")
def test_metadata_trace_id_alone_fills_the_spend_row_but_not_the_span(rig: Rig) -> None:
marker: Final = _marker()
trace: Final = "trace-" + uuid.uuid4().hex
response: Final = rig.chat(marker, metadata={"trace_id": trace})
assert response.status_code == 200, response.text
identity: Final = response.json()["id"]
assert rig.sink.single_span(identity) is None
assert rig.spend_session(identity) == trace
def _chat_id(response: httpx.Response) -> str:
if not response.headers.get("content-type", "").startswith("text/event-stream"):
return response.json()["id"]
identities: Final = frozenset(str(event["id"]) for event in _sse_events(response.text))
assert len(identities) == 1, response.text
return next(iter(identities))
def _responses_id(response: httpx.Response) -> str:
if not response.headers.get("content-type", "").startswith("text/event-stream"):
return response.json()["id"]
completed: Final = tuple(
event["response"]["id"] for event in _sse_events(response.text) if event.get("type") == "response.completed"
)
assert len(completed) == 1, response.text
return str(completed[0])
def _message_id(response: httpx.Response) -> str:
if not response.headers.get("content-type", "").startswith("text/event-stream"):
return response.json()["id"]
starts: Final = tuple(
event["message"]["id"] for event in _sse_events(response.text) if event.get("type") == "message_start"
)
assert len(starts) == 1, response.text
return starts[0]
def _burst(rig: Rig, count: int, session_for: Mapping[int, str]) -> tuple[tuple[int, str, str | None], ...]:
markers: Final = tuple(_marker() for _ in range(count))
def one(index: int) -> tuple[int, str, str | None]:
marker: Final = markers[index]
headers: Final = {"Authorization": f"Bearer {rig.proxy.key}", "x-litellm-session-id": session_for[index]}
route: Final = index % 3
try:
if route == 0:
response: Final = rig.proxy.client.post(
"/v1/chat/completions",
json={
"model": rig.model,
"messages": [{"role": "user", "content": marker}],
"stream": index % 2 == 0,
},
headers=headers,
)
response.read()
if response.status_code != 200:
return index, marker, response.text
return index, _chat_id(response), None
if route == 1:
response = rig.proxy.client.post(
"/v1/responses",
json={"model": rig.model, "input": marker, "stream": index % 2 == 0},
headers=headers,
)
response.read()
if response.status_code != 200:
return index, marker, response.text
return index, _responses_id(response), None
response = rig.proxy.client.post(
"/v1/messages",
json={
"model": rig.model,
"max_tokens": 16,
"messages": [{"role": "user", "content": marker}],
"stream": index % 2 == 0,
},
headers=headers,
)
response.read()
if response.status_code != 200:
return index, marker, response.text
return index, _message_id(response), None
except httpx.HTTPError as error:
return index, marker, repr(error)
with ThreadPoolExecutor(max_workers=10) as pool:
return tuple(pool.map(one, range(count)))
def _is_encrypted_responses_id(identity: str) -> bool:
return identity.startswith("resp_") and _decoded_responses_id(identity) == identity
def _landed(rig: Rig, expected: Mapping[str, str]) -> dict[str, tuple[str, ...]]:
spans: Final = rig.sink.attributes()
return {
session: tuple(
_canonical_id(str(attributes["gen_ai.response.id"]["stringValue"]))
for attributes in spans
if attributes.get(CONVERSATION, {}).get("stringValue") == session and "gen_ai.response.id" in attributes
)
for session in expected.values()
}
def _assert_exactly_once(rig: Rig, expected: Mapping[str, str], landed: Mapping[str, tuple[str, ...]]) -> None:
spend: Final = eventually(
lambda: {
session: tuple(_canonical_id(identity) for identity in rig.spend_request_ids(session))
for session in expected.values()
},
lambda rows: all(len(values) >= 1 for values in rows.values()),
seconds=70,
)
assert landed == spend, (landed, spend)
assert all(len(values) == 1 for values in landed.values()), landed
caller_visible: Final = {
session: (_canonical_id(identity),)
for identity, session in expected.items()
if not _is_encrypted_responses_id(identity)
}
assert {session: landed[session] for session in caller_visible} == caller_visible, landed
@pytest.mark.covers("other.observability.otel.conversation_id_sink_outage_recovers_exactly_once")
def test_sink_outage_during_a_mixed_burst_lands_every_response_exactly_once_after_recovery(rig: Rig) -> None:
sessions: Final = {index: f"conv-{index}-{uuid.uuid4().hex}" for index in range(30)}
rig.sink.outage.set()
try:
health_down: Final = rig.proxy.request("GET", "/health/services", params={"service": "otel"})
results: Final = _burst(rig, 30, sessions)
assert all(error is None for _, _, error in results), [error for _, _, error in results if error]
eventually(lambda: any(True for _ in rig.sink.wire.drain()), lambda seen: seen, seconds=30)
finally:
rig.sink.outage.clear()
assert health_down.status_code == 200, health_down.text
expected: Final = {identity: sessions[index] for index, identity, _ in results}
landed: Final = eventually(
lambda: _landed(rig, expected), lambda seen: all(len(values) >= 1 for values in seen.values()), seconds=80
)
_assert_exactly_once(rig, expected, landed)
@pytest.mark.covers("other.observability.otel.conversation_id_slow_sink_no_duplicates")
def test_slow_sink_during_a_burst_lands_every_response_exactly_once(rig: Rig) -> None:
sessions: Final = {index: f"conv-{index}-{uuid.uuid4().hex}" for index in range(20)}
rig.sink.slow.set()
try:
results: Final = _burst(rig, 20, sessions)
assert all(error is None for _, _, error in results), [error for _, _, error in results if error]
expected: Final = {identity: sessions[index] for index, identity, _ in results}
landed: Final = eventually(
lambda: _landed(rig, expected), lambda seen: all(len(values) >= 1 for values in seen.values()), seconds=80
)
finally:
rig.sink.slow.clear()
_assert_exactly_once(rig, expected, landed)
@pytest.mark.covers("other.observability.otel.conversation_id_survives_worker_kill")
def test_killing_one_of_two_workers_mid_burst_keeps_serving_and_never_duplicates_a_span(two_worker_rig: Rig) -> None:
rig: Final = two_worker_rig
root: Final = psutil.Process(rig.process.process.pid)
workers: Final = eventually(
lambda: tuple(child for child in root.children() if "resource_tracker" not in " ".join(child.cmdline())),
lambda found: len(found) == 2,
seconds=30,
)
sessions: Final = {index: f"conv-{index}-{uuid.uuid4().hex}" for index in range(24)}
markers: Final = tuple(_marker() for _ in range(24))
def one(index: int) -> tuple[str, str | None]:
if index == 8:
os.kill(workers[0].pid, signal.SIGKILL)
try:
response: Final = rig.chat(markers[index], headers={"x-litellm-session-id": sessions[index]})
return f"chatcmpl-{markers[index]}", None if response.status_code == 200 else response.text
except httpx.HTTPError as error:
return f"chatcmpl-{markers[index]}", repr(error)
with ThreadPoolExecutor(max_workers=6) as pool:
results: Final = tuple(pool.map(one, range(24)))
assert rig.process.process.poll() is None, "Proxy root exited after a worker was killed"
after: Final = rig.chat(_marker(), headers={"x-litellm-session-id": "conv-after-kill"})
assert after.status_code == 200, after.text
assert rig.sink.single_span(after.json()["id"]) == "conv-after-kill"
failures: Final = tuple(error for _, error in results if error)
assert all(error.startswith(("ReadError(", "RemoteProtocolError(", "ConnectError(")) for error in failures), (
failures
)
assert len(failures) <= 6, failures
served: Final = {identity: sessions[index] for index, (identity, error) in enumerate(results) if error is None}
assert len(served) >= 18, results
settled: Final = {
identity: sessions[index] for index, (identity, error) in enumerate(results) if index > 14 and not error
}
landed: Final = eventually(
lambda: _landed(rig, settled), lambda seen: all(len(values) >= 1 for values in seen.values()), seconds=60
)
_assert_exactly_once(rig, settled, landed)
assert all(len(values) <= 1 for values in _landed(rig, served).values()), _landed(rig, served)
lost: Final = _landed(rig, {identity: sessions[index] for index, (identity, error) in enumerate(results) if error})
assert all(values == () for values in lost.values()), lost
@pytest.mark.covers("other.observability.otel.conversation_id_flushes_on_shutdown")
def test_terminating_the_proxy_right_after_a_burst_flushes_every_span_before_exit(
provider: Wire, collector: Collector, tmp_path_factory: pytest.TempPathFactory
) -> None:
factory: Final = RigFactory(provider, collector, tmp_path_factory.mktemp("otel-shutdown"), {}, 2)
started: Final = factory.start()
rig: Final = next(started)
sessions: Final = {index: f"conv-{index}-{uuid.uuid4().hex}" for index in range(10)}
markers: Final = tuple(_marker() for _ in range(10))
responses: Final = tuple(
rig.chat(markers[index], headers={"x-litellm-session-id": sessions[index]}) for index in range(10)
)
assert all(response.status_code == 200 for response in responses), [response.text for response in responses]
expected: Final = {f"chatcmpl-{markers[index]}": sessions[index] for index in range(10)}
drained: Final = eventually(
lambda: _landed(rig, expected), lambda seen: all(len(values) >= 1 for values in seen.values()), seconds=60
)
assert drained == {session: (identity,) for identity, session in expected.items()}, drained
rig.process.process.terminate()
assert rig.process.process.wait(timeout=40) in (0, -signal.SIGTERM)
assert _landed(rig, expected) == drained
with pytest.raises(httpx.ConnectError):
next(started)

35
tests/integration/run.py
View file

@ -9,7 +9,18 @@ from pathlib import Path
from types import MappingProxyType
from typing import Final
GROUPS: Final = MappingProxyType(json.loads(Path(__file__).with_name("contracts.json").read_text())["groups"])
GROUPS: Final = MappingProxyType(
{
"management": ("management", "authorization", "configuration"),
"accounting": ("pricing", "spend"),
"database": ("database",),
"providers": ("providers", "routing", "streaming"),
"extensions": ("observability", "compatibility"),
"mcp": ("mcp",),
"sdk": ("sdk",),
"cost": ("cost_calculation",),
}
)
def main() -> int:
@ -27,13 +38,9 @@ def main() -> int:
for path in sorted((root / "tests/integration" / folder).glob("test_*.py"))
)
if not selected:
parser.error(f"No integration contracts selected for {options.group}")
parser.error(f"No integration test files selected for {options.group}")
output: Final = options.results.resolve()
output.mkdir(parents=True, exist_ok=True)
manifest: Final = json.loads((root / "tests/integration/contracts.json").read_text())["tests"]
expected: Final = sorted(node for node in manifest if node.split("::", 1)[0] in selected)
if not expected or set(selected) != {node.split("::", 1)[0] for node in expected}:
parser.error("Every selected file must have canonical manifest nodes")
environment: Final = {
**os.environ,
"PYTHONPATH": os.pathsep.join((str(root), str(root / "tests"), str(root / "tests/e2e"))),
@ -58,11 +65,7 @@ def main() -> int:
f"--hypothesis-seed={options.seed}",
f"--integration-order-seed={options.order_seed}",
f"--junitxml={output / 'junit.xml'}",
*(
("-n", str(options.workers))
if options.workers > 1
else ()
),
*(("-n", str(options.workers)) if options.workers > 1 else ()),
],
cwd=root,
env=environment,
@ -70,9 +73,13 @@ def main() -> int:
if result != 0:
return result
evidence: Final = json.loads((output / "execution.json").read_text())
executed: Final = sorted(evidence["passed"] + evidence["skipped"])
if not evidence["complete"] or executed != expected or sorted(evidence["collected"]) != expected:
print("Executed integration nodes differ from the canonical manifest", file=sys.stderr)
collected_files: Final = {node.split("::", 1)[0] for node in evidence["collected"]}
empty: Final = tuple(path for path in selected if path not in collected_files)
if empty:
sys.stderr.write(f"Selected integration files collected zero tests: {', '.join(empty)}\n")
return 1
if not evidence["complete"]:
sys.stderr.write("Integration run did not complete: a collected node neither passed nor skipped\n")
return 1
return 0

59
tests/test_litellm/integrations/otel/test_otel_v2_logger.py
View file

@ -23,6 +23,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( # noqa: E4
from opentelemetry.trace import SpanKind # noqa: E402
from opentelemetry.trace.status import StatusCode # noqa: E402
from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY # noqa: E402
from litellm.integrations.otel import ( # noqa: E402
GenAI,
LiteLLM,
@ -175,6 +176,64 @@ def test_async_log_success_event_emits_llm_call_span():
assert span.status.status_code is StatusCode.UNSET
def test_llm_call_span_carries_the_callers_conversation_id():
logger, exporter = _logger()
kwargs = {**_kwargs(), "litellm_params": {"litellm_session_id": "conv-42", "metadata": {}}}
_emit_llm(logger, kwargs)
(span,) = exporter.get_finished_spans()
assert span.attributes[GenAI.CONVERSATION_ID] == "conv-42"
def test_llm_call_span_without_a_caller_session_has_no_conversation_id():
"""The proxy stamps ``metadata.trace_id`` with the OTel trace id and
``get_litellm_params`` back-fills ``litellm_session_id`` from it."""
logger, exporter = _logger()
otel_trace_id = "6ca5745ef6780d958f62925747f7a5ee"
kwargs = {
**_kwargs(payload=_payload(trace_id=otel_trace_id)),
"litellm_trace_id": otel_trace_id,
"litellm_params": {
"litellm_session_id": otel_trace_id,
"litellm_trace_id": otel_trace_id,
"metadata": {"trace_id": otel_trace_id},
},
}
_emit_llm(logger, kwargs)
(span,) = exporter.get_finished_spans()
assert GenAI.CONVERSATION_ID not in span.attributes
def test_llm_call_span_keeps_the_header_session_when_the_proxy_generated_a_body_one():
"""``missing_session_id: generate`` mints a body session and marks it, but the
caller's ``langfuse_session_id`` header is still their conversation."""
logger, exporter = _logger()
kwargs = {
**_kwargs(),
"litellm_params": {
"litellm_session_id": "minted-by-proxy",
"metadata": {"session_id": "minted-by-proxy", SESSION_ID_GENERATED_METADATA_KEY: True},
"proxy_server_request": {"headers": {"langfuse_session_id": "conv-header"}},
},
}
_emit_llm(logger, kwargs)
(span,) = exporter.get_finished_spans()
assert span.attributes[GenAI.CONVERSATION_ID] == "conv-header"
def test_replayed_llm_call_span_does_not_take_the_payloads_session_id():
"""``/callback_logs`` replays a finished payload whose ``litellm_params`` hold
only key metadata; a session minted under ``missing_session_id: generate``
lands there without its marker, so ``payload.session_id`` is never trusted."""
logger, exporter = _logger()
kwargs = {
**_kwargs(payload=_payload(session_id="minted-then-replayed", trace_id="minted-then-replayed")),
"litellm_params": {"metadata": {"user_api_key_hash": "hsh"}},
}
_emit_llm(logger, kwargs)
(span,) = exporter.get_finished_spans()
assert GenAI.CONVERSATION_ID not in span.attributes
def test_streaming_span_carries_time_to_first_chunk():
logger, exporter = _logger()
kwargs = {

143
tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py
View file

@ -11,6 +11,7 @@ from typing import Final
import pytest
import litellm
from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY
from litellm.integrations.otel import (
BAGGAGE_PROMOTED_KEYS,
DB,
@ -1345,6 +1346,148 @@ def test_llm_span_data_carries_the_caller_trace_controls():
assert LLMCallSpanData.from_standard_logging_payload(_sample_payload()).trace == TraceControls()
@pytest.mark.parametrize(
("litellm_params", "expected"),
[
({"litellm_session_id": "conv-body"}, "conv-body"),
({"metadata": {"session_id": "conv-meta"}}, "conv-meta"),
({"litellm_metadata": {"session_id": "conv-anthropic"}}, "conv-anthropic"),
({"proxy_server_request": {"headers": {"langfuse_session_id": "conv-header"}}}, "conv-header"),
({"litellm_session_id": "conv-body", "metadata": {"session_id": "conv-meta"}}, "conv-body"),
({"litellm_session_id": "", "metadata": {"session_id": ""}}, None),
({"litellm_trace_id": "trace-only", "metadata": {"trace_id": "trace-only"}}, None),
(
{
"litellm_session_id": "0" * 32,
"litellm_trace_id": "0" * 32,
"metadata": {"trace_id": "0" * 32},
},
None,
),
(
{
"litellm_session_id": "0" * 32,
"litellm_trace_id": "0" * 32,
"metadata": {"trace_id": "0" * 32},
"proxy_server_request": {"headers": {"langfuse_session_id": "conv-header"}},
},
"conv-header",
),
(
{
"litellm_session_id": "minted-by-proxy",
"metadata": {"session_id": "minted-by-proxy", SESSION_ID_GENERATED_METADATA_KEY: True},
},
None,
),
(
{
"litellm_session_id": "minted-by-proxy",
"metadata": {"session_id": "minted-by-proxy", SESSION_ID_GENERATED_METADATA_KEY: True},
"proxy_server_request": {"headers": {"langfuse_session_id": "conv-header"}},
},
"conv-header",
),
(
{
"litellm_session_id": "minted-by-proxy",
"metadata": {"session_id": "conv-other-key"},
"litellm_metadata": {"session_id": "minted-by-proxy", SESSION_ID_GENERATED_METADATA_KEY: True},
},
"conv-other-key",
),
(
{
"litellm_session_id": "minted-by-proxy",
"metadata": {"session_id": "minted-by-proxy", SESSION_ID_GENERATED_METADATA_KEY: True},
"litellm_metadata": {"session_id": "conv-other-key"},
},
"conv-other-key",
),
(
{
"litellm_session_id": "conv-x-header",
"litellm_trace_id": "conv-x-header",
"metadata": {"trace_id": "conv-x-header", "session_id": "conv-x-header"},
},
"conv-x-header",
),
({}, None),
],
ids=[
"litellm_session_id",
"metadata",
"anthropic-metadata",
"langfuse-header",
"litellm_session_id-beats-metadata",
"blank-values",
"trace-id-is-not-a-session",
"backfilled-from-otel-trace-id-is-not-a-conversation",
"backfilled-trace-id-does-not-shadow-the-header",
"proxy-generated-is-not-a-conversation",
"proxy-generated-does-not-shadow-the-header",
"proxy-generated-on-litellm_metadata-does-not-shadow-metadata",
"proxy-generated-on-metadata-does-not-shadow-litellm_metadata",
"x-litellm-session-id-header-sets-trace-and-session",
"empty",
],
)
def test_llm_call_event_resolves_the_callers_conversation_id(litellm_params, expected):
kwargs: Final = {"litellm_params": litellm_params, "litellm_trace_id": "per-request-uuid"}
assert LLMCallEvent.from_dict(kwargs).session_id == expected
@pytest.mark.parametrize(
("litellm_params", "payload", "expected"),
[
(
{"metadata": {"user_api_key_hash": "hsh"}},
{"session_id": "minted-then-replayed", "trace_id": "minted-then-replayed"},
None,
),
(
{"metadata": {"user_api_key_hash": "hsh"}},
{"session_id": "conv-replayed", "trace_id": "0af7651916cd43dd8448eb211c80319c"},
None,
),
({"litellm_session_id": "conv-live"}, {"session_id": "conv-replayed"}, "conv-live"),
(
{
"litellm_session_id": "minted-by-proxy",
"metadata": {"session_id": "minted-by-proxy", SESSION_ID_GENERATED_METADATA_KEY: True},
},
{"session_id": "minted-by-proxy"},
None,
),
],
ids=[
"replayed-minted-session-stays-hidden",
"replayed-payload-is-not-a-source",
"live-params-win",
"generated-stays-hidden",
],
)
def test_llm_call_event_never_reads_the_replayed_payloads_session_id(litellm_params, payload, expected):
"""``/callback_logs`` rebuilds ``litellm_params`` with key metadata only, so a
``StandardLoggingPayload`` minted under ``missing_session_id: generate`` arrives
without its generated marker and is indistinguishable from a caller's session;
the payload is therefore never a source for the conversation id."""
kwargs: Final = {
"litellm_params": litellm_params,
"standard_logging_object": _sample_payload(**payload),
}
assert LLMCallEvent.from_dict(kwargs).session_id == expected
def test_llm_span_stamps_gen_ai_conversation_id_only_when_the_caller_sent_one():
with_session: Final = LLMCallSpanData.from_standard_logging_payload(_sample_payload(), session_id="conv-1")
assert GenAIMapper().map(with_session)[GenAI.CONVERSATION_ID] == "conv-1"
without: Final = LLMCallSpanData.from_standard_logging_payload(_sample_payload(trace_id="per-request-uuid"))
assert without.session_id is None
assert GenAI.CONVERSATION_ID not in GenAIMapper().map(without)
def test_llm_span_carries_proxy_request_route():
"""The LLM span records the proxy route the request arrived on, so it can be
filtered by endpoint (``/v1/responses`` vs ``/v1/chat/completions``) without

55
tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py
View file

@ -903,6 +903,61 @@ def test_authorize_post_accepts_ui_session_cookie(unauthenticated_client):
assert _byok_auth_codes[code]["user_id"] == "browser-user-42"
def test_authorize_post_rejects_cookie_with_revoked_session_key(unauthenticated_client):
"""The cookie JWT stays signature-valid until ``exp``, but logout /
password-change revocation deletes the DB-backed session key sealed
inside it. A cookie whose embedded key no longer resolves must not
authorize BYOK writes."""
import jwt as _jwt
with (
patch("litellm.proxy.proxy_server.master_key", "test-master-key"),
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch(
"litellm.proxy.auth.auth_checks.get_key_object",
new=AsyncMock(side_effect=Exception("key not found")),
),
):
cookie_jwt = _jwt.encode(
{
"user_id": "browser-user-42",
"key": "sk-revoked-session-key",
"login_method": "sso",
"exp": int(time.time()) + 3600,
},
"test-master-key",
algorithm="HS256",
)
resp = _authorize_post_with_cookie(unauthenticated_client, cookie_jwt)
assert resp.status_code == 401
def test_authorize_post_accepts_cookie_with_live_session_key(unauthenticated_client):
"""A cookie whose embedded session key still resolves keeps working."""
import jwt as _jwt
with (
patch("litellm.proxy.proxy_server.master_key", "test-master-key"),
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch(
"litellm.proxy.auth.auth_checks.get_key_object",
new=AsyncMock(return_value=UserAPIKeyAuth(user_id="browser-user-42")),
),
):
cookie_jwt = _jwt.encode(
{
"user_id": "browser-user-42",
"key": "sk-live-session-key",
"login_method": "sso",
"exp": int(time.time()) + 3600,
},
"test-master-key",
algorithm="HS256",
)
resp = _authorize_post_with_cookie(unauthenticated_client, cookie_jwt)
assert resp.status_code == 302
def test_authorize_post_rejects_cookie_signed_with_wrong_key(unauthenticated_client):
"""A cookie JWT signed with a different key than the proxy's master_key
must not grant access — otherwise an attacker who can forge a JWT

4
tests/test_litellm/proxy/auth/test_login_utils.py
View file

@ -2137,7 +2137,7 @@ class TestPasswordResetRequiredSessionMinting:
row = _db_user_row(password="Str0ng!Passw0rd", password_reset_required=True)
result, key_kwargs = await self._login(_prisma_with_user(row))
assert key_kwargs["allowed_routes"] == ["/user/password/change"]
assert key_kwargs["allowed_routes"] == ["/user/password/change", "/session/logout"]
assert key_kwargs["metadata"] == {"login_method": "username_password", "password_reset_required": True}
assert result.password_reset_required is True
@ -2198,7 +2198,7 @@ class TestPasswordResetRequiredSessionMinting:
result, key_kwargs, _ = await self._login_with_screen_result(mock_prisma_client, breached=True)
assert key_kwargs["allowed_routes"] == ["/user/password/change"]
assert key_kwargs["allowed_routes"] == ["/user/password/change", "/session/logout"]
assert key_kwargs["metadata"] == {"login_method": "username_password", "password_reset_required": True}
assert result.password_reset_required is True

66
tests/test_litellm/proxy/auth/test_onboarding.py
View file

@ -477,6 +477,72 @@ async def test_claim_token_sets_accepted_at_after_password_written():
assert outer_claims["key"] == "sk-generated-key"
@pytest.mark.asyncio
async def test_claim_token_revokes_existing_ui_sessions():
"""A claimed invite/reset link changes the password; any UI session minted
under the old password may be in hostile hands and must be revoked. The
sweep runs before the fresh session key is minted, so revoke-all is safe."""
from litellm.proxy.proxy_server import claim_onboarding_link
invite = _make_invite(is_accepted=False)
user = _make_user()
prisma = _make_prisma(invite, user)
request = _make_claim_request(_make_onboarding_token())
data = InvitationClaim(
invitation_link="invite-abc",
user_id="user-123",
password="NewP@ssw0rd123",
)
mock_token_response = {"token": "sk-generated-key", "user_id": "user-123"}
revoke_mock = AsyncMock(return_value=1)
mint_order: list[str] = []
async def _mint(*args, **kwargs):
mint_order.append("mint")
return mock_token_response
async def _revoke(*args, **kwargs):
mint_order.append("revoke")
return 1
revoke_mock.side_effect = _revoke
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma),
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
patch( # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
),
patch("litellm.proxy.proxy_server.premium_user", False),
patch(
"litellm.proxy.proxy_server.generate_key_helper_fn",
new_callable=AsyncMock,
side_effect=_mint,
),
patch(
"litellm.proxy.management_endpoints.session_endpoints.revoke_ui_session_keys",
revoke_mock,
),
patch(
"litellm.proxy.proxy_server.get_custom_url",
return_value="http://localhost:4000/",
),
patch(
"litellm.proxy.proxy_server.get_disabled_non_admin_personal_key_creation",
return_value=False,
),
patch("litellm.proxy.proxy_server.get_server_root_path", return_value=""),
):
await claim_onboarding_link(data=data, request=request)
revoke_mock.assert_awaited_once()
assert revoke_mock.await_args.kwargs["user_id"] == "user-123"
# The sweep must precede the mint or it would kill the fresh session too.
assert mint_order == ["revoke", "mint"]
@pytest.mark.asyncio
async def test_claim_token_rolls_back_invite_when_session_key_mint_fails():
"""A session key failure must not leave the invite permanently consumed."""

66
tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -4905,3 +4905,69 @@ async def test_delete_user_writes_deleted_audit_log_for_user_keys(mocker):
assert audit_row.object_id == user_key.token
assert audit_row.changed_by
assert json.loads(audit_row.before_value)["token"] == user_key.token
@pytest.mark.asyncio
async def test_user_update_password_revokes_target_sessions(_admin_prisma, mocker):
"""An admin-set password implies the old one may be compromised: every UI
session belonging to the target user must be revoked after the write."""
from litellm.proxy.management_endpoints.internal_user_endpoints import (
_update_single_user_helper,
)
mocker.patch( # test-quality-ok: same module-global mocking every test in this file already uses
"litellm.proxy.proxy_server.general_settings",
{"password_policy_check_breached_passwords": False},
)
mock_prisma_client = _admin_prisma
existing_user = mocker.MagicMock()
existing_user.model_dump.return_value = {"user_id": "target-user"}
existing_user.user_id = "target-user"
mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user)
mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "target-user"})
mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x)
revoke_mock = mocker.patch(
"litellm.proxy.management_endpoints.session_endpoints.revoke_ui_session_keys",
new=mocker.AsyncMock(return_value=2),
)
user_request = UpdateUserRequest(user_id="target-user", password="Str0ng!Passw0rd")
admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN)
await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller)
revoke_mock.assert_awaited_once()
revoke_kwargs = revoke_mock.await_args.kwargs
assert revoke_kwargs["user_id"] == "target-user"
# Revoke-all: the admin's own session is not among the target's sessions.
assert revoke_kwargs.get("keep_hashed_token") is None
@pytest.mark.asyncio
async def test_user_update_without_password_revokes_nothing(_admin_prisma, mocker):
"""A non-password /user/update must not touch the target's sessions."""
from litellm.proxy.management_endpoints.internal_user_endpoints import (
_update_single_user_helper,
)
mock_prisma_client = _admin_prisma
existing_user = mocker.MagicMock()
existing_user.model_dump.return_value = {"user_id": "target-user"}
existing_user.user_id = "target-user"
mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user)
mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "target-user"})
mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x)
revoke_mock = mocker.patch(
"litellm.proxy.management_endpoints.session_endpoints.revoke_ui_session_keys",
new=mocker.AsyncMock(return_value=0),
)
user_request = UpdateUserRequest(user_id="target-user", user_email="new@example.com")
admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN)
await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller)
revoke_mock.assert_not_awaited()

24
tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -4796,6 +4796,30 @@ class TestPatchModelCredentialName:
assert "not found" in exc_info.value.message.lower()
credentials_repository.find_by_name.assert_awaited_once_with("ghost-credential")
@pytest.mark.asyncio
async def test_patch_model_resending_unchanged_dangling_credential_name_is_not_validated(self, monkeypatch):
credentials_repository = MagicMock()
db_model: Final = Deployment(
model_name="gpt-4",
litellm_params=LiteLLM_Params(
model="openai/gpt-4o",
api_base="https://api.openai.com/v1",
litellm_credential_name="ghost-credential",
),
model_info=ModelInfo(id="dep-cred-1"),
)
persisted: Final = await self._patch_model(
monkeypatch,
db_model,
self._admin_user(),
"ghost-credential",
credentials_repository=credentials_repository,
)
params: Final = json.loads(persisted[0]["litellm_params"])
assert params["litellm_credential_name"] == "ghost-credential"
credentials_repository.find_by_name.assert_not_awaited()
@pytest.mark.asyncio
async def test_patch_model_accepts_credential_known_only_in_db(self, monkeypatch):
db_model: Final = Deployment(

67
tests/test_litellm/proxy/management_endpoints/test_password_endpoints.py
View file

@ -380,6 +380,73 @@ async def test_change_password_failure_emits_no_audit_log():
audit_mock.assert_not_awaited()
@pytest.mark.asyncio
async def test_change_password_revokes_other_sessions_keeping_callers():
"""A successful change revokes the user's other UI sessions (the old
password may be compromised) while keeping the session that just proved
it holds the current password."""
from litellm.proxy._types import ChangePasswordRequest
prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD)))
revoke_mock = AsyncMock(return_value=0)
caller = UserAPIKeyAuth(
user_id="user-123",
token="hashed-caller-token",
team_id=UI_TEAM_ID,
metadata=dict(PASSWORD_SESSION_METADATA),
)
with (
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.prisma_client", prisma
),
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
),
patch(
"litellm.proxy.management_endpoints.password_endpoints.revoke_ui_session_keys",
revoke_mock,
),
):
await change_password(
data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=NEW_PASSWORD),
user_api_key_dict=caller,
)
revoke_mock.assert_awaited_once()
revoke_kwargs = revoke_mock.await_args.kwargs
assert revoke_kwargs["user_id"] == "user-123"
assert revoke_kwargs["keep_hashed_token"] == "hashed-caller-token"
@pytest.mark.asyncio
async def test_change_password_failure_revokes_no_sessions():
from litellm.proxy._types import ChangePasswordRequest
prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD)))
revoke_mock = AsyncMock(return_value=0)
with (
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.prisma_client", prisma
),
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
),
patch(
"litellm.proxy.management_endpoints.password_endpoints.revoke_ui_session_keys",
revoke_mock,
),
):
with pytest.raises(HTTPException):
await change_password(
data=ChangePasswordRequest(current_password="not-the-password", new_password=NEW_PASSWORD),
user_api_key_dict=_caller(),
)
revoke_mock.assert_not_awaited()
@pytest.mark.asyncio
async def test_change_password_requires_db():
from litellm.proxy._types import ChangePasswordRequest

300
tests/test_litellm/proxy/management_endpoints/test_session_endpoints.py Normal file
View file

@ -0,0 +1,300 @@
"""
Tests for POST /session/logout and revoke_ui_session_keys
(litellm/proxy/management_endpoints/session_endpoints.py).
"""
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException, Response
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.proxy._types import LiteLLM_VerificationToken, UserAPIKeyAuth
from litellm.proxy.management_endpoints.session_endpoints import (
revoke_ui_session_keys,
session_logout,
)
HASHED_TOKEN = "hashed-session-token"
USER_ID = "user-123"
def _session_row(token: str = HASHED_TOKEN, user_id: str = USER_ID) -> LiteLLM_VerificationToken:
return LiteLLM_VerificationToken(token=token, team_id=UI_SESSION_TOKEN_TEAM_ID, user_id=user_id)
def _make_prisma(
find_unique_row: LiteLLM_VerificationToken | None = None,
find_many_rows: list[LiteLLM_VerificationToken] | None = None,
) -> MagicMock:
prisma = MagicMock()
table = prisma.db.litellm_verificationtoken
table.find_unique = AsyncMock(return_value=find_unique_row)
table.find_many = AsyncMock(return_value=find_many_rows or [])
table.delete_many = AsyncMock(return_value=1)
return prisma
def _ui_session_caller(token: str | None = HASHED_TOKEN) -> UserAPIKeyAuth:
return UserAPIKeyAuth(token=token, team_id=UI_SESSION_TOKEN_TEAM_ID, user_id=USER_ID)
def _patched_globals(prisma):
return (
patch( # test-quality-ok: endpoint reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.prisma_client", prisma
),
patch( # test-quality-ok: endpoint reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.proxy_logging_obj", None
),
patch( # test-quality-ok: endpoint reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.user_api_key_cache", MagicMock()
),
)
@pytest.mark.asyncio
async def test_session_logout_revokes_presented_session():
prisma = _make_prisma(find_unique_row=_session_row())
persist_mock = AsyncMock()
evict_mock = AsyncMock()
p1, p2, p3 = _patched_globals(prisma)
with (
p1,
p2,
p3,
patch(
"litellm.proxy.management_endpoints.session_endpoints._persist_deleted_verification_tokens",
persist_mock,
),
patch(
"litellm.proxy.management_endpoints.session_endpoints.delete_cache_key_objects",
evict_mock,
),
):
response = await session_logout(
response=Response(),
user_api_key_dict=_ui_session_caller(),
)
assert response.message == "Session revoked."
delete_kwargs = prisma.db.litellm_verificationtoken.delete_many.call_args.kwargs
assert delete_kwargs["where"] == {"token": HASHED_TOKEN}
# Audit record persisted before the row is gone.
persist_mock.assert_awaited_once()
assert persist_mock.await_args.kwargs["keys"][0].token == HASHED_TOKEN
# Cache evicted + broadcast even on the delete path.
evict_mock.assert_awaited_once()
assert tuple(evict_mock.await_args.kwargs["hashed_tokens"]) == (HASHED_TOKEN,)
@pytest.mark.asyncio
async def test_session_logout_clears_token_cookie():
prisma = _make_prisma(find_unique_row=_session_row())
fastapi_response = Response()
p1, p2, p3 = _patched_globals(prisma)
with (
p1,
p2,
p3,
patch(
"litellm.proxy.management_endpoints.session_endpoints._persist_deleted_verification_tokens",
AsyncMock(),
),
patch(
"litellm.proxy.management_endpoints.session_endpoints.delete_cache_key_objects",
AsyncMock(),
),
):
await session_logout(
response=fastapi_response,
user_api_key_dict=_ui_session_caller(),
)
set_cookie_headers = [v.decode() for k, v in fastapi_response.raw_headers if k == b"set-cookie"]
assert any(h.startswith('token="";') or h.startswith("token=;") for h in set_cookie_headers)
@pytest.mark.asyncio
async def test_session_logout_refuses_non_ui_session_key():
"""The endpoint must not become a generic key-deletion oracle: a normal
virtual key (no UI team id) is refused outright."""
prisma = _make_prisma()
p1, p2, p3 = _patched_globals(prisma)
with p1, p2, p3:
with pytest.raises(HTTPException) as exc_info:
await session_logout(
response=Response(),
user_api_key_dict=UserAPIKeyAuth(token=HASHED_TOKEN, team_id="some-real-team", user_id=USER_ID),
)
assert exc_info.value.status_code == 403
prisma.db.litellm_verificationtoken.delete_many.assert_not_called()
@pytest.mark.asyncio
async def test_session_logout_is_idempotent_when_row_already_gone():
prisma = _make_prisma(find_unique_row=None)
evict_mock = AsyncMock()
p1, p2, p3 = _patched_globals(prisma)
with (
p1,
p2,
p3,
patch(
"litellm.proxy.management_endpoints.session_endpoints.delete_cache_key_objects",
evict_mock,
),
):
response = await session_logout(
response=Response(),
user_api_key_dict=_ui_session_caller(),
)
assert response.message == "Session already revoked."
prisma.db.litellm_verificationtoken.delete_many.assert_not_called()
# The cache entry may outlive the row; evict regardless.
evict_mock.assert_awaited_once()
@pytest.mark.asyncio
async def test_session_logout_requires_db():
p2 = patch("litellm.proxy.proxy_server.proxy_logging_obj", None)
p3 = patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
with (
patch( # test-quality-ok: endpoint reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.prisma_client", None
),
p2,
p3,
):
with pytest.raises(HTTPException) as exc_info:
await session_logout(
response=Response(),
user_api_key_dict=_ui_session_caller(),
)
assert exc_info.value.status_code == 500
@pytest.mark.asyncio
async def test_revoke_ui_session_keys_revokes_all_and_broadcasts():
rows = [_session_row(token="t1"), _session_row(token="t2"), _session_row(token="t3")]
prisma = _make_prisma(find_many_rows=rows)
persist_mock = AsyncMock()
evict_mock = AsyncMock()
p1, p2, p3 = _patched_globals(prisma)
with (
p1,
p2,
p3,
patch(
"litellm.proxy.management_endpoints.session_endpoints._persist_deleted_verification_tokens",
persist_mock,
),
patch(
"litellm.proxy.management_endpoints.session_endpoints.delete_cache_key_objects",
evict_mock,
),
):
revoked = await revoke_ui_session_keys(
user_id=USER_ID,
user_api_key_dict=_ui_session_caller(),
)
assert revoked == 3
find_kwargs = prisma.db.litellm_verificationtoken.find_many.call_args.kwargs
assert find_kwargs["where"] == {"user_id": USER_ID, "team_id": UI_SESSION_TOKEN_TEAM_ID}
delete_kwargs = prisma.db.litellm_verificationtoken.delete_many.call_args.kwargs
assert delete_kwargs["where"] == {"token": {"in": ["t1", "t2", "t3"]}}
persist_mock.assert_awaited_once()
evict_mock.assert_awaited_once()
assert evict_mock.await_args.kwargs["hashed_tokens"] == ["t1", "t2", "t3"]
@pytest.mark.asyncio
async def test_revoke_ui_session_keys_keeps_callers_session():
rows = [_session_row(token="t1"), _session_row(token=HASHED_TOKEN), _session_row(token="t3")]
prisma = _make_prisma(find_many_rows=rows)
p1, p2, p3 = _patched_globals(prisma)
with (
p1,
p2,
p3,
patch(
"litellm.proxy.management_endpoints.session_endpoints._persist_deleted_verification_tokens",
AsyncMock(),
),
patch(
"litellm.proxy.management_endpoints.session_endpoints.delete_cache_key_objects",
AsyncMock(),
),
):
revoked = await revoke_ui_session_keys(
user_id=USER_ID,
user_api_key_dict=_ui_session_caller(),
keep_hashed_token=HASHED_TOKEN,
)
assert revoked == 2
delete_kwargs = prisma.db.litellm_verificationtoken.delete_many.call_args.kwargs
assert delete_kwargs["where"] == {"token": {"in": ["t1", "t3"]}}
@pytest.mark.asyncio
async def test_revoke_ui_session_keys_noop_when_no_sessions():
prisma = _make_prisma(find_many_rows=[])
p1, p2, p3 = _patched_globals(prisma)
with p1, p2, p3:
revoked = await revoke_ui_session_keys(
user_id=USER_ID,
user_api_key_dict=_ui_session_caller(),
)
assert revoked == 0
prisma.db.litellm_verificationtoken.delete_many.assert_not_called()
@pytest.mark.asyncio
async def test_revoke_ui_session_keys_failure_is_swallowed():
"""The password write has already committed when this runs; a revocation
failure must not fail the caller's request."""
prisma = _make_prisma(find_many_rows=[_session_row(token="t1")])
prisma.db.litellm_verificationtoken.delete_many = AsyncMock(side_effect=RuntimeError("db down"))
p1, p2, p3 = _patched_globals(prisma)
with (
p1,
p2,
p3,
patch(
"litellm.proxy.management_endpoints.session_endpoints._persist_deleted_verification_tokens",
AsyncMock(),
),
):
revoked = await revoke_ui_session_keys(
user_id=USER_ID,
user_api_key_dict=_ui_session_caller(),
)
assert revoked == 0
@pytest.mark.asyncio
async def test_revoke_ui_session_keys_noop_without_db():
with patch( # test-quality-ok: helper reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.prisma_client", None
):
revoked = await revoke_ui_session_keys(
user_id=USER_ID,
user_api_key_dict=_ui_session_caller(),
)
assert revoked == 0

35
tests/test_litellm/test_assert_ci_coverage.py
View file

@ -9,7 +9,6 @@ the question neither covers: whether the job that globs a file then deselects it
"""
import importlib.util
import json
import sys
from pathlib import Path
from typing import Final
@ -24,13 +23,14 @@ sys.modules[_spec.name] = coverage # @dataclass(slots=True) rebuilds via sys.mo
_spec.loader.exec_module(coverage)
def test_integration_manifest_requires_exclusive_scheduled_circleci_owner(tmp_path: Path) -> None:
def test_integration_groups_require_exclusive_scheduled_circleci_owner(tmp_path: Path) -> None:
test_path: Final = "tests/integration/management/test_contract.py"
test_file: Final = tmp_path / test_path
test_file.parent.mkdir(parents=True)
test_file.write_text("def test_contract(): pass\n")
(tmp_path / "tests/integration/contracts.json").write_text(
json.dumps({"groups": {"management": ["management"]}, "tests": {f"{test_path}::test_contract": ["mgmt.test"]}})
(tmp_path / "tests/integration/run.py").write_text(
"from types import MappingProxyType\nfrom typing import Final\n"
'GROUPS: Final = MappingProxyType({"management": ("management",)})\n'
)
paths, findings = coverage._integration_ownership(tmp_path)
assert not paths
@ -144,9 +144,7 @@ def test_the_parent_token_alone_does_not_satisfy_any_child(tmp_path):
(root / "billing").mkdir(parents=True)
(root / "billing" / "test_a.py").write_text("def test_a(): assert True\n")
findings = coverage._unassigned_shard_children(
frozenset({"tests/tree"}), roots=("tests/tree",), repo_root=tmp_path
)
findings = coverage._unassigned_shard_children(frozenset({"tests/tree"}), roots=("tests/tree",), repo_root=tmp_path)
assert tuple(f.subject for f in findings) == ("tests/tree/billing",)
@ -181,8 +179,12 @@ def test_the_repo_as_it_stands_has_every_shard_child_assigned():
def _slice(**overrides):
defaults = dict(
job="a_job", globs=("tests/x/**/test_*.py",), named=frozenset(),
required=(), excluded=(), understood=True,
job="a_job",
globs=("tests/x/**/test_*.py",),
named=frozenset(),
required=(),
excluded=(),
understood=True,
)
return coverage.Slice(**{**defaults, **overrides})
@ -224,9 +226,7 @@ def test_an_explicitly_named_file_is_claimed_whatever_the_keywords_say():
def test_an_unparsed_keyword_expression_claims_everything_it_globs():
# Staying silent beats guessing: an expression this parser cannot model must never
# be the reason a file is reported as unrun.
assert _slice(understood=False, excluded=("cache",)).claims(
"tests/x/test_caching.py", frozenset()
) is True
assert _slice(understood=False, excluded=("cache",)).claims("tests/x/test_caching.py", frozenset()) is True
def test_keyword_terms_splits_an_and_chain_into_required_and_excluded():
@ -339,10 +339,7 @@ def test_a_dockerfile_directory_entry_is_stale_because_only_an_exact_path_exempt
def test_a_workflow_that_names_a_file_clears_it_from_the_slice_check():
named = coverage._workflow_named_tokens()
assert named, "the workflows must name some test paths or the check proves nothing"
assert any(
coverage._token_covers(token, "tests/local_testing/test_caching_handler.py")
for token in named
)
assert any(coverage._token_covers(token, "tests/local_testing/test_caching_handler.py") for token in named)
def test_the_slice_check_credits_only_workflows_never_the_circleci_config():
@ -355,6 +352,6 @@ def test_the_slice_check_credits_only_workflows_never_the_circleci_config():
def test_a_file_no_workflow_names_is_still_reported_when_every_slice_drops_it():
named = coverage._workflow_named_tokens()
assert not any(
coverage._token_covers(token, "tests/local_testing/test_caching.py") for token in named
), "test_caching.py is allowlisted, not run; crediting it would hide a real gap"
assert not any(coverage._token_covers(token, "tests/local_testing/test_caching.py") for token in named), (
"test_caching.py is allowlisted, not run; crediting it would hide a real gap"
)

6
ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.tsx
View file

@ -15,7 +15,7 @@ import { changePasswordCall, getProxyBaseUrl } from "@/components/networking";
import { extractProxyErrorMessage } from "@/lib/http/client";
import { useZodForm } from "@/lib/forms/useZodForm";
import { toast } from "@/lib/toast";
import { clearTokenCookies } from "@/utils/cookieUtils";
import { revokeSessionAndClearClientState } from "@/app/(dashboard)/hooks/useLogout";
import { getLoginUrl } from "@/utils/returnUrlUtils";
const changePasswordSchema = z
@ -47,8 +47,10 @@ export function ChangePasswordForm() {
await changePasswordCall(accessToken, values.currentPassword, values.newPassword);
if (passwordResetRequired) {
// The session key was minted restricted; only a fresh login lifts it.
// Revoke it server-side too (best-effort) so it doesn't sit valid
// until the expiry reaper gets to it.
toast.success("Password updated. Please log in with your new password.");
clearTokenCookies();
await revokeSessionAndClearClientState(accessToken);
window.location.replace(getLoginUrl(getProxyBaseUrl()));
return;
}

72
ui/litellm-dashboard/src/app/(dashboard)/hooks/useLogout.test.ts Normal file
View file

@ -0,0 +1,72 @@
import { beforeEach, describe, expect, it, vi } from "vitest";
const sessionLogoutCall = vi.hoisted(() => vi.fn());
const clearTokenCookies = vi.hoisted(() => vi.fn());
const clearStoredReturnUrl = vi.hoisted(() => vi.fn());
vi.mock("@/components/networking", () => ({
sessionLogoutCall,
}));
vi.mock("@/utils/cookieUtils", () => ({
clearTokenCookies,
}));
vi.mock("@/utils/returnUrlUtils", () => ({
clearStoredReturnUrl,
}));
vi.mock("@/app/(dashboard)/hooks/proxySettings/useProxySettings", () => ({
default: vi.fn(() => ({ PROXY_LOGOUT_URL: "" })),
}));
import { revokeSessionAndClearClientState } from "./useLogout";
describe("revokeSessionAndClearClientState", () => {
beforeEach(() => {
vi.clearAllMocks();
sessionLogoutCall.mockResolvedValue({ message: "Session revoked." });
localStorage.setItem("litellm_selected_worker_id", "w1");
localStorage.setItem("litellm_worker_url", "https://worker.example");
});
it("revokes the session server-side before clearing the token cookie", async () => {
const order: string[] = [];
sessionLogoutCall.mockImplementation(async () => {
order.push("revoke");
return { message: "Session revoked." };
});
clearTokenCookies.mockImplementation(() => {
order.push("clearCookies");
});
await revokeSessionAndClearClientState("sk-token");
expect(sessionLogoutCall).toHaveBeenCalledWith("sk-token");
// The cookie holds the credential that authenticates the revoke call, so
// clearing it first would orphan the server-side key.
expect(order).toEqual(["revoke", "clearCookies"]);
});
it("clears all client state", async () => {
await revokeSessionAndClearClientState("sk-token");
expect(clearTokenCookies).toHaveBeenCalled();
expect(clearStoredReturnUrl).toHaveBeenCalled();
expect(localStorage.getItem("litellm_selected_worker_id")).toBeNull();
expect(localStorage.getItem("litellm_worker_url")).toBeNull();
});
it("still clears client state when the revoke call rejects", async () => {
sessionLogoutCall.mockRejectedValue(new Error("proxy unreachable"));
await revokeSessionAndClearClientState("sk-token");
expect(clearTokenCookies).toHaveBeenCalled();
expect(localStorage.getItem("litellm_selected_worker_id")).toBeNull();
});
it("skips the server call without a token but still clears client state", async () => {
await revokeSessionAndClearClientState(null);
expect(sessionLogoutCall).not.toHaveBeenCalled();
expect(clearTokenCookies).toHaveBeenCalled();
});
});

30
ui/litellm-dashboard/src/app/(dashboard)/hooks/useLogout.ts
View file

@ -1,7 +1,29 @@
import { sessionLogoutCall } from "@/components/networking";
import { clearTokenCookies } from "@/utils/cookieUtils";
import { clearStoredReturnUrl } from "@/utils/returnUrlUtils";
import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings";
/**
* Revokes the session key server-side, then clears client state. Exported for
* flows that navigate somewhere other than PROXY_LOGOUT_URL (worker switch,
* forced password reset). The server call must happen BEFORE the cookies are
* cleared (the token authenticates it) and is best-effort: local logout must
* still complete when the server is unreachable.
*/
export async function revokeSessionAndClearClientState(accessToken: string | null): Promise<void> {
if (accessToken) {
try {
await sessionLogoutCall(accessToken);
} catch {
// Best-effort: the key still expires server-side at its session TTL.
}
}
clearTokenCookies();
clearStoredReturnUrl();
localStorage.removeItem("litellm_selected_worker_id");
localStorage.removeItem("litellm_worker_url");
}
/**
* Shared sign-out handler. Used by both the top navbar and the sidebar footer so
* the two entry points can never drift on which client state gets cleared.
@ -10,10 +32,8 @@ export function useLogout(accessToken: string | null): () => void {
const proxySettings = useProxySettings(accessToken);
return () => {
clearTokenCookies();
clearStoredReturnUrl();
localStorage.removeItem("litellm_selected_worker_id");
localStorage.removeItem("litellm_worker_url");
window.location.href = proxySettings.PROXY_LOGOUT_URL || "";
void revokeSessionAndClearClientState(accessToken).finally(() => {
window.location.href = proxySettings.PROXY_LOGOUT_URL || "";
});
};
}

21
ui/litellm-dashboard/src/components/navbar.tsx
View file

@ -5,9 +5,8 @@ import { useWorker } from "@/hooks/useWorker";
import { getProxyBaseUrl } from "@/components/networking";
import { uiHref } from "@/utils/uiHref";
import { useTheme } from "@/contexts/ThemeContext";
import { clearTokenCookies } from "@/utils/cookieUtils";
import { clearStoredReturnUrl, getLoginUrl } from "@/utils/returnUrlUtils";
import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings";
import { revokeSessionAndClearClientState, useLogout } from "@/app/(dashboard)/hooks/useLogout";
import { getLoginUrl } from "@/utils/returnUrlUtils";
import { Badge } from "@/components/ui/badge";
import { PanelLeftClose, PanelLeftOpen } from "lucide-react";
import Link from "next/link";
@ -38,7 +37,6 @@ const Navbar: React.FC<NavbarProps> = ({
onToggleSidebar,
}) => {
const baseUrl = getProxyBaseUrl();
const proxySettings = useProxySettings(accessToken);
const { logoUrl } = useTheme();
const { data: healthData } = useHealthReadinessDetails(accessToken);
const version = healthData?.litellm_version;
@ -50,19 +48,12 @@ const Navbar: React.FC<NavbarProps> = ({
const imageUrl = logoUrl || `${baseUrl}/get_image`;
const darkImageUrl = logoUrl || `${baseUrl}/get_image?theme=dark`;
const handleLogout = () => {
clearTokenCookies();
localStorage.removeItem("litellm_selected_worker_id");
localStorage.removeItem("litellm_worker_url");
window.location.href = proxySettings.PROXY_LOGOUT_URL || "";
};
const handleLogout = useLogout(accessToken);
const handleWorkerSwitch = (workerId: string) => {
clearTokenCookies();
clearStoredReturnUrl();
localStorage.removeItem("litellm_selected_worker_id");
localStorage.removeItem("litellm_worker_url");
window.location.href = `${getLoginUrl()}?worker=${encodeURIComponent(workerId)}`;
void revokeSessionAndClearClientState(accessToken).finally(() => {
window.location.href = `${getLoginUrl()}?worker=${encodeURIComponent(workerId)}`;
});
};
return (

12
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -1595,6 +1595,18 @@ export const claimOnboardingToken = async (
}
};
/**
* Revokes the UI session key server-side (POST /session/logout). Best-effort
* with a short timeout: logout must still complete locally when the server is
* unreachable, so callers swallow rejections.
*/
export const sessionLogoutCall = async (accessToken: string): Promise<{ message: string }> => {
return await apiClient.post(`/session/logout`, {
accessToken,
signal: AbortSignal.timeout(3000),
});
};
export const changePasswordCall = async (
accessToken: string,
currentPassword: string,

50
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -14485,6 +14485,31 @@ export interface paths {
patch?: never;
trace?: never;
};
"/session/logout": {
parameters: {
query?: never;
header?: never;
path?: never;
cookie?: never;
};
get?: never;
put?: never;
/**
* Session Logout
* @description Revoke the UI session key this request authenticated with.
*
* Only accepts UI session keys (minted by dashboard login); any other
* credential is refused, so this can never be used to delete arbitrary keys.
* Revokes only the presented session, not the user's other sessions.
* Idempotent: logging out an already-revoked session succeeds.
*/
post: operations["session_logout_session_logout_post"];
delete?: never;
options?: never;
head?: never;
patch?: never;
trace?: never;
};
"/settings": {
parameters: {
query?: never;
@ -38133,6 +38158,11 @@ export interface components {
/** Timeout */
timeout?: number | null;
};
/** SessionLogoutResponse */
SessionLogoutResponse: {
/** Message */
message: string;
};
/**
* ShadowEvalJobResponse
* @description A shadow-eval job over one or more targets, each with its own budget and stop state;
@ -60958,6 +60988,26 @@ export interface operations {
};
};
};
session_logout_session_logout_post: {
parameters: {
query?: never;
header?: never;
path?: never;
cookie?: never;
};
requestBody?: never;
responses: {
/** @description Successful Response */
200: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["SessionLogoutResponse"];
};
};
};
};
active_callbacks_settings_get: {
parameters: {
query?: never;