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

chore: merge main into litellm_vertex_context_cache_creation_accounting
Some checks are pending
LiteLLM Rust / rust-lint (push) Waiting to run
Details
LiteLLM Rust / rust-test (push) Waiting to run
Details
LiteLLM Rust / rust-wheel (push) Waiting to run
Details
Terraform Modules / fmt, validate, test (aws) (push) Waiting to run
Details
Terraform Modules / fmt, validate, test (gcp) (push) Waiting to run
Details
Terraform Provider / gofmt, vet, build, test (push) Waiting to run
Details
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Waiting to run
Details

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:05 +00:00
commit 03e8b0937f
151 changed files with 13905 additions and 2540 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:

39
.circleci/scripts/run_integration.sh
View file

@ -9,7 +9,6 @@ fi
suite="${1:?integration suite required}"
results="test-results/integration-${suite}"
mkdir -p "$results"
shard_timeout=11m
integration_identity="$(.venv/bin/python -c 'import uuid; print(uuid.uuid4().hex)')"
upstream_pid=""
proxy_pid=""
@ -112,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"
@ -131,10 +139,11 @@ start_proxy() {
fi
setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
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"
@ -176,7 +185,7 @@ if [ "$suite" = browser ]; then
exit 0
fi
timeout --signal=TERM --kill-after=20s "$shard_timeout" env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
INTEGRATION_RUN_ID="$integration_identity" \
DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
INTEGRATION_PROXY_URL="$INTEGRATION_PROXY_URL" INTEGRATION_PEER_URL="$INTEGRATION_PEER_URL" \
@ -187,3 +196,23 @@ timeout --signal=TERM --kill-after=20s "$shard_timeout" env -i PATH="$PATH" HOME
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

675
litellm/model_prices_and_context_window_backup.json
View file

File diff suppressed because it is too large Load diff

1
litellm/models/managed_files.py
View file

@ -61,3 +61,4 @@ class LiteLLM_ManagedVectorStoresTable(LiteLLMPydanticObjectBase):
litellm_params: dict[str, Any] | None = None
team_id: str | None = None
user_id: str | None = None
is_config: bool = False

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)

13
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -45790,6 +45790,10 @@
"title": "Custom Llm Provider",
"type": "string"
},
"is_config": {
"title": "Is Config",
"type": "boolean"
},
"litellm_credential_name": {
"anyOf": [
{
@ -45962,6 +45966,11 @@
"title": "Custom Llm Provider",
"type": "string"
},
"is_config": {
"default": false,
"title": "Is Config",
"type": "boolean"
},
"litellm_credential_name": {
"anyOf": [
{
@ -46203,7 +46212,7 @@
"paths": {
"/v1/vector_store/list": {
"get": {
"description": "List all available vector stores with optional filtering and pagination.\nCombines both in-memory vector stores and those stored in the database.\nDatabase is the source of truth - deleted stores are removed from memory, updated stores sync to memory.\n\nParameters:\n- page: int - Page number for pagination (default: 1)\n- page_size: int - Number of items per page (default: 100)",
"description": "List all available vector stores with optional filtering and pagination.\nCombines both in-memory vector stores and those stored in the database.\nDatabase is the source of truth for stores it owns: deleted stores are removed from memory, updated stores\nsync to memory. Stores declared in the config file are owned by the config file, are always listed, and are\nnever overwritten by database rows.\n\nParameters:\n- page: int - Page number for pagination (default: 1)\n- page_size: int - Number of items per page (default: 100)",
"operationId": "list_vector_stores_v1_vector_store_list_get",
"parameters": [
{
@ -46354,7 +46363,7 @@
},
"/vector_store/list": {
"get": {
"description": "List all available vector stores with optional filtering and pagination.\nCombines both in-memory vector stores and those stored in the database.\nDatabase is the source of truth - deleted stores are removed from memory, updated stores sync to memory.\n\nParameters:\n- page: int - Page number for pagination (default: 1)\n- page_size: int - Number of items per page (default: 100)",
"description": "List all available vector stores with optional filtering and pagination.\nCombines both in-memory vector stores and those stored in the database.\nDatabase is the source of truth for stores it owns: deleted stores are removed from memory, updated stores\nsync to memory. Stores declared in the config file are owned by the config file, are always listed, and are\nnever overwritten by database rows.\n\nParameters:\n- page: int - Page number for pagination (default: 1)\n- page_size: int - Number of items per page (default: 100)",
"operationId": "list_vector_stores_vector_store_list_get",
"parameters": [
{

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

59
litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py
View file

@ -18,6 +18,10 @@ if TYPE_CHECKING:
AUTH_CACHE_INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation"
_POLL_TIMEOUT_SECONDS: Final = 1.0
_MAX_PENDING_PUBLISHES: Final = 1024
_MAX_IN_FLIGHT_PUBLISHES: Final = 16
_pending_publishes: Final[set[asyncio.Task[None]]] = set() # mutable-ok: strong refs keep background publishes alive
_in_flight_publishes: Final = asyncio.Semaphore(_MAX_IN_FLIGHT_PUBLISHES)
_BACKOFF_INITIAL_SECONDS: Final = 5.0
_BACKOFF_MAX_SECONDS: Final = 60.0
@ -67,6 +71,21 @@ def _message_from_data(data: object) -> _CacheInvalidationMessage | None:
)
async def _publish_to_redis(redis_cache: "RedisCache", cache_key: str, message: str) -> None:
try:
client: Final = _pubsub_capable_client(redis_cache)
if client is None:
verbose_proxy_logger.debug(
"auth cache invalidation publish for %s skipped: cluster redis client has no pub/sub support",
cache_key,
)
return
async with _in_flight_publishes:
await client.publish(auth_cache_invalidation_channel(redis_cache), message)
except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors
verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e)
async def publish_auth_cache_invalidation(
cache_key: str, new_value: float | None = None, ttl: float | None = None
) -> None:
@ -80,24 +99,34 @@ async def publish_auth_cache_invalidation(
writes the value into its additional in-memory caches rather than deleting
the key. A spend reset uses this so the handler's self-delivered message
cannot erase the freshly-written post-reset counter or floor marker.
The Redis round trip runs as a background task: this call returns once the
publish has been handed to the event loop, so a Redis that accepts
connections but never replies costs the caller nothing. The DB write has
already committed and the local eviction already happened, so the caller
has nothing to do with the publish result. At most 16 publishes hold a
Redis connection at once; the rest wait in the task set, so a wedge cannot
drain the shared connection pool.
"""
redis_cache: Final = coordination_redis_cache()
if redis_cache is None:
return
try:
client: Final = _pubsub_capable_client(redis_cache)
if client is None:
verbose_proxy_logger.debug(
"auth cache invalidation publish for %s skipped: cluster redis client has no pub/sub support",
cache_key,
)
return
await client.publish(
auth_cache_invalidation_channel(redis_cache),
_cache_invalidation_message_json(cache_key, new_value=new_value, ttl=ttl),
_pending_publishes.difference_update({task for task in _pending_publishes if task.done()})
if len(_pending_publishes) >= _MAX_PENDING_PUBLISHES:
verbose_proxy_logger.warning(
"auth cache invalidation publish for %s dropped: %d publishes already waiting on redis; "
"other workers keep their cached copy until its TTL expires",
cache_key,
len(_pending_publishes),
)
except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors
verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e)
return
task: Final = asyncio.create_task(
_publish_to_redis(
redis_cache, cache_key, _cache_invalidation_message_json(cache_key, new_value=new_value, ttl=ttl)
)
)
_pending_publishes.add(task)
await asyncio.sleep(0)
async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "UserApiKeyCache") -> None:
@ -106,8 +135,8 @@ async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "Us
Every endpoint that mutates a cached object must call this: auth serves those objects
cache-first with no freshness check, so a mutation that leaves the entry in place keeps the
stale object enforced until its TTL expires (LIT-3803). Best-effort on both steps: the DB write
has already committed, so a cache backend error must not fail the endpoint.
stale object enforced until its TTL expires (LIT-3803). Best-effort: the DB write has already
committed, so a cache backend error must not fail the endpoint.
"""
for cache_key in cache_keys:
try:

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

76
litellm/proxy/policy_engine/policy_matcher.py
View file

@ -12,6 +12,7 @@ from typing import Final
from litellm._logging import verbose_proxy_logger
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
from litellm.types.proxy.policy_engine import Policy, PolicyMatchContext, PolicyScope
@ -136,14 +137,46 @@ class PolicyMatcher:
context: PolicyMatchContext,
policies: dict[str, Policy] | None = None,
) -> Callable[[str], bool]:
"""Predicate telling whether a policy exists and its condition matches the context."""
"""
Predicate telling whether a policy exists and any policy in its
inheritance chain applies to the context. Admissions where the
policy's own condition missed but an ancestor applies are logged at
INFO, once per attachment scan.
"""
resolved: Final = policies if policies is not None else PolicyMatcher._registry_policies()
return lambda policy_name: bool(
PolicyMatcher.get_policies_with_matching_conditions(
policy_names=(policy_name,),
context=context,
policies=resolved,
def applies(policy_name: str) -> bool:
applying: Final = PolicyMatcher._applying_chain_members(
policy_name=policy_name, context=context, policies=resolved
)
if applying and policy_name not in applying:
verbose_proxy_logger.info(
"Policy '%s' applied through ancestor '%s' although its own condition did not match "
"(team_alias=%s, key_alias=%s, model=%s)",
policy_name,
applying[0],
context.team_alias,
context.key_alias,
context.model,
)
return bool(applying)
return applies
@staticmethod
def _applying_chain_members(
policy_name: str,
context: PolicyMatchContext,
policies: dict[str, Policy],
) -> tuple[str, ...]:
from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator
chain: Final = PolicyResolver.resolve_inheritance_chain(policy_name=policy_name, policies=policies)
return tuple(
name
for name in chain
if (policy := policies.get(name)) is not None
and (policy.condition is None or ConditionEvaluator.evaluate(policy.condition, context))
)
@staticmethod
@ -160,11 +193,14 @@ class PolicyMatcher:
policies: dict[str, Policy] | None = None,
) -> list[str]:
"""
Filter policies to only those whose conditions match the context.
Filter policies to only those that apply to the given context.
A policy's condition matches if:
- The policy has no condition (condition is None), OR
- The policy's condition evaluates to True for the given context
A policy applies when any policy in its inheritance chain has no
condition or a condition that evaluates to True for the context. The
resolver then drops only the chain members whose own condition fails,
so a child whose condition misses still contributes the guardrails of
its unconditional ancestors. A missing policy resolves to an empty
chain and does not apply.
Args:
policy_names: List of policy names to filter
@ -172,19 +208,11 @@ class PolicyMatcher:
policies: Dictionary of all policies (if None, uses global registry)
Returns:
List of policy names whose conditions match the context
List of policy names that apply to the context
"""
from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator
resolved: Final = policies if policies is not None else PolicyMatcher._registry_policies()
matching_policies: Final = []
for policy_name in policy_names:
policy = resolved.get(policy_name)
if policy is None:
continue
# Policy matches if it has no condition OR condition evaluates to True
if policy.condition is None or ConditionEvaluator.evaluate(policy.condition, context):
matching_policies.append(policy_name)
return matching_policies
return [
policy_name
for policy_name in policy_names
if PolicyMatcher._applying_chain_members(policy_name, context, resolved)
]

6
litellm/proxy/policy_engine/policy_resolver.py
View file

@ -210,6 +210,7 @@ class PolicyResolver:
Returns:
List of (policy_name, GuardrailPipeline) tuples
"""
from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
@ -230,6 +231,11 @@ class PolicyResolver:
policy = policies.get(policy_name)
if policy is None:
continue
if policy.condition is not None and not ConditionEvaluator.evaluate(
condition=policy.condition, context=context
):
verbose_proxy_logger.debug("Policy '%s' condition did not match, skipping pipeline", policy_name)
continue
if policy.pipeline is not None:
pipelines.append((policy_name, policy.pipeline))
verbose_proxy_logger.debug(

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:

43
litellm/proxy/vector_store_endpoints/management_endpoints.py
View file

@ -13,6 +13,7 @@ import json
from typing import TYPE_CHECKING, Any, Final
from fastapi import APIRouter, Depends, HTTPException
from typing_extensions import ReadOnly, TypedDict
if TYPE_CHECKING:
from prisma.models import LiteLLM_ManagedVectorStoresTable as _VectorStoreRow
@ -56,6 +57,32 @@ def _row_to_vector_store(row: "_VectorStoreRow") -> LiteLLM_ManagedVectorStore:
return LiteLLM_ManagedVectorStore(**row.model_dump())
class _ConfigOwnedDetail(TypedDict):
error: ReadOnly[str]
vector_store_id: ReadOnly[str]
def _raise_if_config_owned(vector_store_id: str) -> None:
if litellm.vector_store_registry is None or not litellm.vector_store_registry.is_config_vector_store(
vector_store_id
):
return
detail: Final[_ConfigOwnedDetail] = {
"error": (
f"Vector store {vector_store_id} is defined in the config file, so the config file owns it and it "
"cannot be changed here. Edit the config file to change it, or remove it from the file to let the "
"database own it."
),
"vector_store_id": vector_store_id,
}
raise HTTPException(status_code=400, detail=detail)
def _with_ownership(vector_store: LiteLLM_ManagedVectorStore) -> LiteLLM_ManagedVectorStore:
ownership: Final = LiteLLM_ManagedVectorStore(is_config=vector_store.get("is_config", False))
return vector_store | ownership
_LITELLM_PARAMS_MASKER: Final = SensitiveDataMasker(extra_sensitive_patterns=frozenset(("connection",)))
@ -274,6 +301,7 @@ async def new_vector_store(
status_code=400,
detail="vector_store_id and custom_llm_provider are required",
)
_raise_if_config_owned(vector_store_id)
# Extract and validate metadata
metadata: Final = vector_store.get("vector_store_metadata")
@ -306,6 +334,8 @@ async def new_vector_store(
"message": f"Vector store {vector_store.get('vector_store_id')} created successfully",
"vector_store": response_vs,
}
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception("Error creating vector store: %s", e)
raise HTTPException(status_code=500, detail=str(e))
@ -331,7 +361,9 @@ async def list_vector_stores(
"""
List all available vector stores with optional filtering and pagination.
Combines both in-memory vector stores and those stored in the database.
Database is the source of truth - deleted stores are removed from memory, updated stores sync to memory.
Database is the source of truth for stores it owns: deleted stores are removed from memory, updated stores
sync to memory. Stores declared in the config file are owned by the config file, are always listed, and are
never overwritten by database rows.
Parameters:
- page: int - Page number for pagination (default: 1)
@ -366,8 +398,10 @@ async def list_vector_stores(
if not vector_store_id:
continue
if vector_store.get("is_config", False):
vector_store_map[vector_store_id] = vector_store
# If vector store is in memory but NOT in database, it was deleted
if vector_store_id not in db_vector_store_ids:
elif vector_store_id not in db_vector_store_ids:
verbose_proxy_logger.info(
"Vector store %s exists in memory but not in database - marking for deletion from cache",
vector_store_id,
@ -394,7 +428,7 @@ async def list_vector_stores(
# Filter vector stores based on access control
accessible_vector_stores: Final = []
for vs in await filter_listable_vector_stores(vector_store_map.values(), user_api_key_dict):
redacted = LiteLLM_ManagedVectorStore(**vs)
redacted = _with_ownership(vs)
redacted["litellm_params"] = _redact_sensitive_litellm_params(vs.get("litellm_params"))
accessible_vector_stores.append(redacted)
@ -467,6 +501,7 @@ async def delete_vector_store(
status_code=404,
detail=f"Vector store with ID {data.vector_store_id} not found",
)
_raise_if_config_owned(data.vector_store_id)
# Check access control
if vector_store_to_check and not await _check_vector_store_access(vector_store_to_check, user_api_key_dict):
@ -545,6 +580,7 @@ async def get_vector_store_info(
litellm_params=_redact_sensitive_litellm_params(vector_store.get("litellm_params")),
team_id=vector_store.get("team_id") or None,
user_id=vector_store.get("user_id") or None,
is_config=vector_store.get("is_config", False),
)
return {"vector_store": vector_store_pydantic_obj}
@ -591,6 +627,7 @@ async def update_vector_store(
update_data: Final = data.model_dump(exclude_unset=True)
vector_store_id: Final[str] = data.vector_store_id
update_data.pop("vector_store_id")
_raise_if_config_owned(vector_store_id)
# Per-store access control: anyone authenticated who passes the
# premium-feature gate could otherwise update *any* vector store —

2
litellm/types/vector_stores.py
View file

@ -44,6 +44,8 @@ class LiteLLM_ManagedVectorStore(TypedDict, total=False):
team_id: str | None
user_id: str | None
is_config: ReadOnly[bool]
class LiteLLM_ManagedVectorStoreListResponse(TypedDict, total=False):
"""Response format for listing vector stores"""

12
litellm/vector_stores/vector_store_registry.py
View file

@ -340,7 +340,7 @@ class VectorStoreRegistry:
# Verify vector store still exists in database (if we have DB access)
# This ensures deleted vector stores are removed from cache
if vector_store is not None and prisma_client is not None:
if vector_store is not None and prisma_client is not None and not vector_store.get("is_config", False):
try:
# Check if it still exists in database
db_vector_store = await ManagedVectorStoresRepository(prisma_client).table.find_unique(
@ -426,6 +426,7 @@ class VectorStoreRegistry:
vector_store_metadata=vector_store_litellm_params.get("vector_store_metadata"),
created_at=datetime.now(timezone.utc),
updated_at=datetime.now(timezone.utc),
is_config=True,
)
self.vector_stores.append(litellm_managed_vector_store)
@ -452,6 +453,10 @@ class VectorStoreRegistry:
return response
def is_config_vector_store(self, vector_store_id: str) -> bool:
vector_store: Final = self.get_litellm_managed_vector_store_from_registry(vector_store_id=vector_store_id)
return vector_store is not None and vector_store.get("is_config", False)
def add_vector_store_to_registry(self, vector_store: LiteLLM_ManagedVectorStore):
"""
Add a vector store to the registry
@ -475,10 +480,11 @@ class VectorStoreRegistry:
]
def update_vector_store_in_registry(self, vector_store_id: str, updated_data: LiteLLM_ManagedVectorStore):
"""Update or add a vector store in the registry"""
"""Update or add a vector store in the registry. Config-defined stores are left untouched"""
for i, vector_store in enumerate(self.vector_stores):
if vector_store.get("vector_store_id") == vector_store_id:
self.vector_stores[i] = updated_data
if not vector_store.get("is_config", False):
self.vector_stores[i] = updated_data
return
self.vector_stores.append(updated_data)

675
model_prices_and_context_window.json
View file

File diff suppressed because it is too large Load diff

74
tests/e2e/guardrails/guardrails_client.py
View file

@ -17,6 +17,7 @@ from models import (
AnthropicMessagesResponse,
ChatBody,
ChatMessage,
ChatMetadata,
ChatResponse,
ChatTool,
KeyGenerateBody,
@ -133,6 +134,31 @@ class GuardrailCreateResponse(BaseModel):
guardrail_id: str
class PolicyConditionBody(BaseModel):
model: str
class PolicyCreateBody(BaseModel):
policy_name: str
inherit: str | None = None
guardrails_add: list[str]
condition: PolicyConditionBody | None = None
class PolicyCreateResponse(BaseModel):
policy_id: str
policy_name: str
class PolicyAttachmentCreateBody(BaseModel):
policy_name: str
tags: list[str]
class PolicyAttachmentCreateResponse(BaseModel):
attachment_id: str
class ApplyGuardrailRequest(BaseModel):
guardrail_name: str
text: str
@ -243,6 +269,49 @@ class GuardrailsClient:
response_type=NoBody,
)
def create_policy(self, body: PolicyCreateBody) -> str:
"""Create a policy via POST /policies and return its name once every replica
can be expected to serve it (policies reach the data plane on the periodic
DB sync, same as guardrails)."""
created = unwrap(
self.proxy.transport.post(
"/policies",
headers=self.proxy.transport.master,
json=body,
response_type=PolicyCreateResponse,
)
)
settle_propagation(time.monotonic())
return created.policy_name
def delete_policy(self, policy_name: str) -> None:
_ = self.proxy.transport.delete(
f"/policies/name/{policy_name}/all-versions",
headers=self.proxy.transport.master,
json=NoBody(),
response_type=NoBody,
)
def attach_policy_to_tags(self, policy_name: str, tags: list[str]) -> str:
attachment_id = unwrap(
self.proxy.transport.post(
"/policies/attachments",
headers=self.proxy.transport.master,
json=PolicyAttachmentCreateBody(policy_name=policy_name, tags=tags),
response_type=PolicyAttachmentCreateResponse,
)
).attachment_id
settle_propagation(time.monotonic())
return attachment_id
def delete_policy_attachment(self, attachment_id: str) -> None:
_ = self.proxy.transport.delete(
f"/policies/attachments/{attachment_id}",
headers=self.proxy.transport.master,
json=NoBody(),
response_type=NoBody,
)
def create_team_opted_out_of_global_guardrails(self, alias: str) -> str:
team_id = unwrap(
self.proxy.transport.post(
@ -322,11 +391,13 @@ class GuardrailsClient:
max_tokens: int = 16,
tools: list[ChatTool] | None = None,
tool_choice: str | None = None,
tags: list[str] | None = None,
) -> StreamingResponse:
"""Drive /chat/completions returning the raw HTTP outcome, for the
assertions a typed body cannot carry: the `x-litellm-applied-guardrails`
response header, which is how an ALLOW scenario proves the guardrail ran
rather than being absent."""
rather than being absent. `tags` land in `metadata.tags`, which is what a
tag-scoped policy attachment matches on."""
return self.proxy.transport.send(
"/chat/completions",
headers=self.proxy.transport.bearer(key),
@ -337,6 +408,7 @@ class GuardrailsClient:
guardrails=guardrails,
tools=tools,
tool_choice=tool_choice,
metadata=ChatMetadata(tags=tags) if tags is not None else None,
),
)

127
tests/e2e/guardrails/test_policy_inherited_guardrail_e2e.py Normal file
View file

@ -0,0 +1,127 @@
"""Live e2e: a policy attached to a request keeps its inherited parent guardrails
when only the child's own `condition` fails to match the request model.
The parent policy has no condition and adds a content filter. The child inherits
it, adds a second content filter, and carries a model condition. The attachment
points at the child only, so the parent is reachable through inheritance alone.
A request the child condition does not match must still be blocked by the
parent's filter; a request it does match must be blocked by both.
Uses litellm_content_filter (keyword match, no external service) so the block is
deterministic and free, with the request model routed to a real provider.
"""
from __future__ import annotations
import pytest
from e2e_config import CHEAP_OPENAI_MODEL, unique_marker
from e2e_http import StreamingResponse
from guardrails_client import (
GuardrailsClient,
PolicyConditionBody,
PolicyCreateBody,
)
from lifecycle import ResourceManager
pytestmark = pytest.mark.e2e
MODEL = CHEAP_OPENAI_MODEL
def _applied_guardrails(outcome: StreamingResponse) -> frozenset[str]:
return frozenset(
name.strip() for name in outcome.headers.get("x-litellm-applied-guardrails", "").split(",") if name.strip()
)
def _setup_child_policy_attached_to_tag(
client: GuardrailsClient,
resources: ResourceManager,
*,
child_condition_model: str,
parent_banned: str,
child_banned: str,
tag: str,
) -> tuple[str, str]:
"""Register parent and child content filters, a parent policy adding the parent
filter, a child policy inheriting it with `child_condition_model`, and attach
only the child to `tag`. Returns (parent_guardrail_name, child_guardrail_name)."""
parent_guardrail = f"e2e-parent-guard-{parent_banned}"
child_guardrail = f"e2e-child-guard-{child_banned}"
parent_guardrail_id = client.create_content_filter_guardrail(parent_guardrail, parent_banned, default_on=False)
resources.defer(lambda: client.delete_guardrail(parent_guardrail_id))
child_guardrail_id = client.create_content_filter_guardrail(child_guardrail, child_banned, default_on=False)
resources.defer(lambda: client.delete_guardrail(child_guardrail_id))
parent_policy = client.create_policy(
PolicyCreateBody(policy_name=f"e2e-parent-policy-{parent_banned}", guardrails_add=[parent_guardrail])
)
resources.defer(lambda: client.delete_policy(parent_policy))
child_policy = client.create_policy(
PolicyCreateBody(
policy_name=f"e2e-child-policy-{child_banned}",
inherit=parent_policy,
guardrails_add=[child_guardrail],
condition=PolicyConditionBody(model=child_condition_model),
)
)
resources.defer(lambda: client.delete_policy(child_policy))
attachment_id = client.attach_policy_to_tags(child_policy, [tag])
resources.defer(lambda: client.delete_policy_attachment(attachment_id))
return parent_guardrail, child_guardrail
class TestPolicyInheritedGuardrail:
def test_child_condition_miss_still_applies_inherited_parent_guardrail(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
parent_banned = unique_marker()
child_banned = unique_marker()
tag = f"e2e-policy-tag-{unique_marker()}"
parent_guardrail, child_guardrail = _setup_child_policy_attached_to_tag(
client,
resources,
child_condition_model=f"never-matches-{unique_marker()}",
parent_banned=parent_banned,
child_banned=child_banned,
tag=tag,
)
outcome = client.chat_raw(scoped_key, MODEL, f"Reply with the single word OK. {parent_banned}", tags=[tag])
assert outcome.status_code == 400, (
f"the inherited parent content filter must block the banned keyword even though the child "
f"policy's own model condition does not match {MODEL}; got {outcome.status_code}: {outcome.body[:300]}"
)
assert parent_guardrail in _applied_guardrails(outcome), (
f"x-litellm-applied-guardrails must name the inherited parent guardrail; got {outcome.headers}"
)
assert child_guardrail not in _applied_guardrails(outcome), (
f"the child's own guardrail must not run when its condition fails; got {outcome.headers}"
)
def test_child_condition_match_applies_child_and_inherited_parent_guardrails(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
parent_banned = unique_marker()
child_banned = unique_marker()
tag = f"e2e-policy-tag-{unique_marker()}"
parent_guardrail, child_guardrail = _setup_child_policy_attached_to_tag(
client,
resources,
child_condition_model=MODEL,
parent_banned=parent_banned,
child_banned=child_banned,
tag=tag,
)
outcome = client.chat_raw(scoped_key, MODEL, f"Reply with the single word OK. {child_banned}", tags=[tag])
assert outcome.status_code == 400, (
f"the child's own content filter must block its banned keyword when the condition matches {MODEL}; "
f"got {outcome.status_code}: {outcome.body[:300]}"
)
assert {parent_guardrail, child_guardrail} <= _applied_guardrails(outcome), (
f"both the child and inherited parent guardrails must run; got {outcome.headers}"
)

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`

16
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
@ -12,9 +12,9 @@ The generated lifecycle models use 20 examples, eight steps, generation and shri
Reuse the existing canned provider handlers through `_support/upstream.py`. It rejects internal request fields and exposes actual received requests for independent assertions. Register every created resource for cleanup immediately, keep expected values independent of production calculations, and assert readback plus the runtime effect of a change
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, skipped tests, failed cleanup or a selected test without a passed call fail qualification. Existing GitHub Actions jobs do not own these tests
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 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
@ -22,7 +22,7 @@ Fixtures must contain synthetic data only. Keep private incident records and sou
Database cases own their temporary schemas, roles, constraints and proxy processes. They prove reader-versus-writer execution with PostgreSQL lock observations, exercise real transaction wait limits and verify rollback after a reached database failure
Accounting cases compare persisted input and output cost components against literal rates, including zero and default prices. Cache state models assert actual upstream calls, response identity and every persisted charge. Generated accounting tests have a 180-second test limit to accommodate the asynchronous spend writer; CircleCI keeps the whole shard capped at 11 minutes
Accounting cases compare persisted input and output cost components against literal rates, including zero and default prices. Cache state models assert actual upstream calls, response identity and every persisted charge. Generated accounting tests have a 180-second test limit to accommodate the asynchronous spend writer
Provider contracts exercise actual TCP requests with synthetic credentials and local protocol peers. The S3 verifier uses independently implemented equations, a published known-answer vector, a fixed signing clock and deliberately invalid signed requests. Bedrock cases clear ambient AWS credential sources and check the literal model path, loaded role references, STS requests and bearer-only behavior
@ -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]

11
tests/integration/_support/process.py
View file

@ -1,6 +1,6 @@
import os
import socket
import signal
import socket
import subprocess
import sys
import time
@ -13,7 +13,6 @@ from typing import Final
import httpx
import psutil
from integration._support.client import Gateway
@ -61,9 +60,10 @@ def owned_proxy(
*,
config: Path | None = None,
remove_environment: tuple[str, ...] = (),
workers: int = 1,
) -> Iterator[Gateway]:
with owned_proxy_process(
gateway, directory, overrides, config=config, remove_environment=remove_environment
gateway, directory, overrides, config=config, remove_environment=remove_environment, workers=workers
) as owned:
yield owned.gateway
@ -76,11 +76,12 @@ def owned_proxy_process(
*,
config: Path | None = None,
remove_environment: tuple[str, ...] = (),
workers: int = 1,
) -> Iterator[OwnedProxy]:
with socket.socket() as reserve:
reserve.bind(("127.0.0.1", 0))
port: Final = reserve.getsockname()[1]
root: Final = Path(__file__).resolve().parents[3]
root: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3])
environment: Final = {
**{name: value for name, value in os.environ.items() if name not in remove_environment},
"LITELLM_MASTER_KEY": gateway.key,
@ -104,7 +105,7 @@ def owned_proxy_process(
"--port",
str(port),
"--num_workers",
"1",
str(workers),
"--use_prisma_db_push",
"--enforce_prisma_migration_check",
],

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

26
tests/integration/_support/upstream.py
View file

@ -152,6 +152,31 @@ class Provider:
)
return await chat_completions(request)
async def vector_store_search(self, request: Request) -> Response:
body: Final = JSON_OBJECT.validate_json(await request.body())
self.observations.put(Observation(request.url.path, request.headers.get("authorization", ""), body))
query: Final = body.get("query")
if not isinstance(query, str) or not query:
return JSONResponse({"error": {"message": "query is required"}}, status_code=400)
vector_store_id: Final = cast(str, request.path_params["vector_store_id"])
return JSONResponse(
{
"object": "vector_store.search_results.page",
"search_query": query,
"data": [
{
"file_id": f"file_{vector_store_id}",
"filename": "scripted.txt",
"score": 0.9,
"attributes": {},
"content": [{"type": "text", "text": f"scripted context for {query}"}],
}
],
"has_more": False,
"next_page": None,
}
)
async def script(self, request: Request) -> Response:
name: Final = cast(str, request.path_params["model"])
if request.method in {"DELETE", "GET"} and name not in self.scripts:
@ -338,6 +363,7 @@ class Provider:
Route("/v1/completions", completions, methods=["POST"]),
Route("/v1/embeddings", embeddings, methods=["POST"]),
Route("/v1/moderations", moderations, methods=["POST"]),
Route("/vector_stores/{vector_store_id}/search", self.vector_store_search, methods=["POST"]),
Route("/{path:path}", self.scripted, methods=["POST"]),
Route("/{path:path}", self.scripted, methods=["GET"]),
WebSocketRoute("/v1/realtime", self.realtime),

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

@ -2,11 +2,13 @@ from __future__ import annotations
import ssl
import threading
import time
from collections.abc import Callable, Generator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from queue import SimpleQueue
from types import MappingProxyType
from typing import Final
@ -26,6 +28,8 @@ class Reply:
chunks: tuple[bytes, ...] | None = None
abort_after: int | None = None
gate_after_first: threading.Event | None = None
pause_between_chunks: float = 0
headers: Mapping[str, str] = MappingProxyType({})
@dataclass(frozen=True, slots=True)
@ -64,6 +68,8 @@ def wire_server(respond: Callable[[Request], Reply], tls: ssl.SSLContext | None
reply = Reply(status=500)
self.send_response(reply.status)
self.send_header("content-type", reply.content_type)
for name, value in reply.headers.items():
self.send_header(name, value)
if reply.chunks is None:
self.send_header("content-length", str(len(reply.body)))
else:
@ -81,6 +87,8 @@ def wire_server(respond: Callable[[Request], Reply], tls: ssl.SSLContext | None
self.wfile.flush()
if index == 0 and reply.gate_after_first is not None:
assert reply.gate_after_first.wait(timeout=5), "Stream barrier was never released"
if reply.pause_between_chunks and index + 1 < len(reply.chunks):
time.sleep(reply.pause_between_chunks)
else:
self.wfile.write(b"0\r\n\r\n")
self.wfile.flush()

44
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)
@ -81,27 +75,35 @@ def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None:
collected: Final = session.config.stash.get(COLLECTED, ())
reports: Final = tuple(report for report in session.config.stash[REPORTS] if report.nodeid in collected)
passed: Final = tuple(report.nodeid for report in reports if report.when == "call" and report.passed)
skipped: Final = tuple(report.nodeid for report in reports if report.skipped)
complete: Final = (
exitstatus == 0
and bool(collected)
and sorted(collected) == sorted(passed)
and all(report.passed for report in reports)
and sorted(collected) == sorted(passed + skipped)
and not any(report.failed for report in reports)
)
output: Final = Path(destination)
output.mkdir(parents=True, exist_ok=True)
(output / "execution.json").write_text(
json.dumps({
"collected": collected, "passed": passed, "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:

1813
tests/integration/contracts.json
View file

File diff suppressed because it is too large Load diff

12
tests/integration/coordination_redis_proxy_config.yaml Normal file
View file

@ -0,0 +1,12 @@
model_list: []
general_settings:
master_key: os.environ/LITELLM_MASTER_KEY
database_url: os.environ/DATABASE_URL
store_model_in_db: true
disable_spend_logs: false
proxy_batch_write_at: 1
coordination_redis:
host: os.environ/REDIS_HOST
port: os.environ/REDIS_PORT
router_settings:
disable_cooldowns: true

29
tests/integration/management/test_budget_updates.py Normal file
View file

@ -0,0 +1,29 @@
from datetime import datetime, timedelta, timezone
from typing import Final
import pytest
from tests.integration._support.client import Gateway, string_value
from tests.integration._support.database import read_rows
def _persisted_reset_at(budget_id: str) -> datetime:
rows: Final = read_rows(
'SELECT budget_reset_at::text AS reset_at FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (budget_id,)
)
assert len(rows) == 1, rows
reset_at: Final = datetime.fromisoformat(string_value(rows[0]["reset_at"]))
return reset_at if reset_at.tzinfo is not None else reset_at.replace(tzinfo=timezone.utc)
@pytest.mark.covers("mgmt.budget.update.duration_change_recomputes_reset_at")
def test_shortening_budget_duration_moves_reset_at_onto_the_new_schedule(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
budget_id: Final = scenario.budget(max_budget=10.0, budget_duration="10d")
ten_day_reset_at: Final = _persisted_reset_at(budget_id)
before: Final = datetime.now(timezone.utc)
response: Final = gateway.request("POST", "/budget/update", {"budget_id": budget_id, "budget_duration": "1d"})
assert response.status_code == 200, response.text
updated: Final = _persisted_reset_at(budget_id)
assert updated < ten_day_reset_at, f"{updated} not before {ten_day_reset_at}"
assert before < updated <= before + timedelta(days=1, minutes=5), f"{updated} not within 1d of {before}"

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

59
tests/integration/management/test_organization_budget_clear.py Normal file
View file

@ -0,0 +1,59 @@
import uuid
from typing import Final
import pytest
from tests.integration._support.client import Gateway, object_value, string_value
from tests.integration._support.database import read_rows
def _budget_rows(budget_id: str) -> list[dict[str, object]]:
return read_rows(
'SELECT tpm_limit, rpm_limit, max_budget FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (budget_id,)
)
@pytest.mark.covers("mgmt.organization.update.null_clears_budget_limit")
def test_patch_organization_update_with_null_tpm_limit_clears_it_and_keeps_sibling_limits(gateway: Gateway) -> None:
created: Final = gateway.post(
"/organization/new",
{
"organization_alias": f"integration-{uuid.uuid4().hex}",
"tpm_limit": 4000,
"rpm_limit": 40,
"max_budget": 12.5,
},
)
organization_id: Final = string_value(created["organization_id"])
budget_id: Final = string_value(created["budget_id"])
try:
assert _budget_rows(budget_id) == [{"tpm_limit": 4000, "rpm_limit": 40, "max_budget": 12.5}]
updated: Final = gateway.request(
"PATCH", "/organization/update", {"organization_id": organization_id, "tpm_limit": None}
)
assert updated.status_code == 200, updated.text
updated_budget: Final = object_value(object_value(updated.json())["litellm_budget_table"])
assert (updated_budget["tpm_limit"], updated_budget["rpm_limit"], updated_budget["max_budget"]) == (
None,
40,
12.5,
), updated.text
assert _budget_rows(budget_id) == [{"tpm_limit": None, "rpm_limit": 40, "max_budget": 12.5}]
info: Final = gateway.request("GET", "/organization/info", params={"organization_id": organization_id})
assert info.status_code == 200, info.text
info_budget: Final = object_value(object_value(info.json())["litellm_budget_table"])
assert (info_budget["tpm_limit"], info_budget["rpm_limit"], info_budget["max_budget"]) == (
None,
40,
12.5,
), info.text
finally:
deleted: Final = gateway.request("DELETE", "/organization/delete", {"organization_ids": [organization_id]})
assert deleted.status_code == 200, deleted.text
gateway.post("/budget/delete", {"id": budget_id})
assert (
read_rows(
'SELECT organization_id FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (organization_id,)
)
== []
)

61
tests/integration/management/test_team_budget_duration_defaults.py Normal file
View file

@ -0,0 +1,61 @@
import uuid
from pathlib import Path
from typing import Final
import pytest
import yaml
from pydantic import JsonValue
from tests.integration._support.client import Gateway, string_value
from tests.integration._support.database import read_rows
from tests.integration._support.process import owned_proxy
def _budget_row(team_id: str) -> dict[str, JsonValue]:
rows: Final = read_rows(
'SELECT max_budget, budget_duration, budget_reset_at::text FROM "LiteLLM_TeamTable" WHERE team_id = %s',
(team_id,),
)
assert len(rows) == 1, rows
return rows[0]
@pytest.mark.covers("mgmt.team.new.explicit_null_budget_duration_overrides_default")
def test_team_new_explicit_null_budget_duration_is_not_replaced_by_default(gateway: Gateway, tmp_path: Path) -> None:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["litellm_settings"]["default_team_params"] = {"budget_duration": "30d"}
path: Final = tmp_path / "team-defaults.yaml"
path.write_text(yaml.safe_dump(config))
with (
owned_proxy(gateway, tmp_path, {"STORE_MODEL_IN_DB": "False"}, config=path) as candidate,
candidate.scenario() as scenario,
):
never_resetting: Final = candidate.request(
"POST",
"/team/new",
{"team_alias": f"integration-{uuid.uuid4().hex}", "max_budget": 500, "budget_duration": None},
)
assert never_resetting.status_code == 200, never_resetting.text
never_resetting_id: Final = string_value(never_resetting.json()["team_id"])
scenario.cleanups.callback(scenario.delete_team, never_resetting_id)
assert never_resetting.json()["max_budget"] == 500.0, never_resetting.text
assert never_resetting.json()["budget_duration"] is None, never_resetting.text
assert never_resetting.json()["budget_reset_at"] is None, never_resetting.text
assert _budget_row(never_resetting_id) == {
"max_budget": 500.0,
"budget_duration": None,
"budget_reset_at": None,
}
inheriting: Final = candidate.request(
"POST", "/team/new", {"team_alias": f"integration-{uuid.uuid4().hex}", "max_budget": 500}
)
assert inheriting.status_code == 200, inheriting.text
inheriting_id: Final = string_value(inheriting.json()["team_id"])
scenario.cleanups.callback(scenario.delete_team, inheriting_id)
assert inheriting.json()["budget_duration"] == "30d", inheriting.text
assert inheriting.json()["budget_reset_at"] is not None, inheriting.text
inheriting_row: Final = _budget_row(inheriting_id)
assert inheriting_row["max_budget"] == 500.0, inheriting_row
assert inheriting_row["budget_duration"] == "30d", inheriting_row
assert inheriting_row["budget_reset_at"] is not None, inheriting_row

40
tests/integration/management/test_team_member_budget_cache.py Normal file
View file

@ -0,0 +1,40 @@
import os
from typing import Final
import pytest
from pydantic import JsonValue, TypeAdapter
from redis import Redis
from tests.integration._support.client import Gateway, eventually, object_value, string_value
from tests.integration._support.database import read_rows
_CACHED_BUDGET: Final = TypeAdapter(dict[str, JsonValue])
@pytest.mark.covers("mgmt.team_member_budget.default_budget_is_cached_in_redis_as_json")
def test_team_member_default_budget_lands_in_redis_after_first_member_call(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model()
user: Final = scenario.user()
team: Final = scenario.team(team_member_budget=25)
key: Final = scenario.key(team_id=team, user_id=user, models=[model])
teams: Final = read_rows('SELECT metadata FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team,))
assert len(teams) == 1, teams
budget_id: Final = string_value(object_value(teams[0]["metadata"])["team_member_budget_id"])
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": "member budget cache"}]},
key=key,
)
assert response.status_code == 200, response.text
with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache:
cached: Final = eventually(
lambda: cache.get(f"team_member_default_budget:{budget_id}"),
lambda value: value is not None,
seconds=10,
)
assert isinstance(cached, bytes), cached
budget: Final = _CACHED_BUDGET.validate_json(cached)
assert budget["budget_id"] == budget_id, cached
assert budget["max_budget"] == 25, cached

379
tests/integration/management/test_user_updates_wedged_coordination_redis.py Normal file
View file

@ -0,0 +1,379 @@
import os
import signal
import time
import uuid
from collections.abc import Callable, Mapping
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Final
from urllib.parse import urlsplit, urlunsplit
import psutil
import psycopg
import pytest
from psycopg import sql
from pydantic import JsonValue
from redis import Redis
from redis.client import PubSub
from tests.integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value
from tests.integration._support.process import owned_proxy
from tests.integration._support.redis_process import owned_redis
_USERS: Final = 60
_BURST: Final = 30
_HANDLER_BUDGET_SECONDS: Final = 0.75
_BULK_BUDGET_SECONDS: Final = 2.0
_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation"
def _timed_post(candidate: Gateway, path: str, body: Mapping[str, JsonValue], timeout: float = 15) -> float:
started: Final = time.monotonic()
response: Final = candidate.client.request(
"POST",
path,
json=body,
headers={"Authorization": f"Bearer {candidate.key}"},
timeout=timeout,
)
elapsed: Final = time.monotonic() - started
assert response.status_code == 200, f"POST {path}: {response.status_code} {response.text} after {elapsed:.3f}s"
return elapsed
def _received(pubsub: PubSub) -> tuple[dict[str, JsonValue], ...]:
messages: list[dict[str, JsonValue]] = []
while True:
message = pubsub.get_message(ignore_subscribe_messages=True, timeout=0)
if message is None:
return tuple(messages)
data = message.get("data")
if isinstance(data, (bytes, str)):
messages.append(JSON_OBJECT.validate_json(data))
def _worker_pid(port: int) -> int:
for process in psutil.process_iter():
parent = process.parent()
if parent is None:
continue
try:
cmdline = parent.cmdline()
own_cmdline = process.cmdline()
except (psutil.NoSuchProcess, psutil.AccessDenied):
continue
if (
"integration._support.proxy" in cmdline
and "--port" in cmdline
and str(port) in cmdline
and not any("prisma" in part for part in own_cmdline)
):
return process.pid
raise AssertionError(f"no uvicorn worker found under the owned proxy on port {port}")
def _burst_call(
index: int, users: tuple[str, ...], key: str, team_id: str, customer_id: str
) -> tuple[str, dict[str, JsonValue]]:
match index % 5:
case 0:
return "/user/update", {"user_id": users[index], "max_budget": 200.0 + index}
case 1:
return "/user/update", {"user_id": users[index], "tpm_limit": 1000 + index}
case 2:
return "/key/update", {"key": key, "max_budget": 7.0 + index}
case 3:
return "/team/update", {"team_id": team_id, "max_budget": 7.0 + index}
case _:
return "/customer/update", {"user_id": customer_id, "max_budget": 7.0 + index}
@pytest.mark.timeout(240)
@pytest.mark.covers(
"mgmt.user.update.budget_change_returns_promptly_with_wedged_coordination_redis",
"mgmt.user.bulk_update.budget_change_returns_promptly_with_wedged_coordination_redis",
"mgmt.customer.update.budget_change_returns_promptly_with_wedged_coordination_redis",
"mgmt.key.reset_spend.returns_promptly_with_wedged_coordination_redis",
"mgmt.auth_cache_invalidation.publish_parked_by_short_redis_wedge_lands_after_recovery",
"mgmt.auth_cache_invalidation.burst_with_worker_kill_keeps_serving_while_redis_wedged",
)
def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged(
gateway: Gateway, tmp_path: Path, record_property: Callable[[str, object], None]
) -> None:
original: Final = os.environ["DATABASE_URL"]
identity: Final = "integration_wedged_redis_" + uuid.uuid4().hex
parsed: Final = urlsplit(original)
database_url: Final = urlunsplit((parsed.scheme, parsed.netloc, "/" + identity, "", ""))
timings: dict[str, float] = {}
with psycopg.connect(original, autocommit=True) as admin:
admin.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(identity)))
try:
results_dir: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(tmp_path)))
prior_logs: Final = frozenset(results_dir.glob("owned-proxy-*.log"))
with (
owned_redis(tmp_path) as coordination,
owned_proxy(
gateway,
tmp_path,
{
"DATABASE_URL": database_url,
"REDIS_HOST": coordination.host,
"REDIS_PORT": str(coordination.port),
},
config=Path("tests/integration/coordination_redis_proxy_config.yaml"),
workers=2,
) as candidate,
Redis(host=coordination.host, port=coordination.port, socket_timeout=1) as subscriber_client,
):
pubsub: Final = subscriber_client.pubsub()
pubsub.subscribe(_CHANNEL)
received: list[dict[str, JsonValue]] = []
def drained() -> tuple[dict[str, JsonValue], ...]:
received.extend(_received(pubsub))
return tuple(received)
eventually(
lambda: subscriber_client.pubsub_numsub(_CHANNEL)[0][1],
lambda count: count >= 3,
seconds=15,
)
users: Final = tuple(f"{identity}_u{index}" for index in range(_USERS))
for user_id in users:
candidate.post("/user/new", {"user_id": user_id, "auto_create_key": False, "max_budget": 10.0})
key: Final = string_value(
candidate.post("/key/generate", {"user_id": users[0], "max_budget": 5.0})["key"]
)
team_id: Final = string_value(
candidate.post("/team/new", {"team_alias": identity, "max_budget": 5.0})["team_id"]
)
customer_id: Final = identity + "_cust"
candidate.post("/customer/new", {"user_id": customer_id, "max_budget": 5.0})
drained()
timings["h1_healthy"] = _timed_post(
candidate, "/user/update", {"user_id": users[0], "max_budget": 11.0}
)
assert timings["h1_healthy"] < _HANDLER_BUDGET_SECONDS, (
f"healthy /user/update took {timings['h1_healthy']:.3f}s"
)
eventually(
drained,
lambda messages: any(message.get("cache_key") == users[0] for message in messages),
seconds=10,
)
timings["h2_healthy_control"] = _timed_post(
candidate, "/user/update", {"user_id": users[0], "tpm_limit": 1000}
)
assert timings["h2_healthy_control"] < _HANDLER_BUDGET_SECONDS, (
f"healthy control update took {timings['h2_healthy_control']:.3f}s"
)
coordination.signal(signal.SIGSTOP)
try:
timings["s2_wedged_control"] = _timed_post(
candidate, "/user/update", {"user_id": users[0], "tpm_limit": 1000}
)
assert timings["s2_wedged_control"] < _HANDLER_BUDGET_SECONDS, (
f"control update without a cache-relevant field took {timings['s2_wedged_control']:.3f}s"
)
timings["s1_user_update"] = _timed_post(
candidate, "/user/update", {"user_id": users[0], "max_budget": 98.0}
)
assert timings["s1_user_update"] < _HANDLER_BUDGET_SECONDS, (
f"/user/update with max_budget took {timings['s1_user_update']:.3f}s "
"with a wedged coordination Redis"
)
timings["s3_bulk_update"] = _timed_post(
candidate, "/user/bulk_update", {"all_users": True, "user_updates": {"max_budget": 79.0}}
)
assert timings["s3_bulk_update"] < _BULK_BUDGET_SECONDS, (
f"/user/bulk_update over {_USERS} users took {timings['s3_bulk_update']:.3f}s "
"with a wedged coordination Redis"
)
timings["s4_key_update"] = _timed_post(
candidate, "/key/update", {"key": key, "max_budget": 6.0}, timeout=60
)
assert timings["s4_key_update"] < 30, (
f"/key/update hung for {timings['s4_key_update']:.3f}s with a wedged coordination Redis"
)
timings["s5_team_update"] = _timed_post(
candidate, "/team/update", {"team_id": team_id, "max_budget": 6.0}, timeout=60
)
assert timings["s5_team_update"] < 30, (
f"/team/update hung for {timings['s5_team_update']:.3f}s with a wedged coordination Redis"
)
timings["s6_customer_update"] = _timed_post(
candidate, "/customer/update", {"user_id": customer_id, "max_budget": 6.0}
)
assert timings["s6_customer_update"] < _HANDLER_BUDGET_SECONDS, (
f"/customer/update took {timings['s6_customer_update']:.3f}s with a wedged coordination Redis"
)
timings["s7_reset_spend"] = _timed_post(candidate, f"/key/{key}/reset_spend", {"reset_to": 0})
assert timings["s7_reset_spend"] < _HANDLER_BUDGET_SECONDS, (
f"/key/<key>/reset_spend took {timings['s7_reset_spend']:.3f}s with a wedged coordination Redis"
)
missing_started: Final = time.monotonic()
missing: Final = candidate.request(
"POST", "/user/update", {"user_id": users[0], "max_budget": "not-a-number"}
)
timings["s8_invalid_body"] = time.monotonic() - missing_started
assert missing.status_code // 100 == 4, (
f"/user/update with an invalid body returned {missing.status_code} "
f"in {timings['s8_invalid_body']:.3f}s"
)
assert timings["s8_invalid_body"] < _HANDLER_BUDGET_SECONDS, (
f"/user/update with an invalid body took {timings['s8_invalid_body']:.3f}s"
)
def burst_request(path: str, body: Mapping[str, JsonValue]) -> tuple[object, float]:
started: Final = time.monotonic()
try:
response: Final = candidate.client.request(
"POST",
path,
json=body,
headers={"Authorization": f"Bearer {candidate.key}"},
timeout=60,
)
return response.status_code, time.monotonic() - started
except Exception as error: # noqa: BLE001 # the killed worker drops in-flight requests
return error, time.monotonic() - started
port: Final = candidate.client.base_url.port
assert port is not None, f"owned proxy client has no port: {candidate.client.base_url}"
with ThreadPoolExecutor(_BURST) as pool:
futures: Final = [
pool.submit(
burst_request,
*_burst_call(i, users, key, team_id, customer_id),
)
for i in range(_BURST)
]
os.kill(_worker_pid(port), signal.SIGKILL)
results: Final = [future.result() for future in futures]
responses: Final = [(status, elapsed) for status, elapsed in results if isinstance(status, int)]
failures: Final = [status for status, _elapsed in responses if status != 200]
assert not failures, f"burst responses that were not 200: {failures}"
transport_errors: Final = [status for status, _elapsed in results if not isinstance(status, int)]
assert len(transport_errors) <= 3, (
f"{len(transport_errors)} requests raised transport errors: {transport_errors!r}"
)
elapsed_sorted: Final = sorted(
elapsed for i, (status, elapsed) in enumerate(results) if i % 5 in (0, 1, 4) and status == 200
)
timings["c1_burst_p95"] = elapsed_sorted[int(len(elapsed_sorted) * 0.95) - 1]
assert timings["c1_burst_p95"] < _HANDLER_BUDGET_SECONDS, (
f"burst p95 {timings['c1_burst_p95']:.3f}s"
)
eventually(
lambda: candidate.request("GET", "/health/liveliness").status_code,
lambda status: status == 200,
seconds=15,
)
timings["c1_survivor"] = _timed_post(
candidate, "/user/update", {"user_id": users[0], "tpm_limit": 2000}
)
assert timings["c1_survivor"] < _HANDLER_BUDGET_SECONDS, (
f"control update on the surviving worker took {timings['c1_survivor']:.3f}s"
)
finally:
coordination.signal(signal.SIGCONT)
wedged_keys: Final = {users[i] for i in range(_BURST) if i % 5 == 0 and i != 0} | {f"team_id:{team_id}"}
def proxy_log() -> str:
return "".join(
path.read_text() for path in results_dir.glob("owned-proxy-*.log") if path not in prior_logs
)
team_wedged_key: Final = f"team_id:{team_id}"
eventually(
proxy_log,
lambda text: (
all(
f"publish for {wedged_key} failed" in text
for wedged_key in wedged_keys
if wedged_key != team_wedged_key
)
and (
f"publish for {team_wedged_key} failed" in text
or f"internal usage cache entry {team_wedged_key}" in text
)
),
seconds=45,
)
marker: Final = len(received)
drained()
recovered_keys: Final = {str(message.get("cache_key")) for message in received[marker:]}
assert recovered_keys.isdisjoint(wedged_keys), (
f"wedged publishes unexpectedly landed after recovery: {sorted(recovered_keys & wedged_keys)}"
)
coordination.signal(signal.SIGSTOP)
try:
timings["r1b_short_wedge_a"] = _timed_post(
candidate, "/user/update", {"user_id": users[4], "max_budget": 15.0}
)
assert timings["r1b_short_wedge_a"] < _HANDLER_BUDGET_SECONDS, (
f"/user/update inside a short wedge took {timings['r1b_short_wedge_a']:.3f}s"
)
timings["r1b_short_wedge_b"] = _timed_post(
candidate, "/user/update", {"user_id": users[5], "max_budget": 16.0}
)
assert timings["r1b_short_wedge_b"] < _HANDLER_BUDGET_SECONDS, (
f"/user/update inside a short wedge took {timings['r1b_short_wedge_b']:.3f}s"
)
finally:
coordination.signal(signal.SIGCONT)
eventually(
drained,
lambda messages: {str(message.get("cache_key")) for message in messages} >= {users[4], users[5]},
seconds=10,
)
timings["r2_resumed"] = _timed_post(
candidate, "/user/update", {"user_id": users[1], "max_budget": 12.0}
)
assert timings["r2_resumed"] < _HANDLER_BUDGET_SECONDS, (
f"post-recovery /user/update took {timings['r2_resumed']:.3f}s"
)
eventually(
drained,
lambda messages: any(message.get("cache_key") == users[1] for message in messages),
seconds=10,
)
coordination.stop()
timings["f1_refused"] = _timed_post(
candidate, "/user/update", {"user_id": users[2], "max_budget": 13.0}
)
assert timings["f1_refused"] < _HANDLER_BUDGET_SECONDS, (
f"/user/update with refused coordination Redis took {timings['f1_refused']:.3f}s"
)
coordination.start()
restarted_pubsub: Final = subscriber_client.pubsub()
restarted_pubsub.subscribe(_CHANNEL)
restarted_received: list[dict[str, JsonValue]] = []
def drained_after_restart() -> tuple[dict[str, JsonValue], ...]:
restarted_received.extend(_received(restarted_pubsub))
return tuple(restarted_received)
eventually(
lambda: subscriber_client.pubsub_numsub(_CHANNEL)[0][1],
lambda count: count >= 3,
seconds=30,
)
timings["f2_restarted"] = _timed_post(
candidate, "/user/update", {"user_id": users[3], "max_budget": 14.0}
)
assert timings["f2_restarted"] < _HANDLER_BUDGET_SECONDS, (
f"/user/update after Redis restart took {timings['f2_restarted']:.3f}s"
)
eventually(
drained_after_restart,
lambda messages: any(message.get("cache_key") == users[3] for message in messages),
seconds=10,
)
info_last: Final = object_value(candidate.get("/user/info", {"user_id": users[-1]})["user_info"])
assert info_last["max_budget"] == 79.0, info_last
info_user3: Final = object_value(candidate.get("/user/info", {"user_id": users[3]})["user_info"])
assert info_user3["max_budget"] == 14.0, info_user3
finally:
admin.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(identity)))
record_property("cell_elapsed_seconds", timings)

351
tests/integration/management/test_vector_store_config_ownership.py Normal file
View file

@ -0,0 +1,351 @@
import os
import uuid
from collections.abc import Mapping
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Final
from urllib.parse import urlsplit, urlunsplit
import httpx
import psycopg
import pytest
from psycopg import sql
from pydantic import JsonValue
from tests.integration._support.client import Gateway, eventually, object_value
from tests.integration._support.database import read_rows
from tests.integration._support.process import owned_proxy
from tests.integration._support.redis_process import owned_redis
CONFIG_STORE_ID: Final = "vs_integration_config_store"
CONFIG_STORE_NAME: Final = "integration-config-store"
SEARCH_PATH: Final = f"/vector_stores/{CONFIG_STORE_ID}/search"
PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
def listed_rows(response: httpx.Response) -> tuple[dict[str, JsonValue], ...]:
rows: Final = object_value(response.json()).get("data")
assert isinstance(rows, list), response.text
return tuple(object_value(row) for row in rows)
def listed_store(gateway: Gateway, vector_store_id: str, *, key: str | None = None) -> dict[str, JsonValue]:
listed: Final = gateway.request("GET", "/vector_store/list", key=key)
assert listed.status_code == 200, listed.text
matches: Final = tuple(row for row in listed_rows(listed) if row["vector_store_id"] == vector_store_id)
assert len(matches) == 1, f"{vector_store_id} appears {len(matches)} times in {listed.text}"
return matches[0]
def listed_ids(gateway: Gateway) -> tuple[str, ...]:
rows: Final = gateway.get("/vector_store/list")["data"]
assert isinstance(rows, list)
return tuple(str(object_value(row)["vector_store_id"]) for row in rows)
def config_store_info(gateway: Gateway) -> dict[str, JsonValue]:
return object_value(gateway.post("/vector_store/info", {"vector_store_id": CONFIG_STORE_ID})["vector_store"])
def store_rows(vector_store_id: str) -> list[dict[str, JsonValue]]:
return read_rows(
'SELECT vector_store_id, vector_store_name FROM "LiteLLM_ManagedVectorStoresTable" WHERE vector_store_id = %s',
(vector_store_id,),
)
def assert_config_write_refused(gateway: Gateway) -> None:
for path, body in (
("/vector_store/update", {"vector_store_id": CONFIG_STORE_ID, "vector_store_name": "renamed"}),
("/vector_store/delete", {"vector_store_id": CONFIG_STORE_ID}),
("/vector_store/new", {"vector_store_id": CONFIG_STORE_ID, "custom_llm_provider": "openai"}),
):
refused = gateway.request("POST", path, body)
assert refused.status_code == 400, f"{path}: {refused.status_code} {refused.text}"
error = object_value(object_value(refused.json())["detail"])
assert error["vector_store_id"] == CONFIG_STORE_ID, refused.text
assert "config file" in str(error["error"]), refused.text
def burst_list(gateway: Gateway) -> tuple[int, str]:
response: Final = gateway.request("GET", "/vector_store/list")
if response.status_code != 200:
return response.status_code, response.text
ids: Final = tuple(str(row["vector_store_id"]) for row in listed_rows(response))
return response.status_code, "config" if CONFIG_STORE_ID in ids else response.text
def burst_post(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[int, str]:
response: Final = gateway.request("POST", path, body)
return response.status_code, response.text
def upstream_requests(upstream: httpx.Client, marker: str) -> list[dict[str, JsonValue]]:
observed: Final = upstream.get("/__observations")
observed.raise_for_status()
requests: Final = object_value(observed.json())["requests"]
assert isinstance(requests, list), observed.text
return [object_value(value) for value in requests if marker in str(object_value(value)["body"])]
@pytest.mark.covers("mgmt.vector_store.list.keeps_config_store_beside_db_stores")
def test_config_store_is_listed_beside_db_store_and_survives_listing(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
db_store_id: Final = f"vs_db_{uuid.uuid4().hex}"
gateway.post("/vector_store/new", {"vector_store_id": db_store_id, "custom_llm_provider": "openai"})
scenario.cleanups.callback(gateway.post, "/vector_store/delete", {"vector_store_id": db_store_id})
before: Final = config_store_info(gateway)
assert before["vector_store_id"] == CONFIG_STORE_ID, before
config_row: Final = listed_store(gateway, CONFIG_STORE_ID)
assert config_row["is_config"] is True, config_row
assert config_row["vector_store_name"] == CONFIG_STORE_NAME, config_row
assert object_value(config_row["litellm_params"])["api_key"] != "integration-provider-key", config_row
db_row: Final = listed_store(gateway, db_store_id)
assert db_row["is_config"] is False, db_row
after: Final = config_store_info(gateway)
assert after["vector_store_id"] == CONFIG_STORE_ID, after
assert after["is_config"] is True, after
assert after["vector_store_description"] == "declared in tests/integration/proxy_config.yaml", after
assert store_rows(CONFIG_STORE_ID) == [], "config store must not need a database row"
assert listed_store(gateway, CONFIG_STORE_ID)["is_config"] is True
@pytest.mark.covers("mgmt.vector_store.write.config_store_is_read_only")
def test_config_store_refuses_new_update_and_delete(gateway: Gateway) -> None:
assert_config_write_refused(gateway)
row: Final = listed_store(gateway, CONFIG_STORE_ID)
assert row["vector_store_name"] == CONFIG_STORE_NAME, row
assert row["is_config"] is True, row
assert config_store_info(gateway)["vector_store_name"] == CONFIG_STORE_NAME
@pytest.mark.covers("mgmt.vector_store.write.db_store_lifecycle_unchanged_beside_config_store")
def test_db_store_lifecycle_is_unchanged_beside_config_store(gateway: Gateway) -> None:
incomplete: Final = gateway.request("POST", "/vector_store/new", {"custom_llm_provider": "openai"})
assert incomplete.status_code == 400, incomplete.text
db_store_id: Final = f"vs_db_{uuid.uuid4().hex}"
created: Final = gateway.request(
"POST",
"/vector_store/new",
{"vector_store_id": db_store_id, "custom_llm_provider": "openai", "vector_store_name": "first"},
)
assert created.status_code == 200, created.text
assert store_rows(db_store_id) == [{"vector_store_id": db_store_id, "vector_store_name": "first"}]
updated: Final = gateway.post(
"/vector_store/update", {"vector_store_id": db_store_id, "vector_store_name": "second"}
)
assert object_value(updated["vector_store"])["vector_store_name"] == "second", updated
assert store_rows(db_store_id) == [{"vector_store_id": db_store_id, "vector_store_name": "second"}]
row: Final = listed_store(gateway, db_store_id)
assert row["vector_store_name"] == "second" and row["is_config"] is False, row
info: Final = object_value(gateway.post("/vector_store/info", {"vector_store_id": db_store_id})["vector_store"])
assert info["vector_store_name"] == "second" and info["is_config"] is False, info
gateway.post("/vector_store/delete", {"vector_store_id": db_store_id})
assert store_rows(db_store_id) == []
assert db_store_id not in listed_ids(gateway)
assert CONFIG_STORE_ID in listed_ids(gateway)
missing: Final = gateway.request("POST", "/vector_store/info", {"vector_store_id": db_store_id})
assert missing.status_code == 404, missing.text
@pytest.mark.covers("other.vector_store.chat.config_store_search_reaches_upstream_after_listing")
def test_chat_with_config_store_searches_upstream_and_injects_context_after_listing(gateway: Gateway) -> None:
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
):
model: Final = scenario.model()
marker: Final = f"lit6337 {uuid.uuid4().hex}"
assert CONFIG_STORE_ID in listed_ids(gateway)
upstream.get("/__observations").raise_for_status()
completion: Final = gateway.post(
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": marker}], "vector_store_ids": [CONFIG_STORE_ID]},
)
assert object_value(completion["usage"])["total_tokens"] == 40, completion
requests: Final = upstream_requests(upstream, marker)
searches: Final = [value for value in requests if value["path"] == SEARCH_PATH]
assert len(searches) == 1, requests
assert object_value(searches[0]["body"])["query"] == marker, searches
assert searches[0]["authorization"] == "Bearer integration-provider-key", searches
chats: Final = [value for value in requests if value["path"] == "/v1/chat/completions"]
assert len(chats) == 1, requests
messages: Final = object_value(chats[0]["body"])["messages"]
assert isinstance(messages, list), chats
contents: Final = tuple(str(object_value(message)["content"]) for message in messages)
assert contents == (f"Context:\n\nscripted context for {marker}\n\n", marker), contents
@pytest.mark.covers("other.vector_store.search.config_store_passthrough_uses_yaml_credentials_after_listing")
def test_passthrough_search_on_config_store_uses_yaml_credentials_after_listing(gateway: Gateway) -> None:
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream:
marker: Final = f"lit6337 passthrough {uuid.uuid4().hex}"
assert CONFIG_STORE_ID in listed_ids(gateway)
upstream.get("/__observations").raise_for_status()
searched: Final = gateway.request("POST", f"/v1/vector_stores/{CONFIG_STORE_ID}/search", {"query": marker})
assert searched.status_code == 200, searched.text
data: Final = listed_rows(searched)
assert len(data) == 1, searched.text
content: Final = data[0]["content"]
assert isinstance(content, list), searched.text
assert object_value(content[0])["text"] == f"scripted context for {marker}", searched.text
requests: Final = upstream_requests(upstream, marker)
assert [value["path"] for value in requests] == [SEARCH_PATH], requests
assert requests[0]["authorization"] == "Bearer integration-provider-key", requests
@pytest.mark.covers("authz.vector_store.list.non_admin_key_access_to_config_store_follows_grants")
def test_non_admin_key_access_to_config_store_follows_grants_after_admin_listing(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
granted: Final = scenario.key(object_permission={"vector_stores": [CONFIG_STORE_ID]})
plain: Final = scenario.key()
assert CONFIG_STORE_ID in listed_ids(gateway)
row: Final = listed_store(gateway, CONFIG_STORE_ID, key=granted)
assert row["is_config"] is True and row["vector_store_name"] == CONFIG_STORE_NAME, row
unlisted: Final = gateway.request("GET", "/vector_store/list", key=plain)
assert unlisted.status_code == 200, unlisted.text
assert CONFIG_STORE_ID not in {value["vector_store_id"] for value in listed_rows(unlisted)}, unlisted.text
for key in (granted, plain):
info = gateway.request("POST", "/vector_store/info", {"vector_store_id": CONFIG_STORE_ID}, key=key)
assert info.status_code == 200, info.text
assert object_value(object_value(info.json())["vector_store"])["is_config"] is True, info.text
forbidden: Final = gateway.request(
"POST", "/vector_store/delete", {"vector_store_id": CONFIG_STORE_ID}, key=granted
)
assert forbidden.status_code in {400, 401, 403}, forbidden.text
assert CONFIG_STORE_ID in listed_ids(gateway)
@pytest.mark.covers("mgmt.vector_store.list.peer_process_keeps_config_store_and_sees_db_store")
def test_peer_process_keeps_config_store_and_sees_db_store_created_elsewhere(gateway: Gateway, peer: Gateway) -> None:
with gateway.scenario() as scenario:
db_store_id: Final = f"vs_db_{uuid.uuid4().hex}"
gateway.post("/vector_store/new", {"vector_store_id": db_store_id, "custom_llm_provider": "openai"})
scenario.cleanups.callback(gateway.request, "POST", "/vector_store/delete", {"vector_store_id": db_store_id})
for side in (gateway, peer, gateway, peer):
assert listed_store(side, CONFIG_STORE_ID)["is_config"] is True
assert listed_store(side, db_store_id)["is_config"] is False
assert config_store_info(side)["is_config"] is True
assert_config_write_refused(side)
gateway.post("/vector_store/delete", {"vector_store_id": db_store_id})
assert db_store_id not in listed_ids(peer)
assert CONFIG_STORE_ID in listed_ids(peer)
@pytest.mark.covers("mgmt.vector_store.chaos.concurrent_burst_keeps_config_store_across_workers")
def test_concurrent_burst_keeps_config_store_and_refuses_every_config_write(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model()
db_store_ids: Final = tuple(f"vs_db_{uuid.uuid4().hex}" for _ in range(6))
for db_store_id in db_store_ids:
scenario.cleanups.callback(
gateway.request, "POST", "/vector_store/delete", {"vector_store_id": db_store_id}
)
def act(index: int) -> tuple[str, int, str]:
match index % 5:
case 0:
return ("list", *burst_list(gateway))
case 1:
return ("info", *burst_post(gateway, "/vector_store/info", {"vector_store_id": CONFIG_STORE_ID}))
case 2:
return (
"config-update",
*burst_post(
gateway,
"/vector_store/update",
{"vector_store_id": CONFIG_STORE_ID, "vector_store_name": str(index)},
),
)
case 3:
return (
"db-new",
*burst_post(
gateway,
"/vector_store/new",
{
"vector_store_id": db_store_ids[index % len(db_store_ids)],
"custom_llm_provider": "openai",
},
),
)
case _:
return (
"chat",
*burst_post(
gateway,
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": f"burst {index}"}],
"vector_store_ids": [CONFIG_STORE_ID],
},
),
)
with ThreadPoolExecutor(max_workers=10) as pool:
outcomes: Final = tuple(pool.map(act, range(30)))
expected: Final = {"list": 200, "info": 200, "config-update": 400, "db-new": 200, "chat": 200}
assert [(kind, status) for kind, status, _ in outcomes] == [
(kind, expected[kind]) for kind, _, _ in outcomes
], outcomes
assert all(detail == "config" for kind, _, detail in outcomes if kind == "list"), outcomes
assert listed_store(gateway, CONFIG_STORE_ID)["vector_store_name"] == CONFIG_STORE_NAME
assert config_store_info(gateway)["vector_store_name"] == CONFIG_STORE_NAME
assert store_rows(CONFIG_STORE_ID) == []
assert all(len(store_rows(db_store_id)) == 1 for db_store_id in db_store_ids), "each DB store exactly once"
@pytest.mark.timeout(180)
@pytest.mark.covers("mgmt.vector_store.chaos.redis_outage_keeps_config_store_and_recovers")
def test_redis_outage_keeps_config_store_served_and_recovers(
gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
original: Final = os.environ["DATABASE_URL"]
identity: Final = "integration_vs_outage_" + uuid.uuid4().hex
parsed: Final = urlsplit(original)
database_url: Final = urlunsplit((parsed.scheme, parsed.netloc, "/" + identity, "", ""))
with psycopg.connect(original, autocommit=True) as admin:
admin.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(identity)))
try:
with owned_redis(tmp_path) as cache, monkeypatch.context() as environment:
environment.setenv("DATABASE_URL", database_url)
overrides: Final = {
"DATABASE_URL": database_url,
"REDIS_HOST": cache.host,
"REDIS_PORT": str(cache.port),
"REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT": "1",
}
with owned_proxy(gateway, tmp_path, overrides, config=PROXY_CONFIG, workers=2) as candidate:
db_store_id: Final = f"vs_db_{uuid.uuid4().hex}"
for phase in ("before", "during", "after"):
if phase == "during":
cache.stop()
if phase == "after":
cache.start()
for _ in range(4):
assert listed_store(candidate, CONFIG_STORE_ID)["is_config"] is True, phase
assert config_store_info(candidate)["vector_store_name"] == CONFIG_STORE_NAME, phase
assert_config_write_refused(candidate)
created = candidate.request(
"POST",
"/vector_store/new",
{"vector_store_id": f"{db_store_id}_{phase}", "custom_llm_provider": "openai"},
)
assert created.status_code == 200, (phase, created.text)
assert eventually(
lambda phase=phase: store_rows(f"{db_store_id}_{phase}"), lambda rows: len(rows) == 1
), phase
assert f"{db_store_id}_{phase}" in listed_ids(candidate), phase
assert store_rows(CONFIG_STORE_ID) == []
with psycopg.connect(database_url) as fresh:
counted: Final = fresh.execute(
'SELECT count(*) FROM "LiteLLM_ManagedVectorStoresTable" WHERE vector_store_id LIKE %s',
(f"{db_store_id}%",),
).fetchone()
assert counted is not None and counted[0] == 3, counted
finally:
admin.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(identity)))
assert admin.execute("SELECT datname FROM pg_database WHERE datname=%s", (identity,)).fetchall() == []

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

116
tests/integration/observability/test_callback_delivery.py
View file

@ -102,7 +102,7 @@ def test_concurrent_success_and_failure_join_callbacks_and_rows_without_credenti
responses: Final = tuple(pool.map(request, tags))
assert tuple(response.status_code for response in responses) == (200, 400, 200, 400)
assert len(provider.drain()) == 4
batches = []
batches: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches
def delivered() -> tuple[dict, ...]:
batches.extend(endpoint.drain())
@ -135,7 +135,7 @@ def test_concurrent_success_and_failure_join_callbacks_and_rows_without_credenti
assert "synthetic callback failure" in json.dumps(event["error_information"])
rows: Final = eventually(
lambda identity=event["id"]: read_rows(
'SELECT request_id, spend, prompt_tokens, completion_tokens, request_tags '
"SELECT request_id, spend, prompt_tokens, completion_tokens, request_tags "
'FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
(identity,),
),
@ -156,6 +156,118 @@ def test_concurrent_success_and_failure_join_callbacks_and_rows_without_credenti
assert event["prompt_tokens"] == event["completion_tokens"] == rows[0]["completion_tokens"] == 0
def _responses_frames(identity: str, text: str) -> tuple[bytes, ...]:
output: Final = [
{
"type": "message",
"id": f"msg_{identity}",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": text, "annotations": []}],
}
]
completed: Final = {
"id": identity,
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-4o-mini",
"output": output,
"usage": {
"input_tokens": 11,
"output_tokens": 4,
"total_tokens": 15,
"input_tokens_details": {"cached_tokens": 0},
"output_tokens_details": {"reasoning_tokens": 0},
},
}
events: Final = (
{"type": "response.created", "response": {**completed, "status": "in_progress", "output": [], "usage": None}},
{
"type": "response.output_text.delta",
"item_id": f"msg_{identity}",
"output_index": 0,
"content_index": 0,
"delta": text,
},
{"type": "response.completed", "response": completed},
)
return tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events)
@pytest.mark.covers("other.observability.callbacks.streamed_responses_events_carry_provider_response_headers")
def test_streamed_responses_success_callback_carries_provider_apim_request_id(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "resp_" + uuid.uuid4().hex
correlation: Final = "azure-correlation-" + marker
region: Final = "East US 2"
secret: Final = "synthetic-provider-secret-" + marker
sink_secret: Final = "synthetic-sink-secret-" + marker
def upstream(request: Request) -> Reply:
assert request.target.endswith("/responses"), request.target
assert request.headers["authorization"] == f"Bearer {secret}"
assert json.loads(request.body) == {
"model": "gpt-4o-mini",
"input": "header control " + marker,
"stream": True,
}, request.body
return Reply(
content_type="text/event-stream",
chunks=_responses_frames(marker, "streamed control"),
headers={"apim-request-id": correlation, "x-ms-region": region},
)
def sink(request: Request) -> Reply:
assert request.headers["authorization"] == f"Bearer {sink_secret}"
return Reply()
with wire_server(upstream) as provider, wire_server(sink) as endpoint:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["litellm_settings"].update({"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1})
path: Final = tmp_path / "callbacks.yaml"
path.write_text(yaml.safe_dump(config))
with (
owned_proxy(
gateway,
tmp_path,
{
"GENERIC_LOGGER_ENDPOINT": endpoint.url,
"GENERIC_LOGGER_HEADERS": f"Authorization=Bearer {sink_secret}",
},
config=path,
) as candidate,
candidate.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key=secret)
response: Final = candidate.request(
"POST", "/v1/responses", {"model": model, "input": "header control " + marker, "stream": True}
)
assert response.status_code == 200, response.text
assert f'"item_id":"msg_{marker}"' in response.text, response.text
assert '"type":"response.completed"' in response.text, response.text
assert len(provider.drain()) == 1
batches: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches
def delivered() -> tuple[dict, ...]:
batches.extend(endpoint.drain())
return tuple(
event for batch in batches for event in json.loads(batch.body) if event.get("model_group") == model
)
events: Final = eventually(delivered, lambda values: len(values) == 1, seconds=10)
assert (events[0]["status"], events[0]["stream"], events[0]["call_type"]) == ("success", True, "aresponses")
additional_headers: Final = events[0]["hidden_params"]["additional_headers"] or {}
provider_headers: Final = {
name: value
for name, value in additional_headers.items()
if name in ("llm_provider-apim-request-id", "llm_provider-x-ms-region")
}
assert provider_headers == {
"llm_provider-apim-request-id": correlation,
"llm_provider-x-ms-region": region,
}, json.dumps(events[0]["hidden_params"])
_RAISING_HOOK: Final = """
from litellm.integrations.custom_logger import CustomLogger

102
tests/integration/observability/test_guardrail_effects.py
View file

@ -5,9 +5,7 @@ from typing import Final
import pytest
import yaml
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.client import Gateway
from integration._support.mcp import mcp_peer, register_mcp, tool_names
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
@ -146,6 +144,104 @@ def test_guardrail_denial_prevents_provider_and_preserves_allowed_control(gatewa
assert len(policy.drain()) == 2
@pytest.mark.covers("other.observability.guardrails.bedrock_passthrough_converse_scans_only_caller_content")
def test_bedrock_passthrough_converse_guardrail_ignores_denied_term_in_tool_definition(
gateway: Gateway, tmp_path: Path
) -> None:
identity: Final = "guardrail" + uuid.uuid4().hex
denied: Final = "synthetic denied marker"
allowed: Final = "synthetic allowed weather question"
access_key: Final = "AKIASYNTHETICPASSTHROUGH"
tool_config: Final = {
"tools": [
{
"toolSpec": {
"name": "lookup_weather",
"description": f"Look up the forecast, never answer a {denied}",
"inputSchema": {
"json": {
"type": "object",
"properties": {"city": {"type": "string", "enum": [denied]}},
"required": ["city"],
}
},
}
}
]
}
def guardrail(request: Request) -> Reply:
assert request.target == "/beta/litellm_basic_guardrail_api"
texts: Final = json.loads(request.body)["texts"]
result: Final = (
{"action": "BLOCKED", "blocked_reason": "synthetic policy denial"}
if any(denied in text for text in texts)
else {"action": "NONE"}
)
return Reply(body=json.dumps(result).encode())
def runtime(request: Request) -> Reply:
assert request.target == "/model/anthropic.claude-3-haiku-20240307-v1:0/converse"
assert request.headers["authorization"].startswith(f"AWS4-HMAC-SHA256 Credential={access_key}/"), (
request.headers
)
return Reply(
body=json.dumps(
{
"output": {"message": {"role": "assistant", "content": [{"text": "sunny passthrough control"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15},
"metrics": {"latencyMs": 1},
}
).encode()
)
with wire_server(guardrail) as policy, wire_server(runtime) as bedrock, gateway.scenario() as scenario:
model: Final = scenario.model(
model="bedrock/anthropic.claude-3-haiku-20240307-v1:0",
api_key=None,
api_base=bedrock.url,
aws_access_key_id=access_key,
aws_secret_access_key="synthetic-secret",
aws_region_name="us-east-1",
)
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["guardrails"] = [
{
"guardrail_name": identity,
"litellm_params": {
"guardrail": "generic_guardrail_api",
"mode": "pre_call",
"default_on": True,
"api_base": policy.url,
"api_key": "synthetic-guardrail-key",
},
}
]
path: Final = tmp_path / "bedrock-passthrough.yaml"
path.write_text(yaml.safe_dump(config))
with owned_proxy(gateway, tmp_path, {}, config=path) as candidate:
route: Final = f"/bedrock/model/{model}/converse"
passed: Final = candidate.request(
"POST",
route,
{"messages": [{"role": "user", "content": [{"text": allowed}]}], "toolConfig": tool_config},
)
assert passed.status_code == 200, passed.text
assert passed.json()["output"]["message"]["content"] == [{"text": "sunny passthrough control"}]
forwarded: Final = bedrock.drain()
assert len(forwarded) == 1, "the runtime peer must see exactly the allowed request"
assert json.loads(forwarded[0].body)["toolConfig"] == tool_config
blocked: Final = candidate.request(
"POST",
route,
{"messages": [{"role": "user", "content": [{"text": denied}]}], "toolConfig": tool_config},
)
assert blocked.status_code == 400 and "synthetic policy denial" in blocked.text, blocked.text
assert bedrock.drain() == ()
assert [json.loads(request.body)["texts"] for request in policy.drain()] == [[allowed], [denied]]
@pytest.mark.covers("other.mcp.guardrails.request_selection_blocks_resolved_tool_without_execution")
def test_request_selected_mcp_guardrail_blocks_direct_and_virtual_calls(gateway: Gateway, tmp_path: Path) -> None:
guardrail = "mcp-policy-" + uuid.uuid4().hex

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)

72
tests/integration/pricing/test_configured_prices.py
View file

@ -1,10 +1,12 @@
from collections.abc import Iterator, Mapping
from typing import Final
from pathlib import Path
import json
import uuid
from collections.abc import Iterator, Mapping
from pathlib import Path
from typing import Final
import pytest
import yaml
from pydantic import JsonValue
from tests.integration._support.client import Gateway, eventually, object_value, string_value
from tests.integration._support.database import read_rows
@ -28,6 +30,31 @@ def test_custom_price_is_reported_and_charged(gateway: Gateway) -> None:
assert params["output_cost_per_token"] == 0.002
@pytest.mark.covers("quota_management.cost_estimate.configured_price.reported_for_model_absent_from_cost_map")
def test_cost_estimate_reports_configured_prices_for_model_absent_from_cost_map(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"openai/integration-on-prem-{uuid.uuid4().hex}",
input_cost_per_token=0.003,
output_cost_per_token=0.007,
)
response: Final = gateway.request(
"POST",
"/cost/estimate",
{"model": model, "input_tokens": 1000, "output_tokens": 500, "num_requests_per_day": 10},
)
assert response.status_code == 200, response.text
body: Final = object_value(response.json())
assert body["input_cost_per_token"] == pytest.approx(0.003), response.text
assert body["output_cost_per_token"] == pytest.approx(0.007), response.text
assert body["input_cost_per_request"] == pytest.approx(1000 * 0.003), response.text
assert body["output_cost_per_request"] == pytest.approx(500 * 0.007), response.text
margin: Final = body["margin_cost_per_request"]
assert isinstance(margin, float), response.text
assert body["cost_per_request"] == pytest.approx(1000 * 0.003 + 500 * 0.007 + margin), response.text
assert body["daily_cost"] == pytest.approx(10 * (1000 * 0.003 + 500 * 0.007 + margin)), response.text
@pytest.mark.covers("quota_management.spend_tracking.default_prices.survive_nullable_sibling_reload")
def test_default_prices_survive_nullable_sibling_and_reload(gateway: Gateway) -> None:
for registration_order in (("custom", "omitted", "nullable"), ("nullable", "omitted", "custom")):
@ -103,6 +130,45 @@ def test_default_prices_survive_nullable_sibling_and_reload(gateway: Gateway) ->
assert float(rows[0]["spend"]) == pytest.approx(expected, rel=1e-6)
COST_MAP_DISPLAY_PRICING_KEYS: Final = frozenset(
{
"input_cost_per_token",
"output_cost_per_token",
"cache_read_input_token_cost",
"cache_creation_input_token_cost",
}
)
def persisted_model_info(identity: str) -> dict[str, JsonValue]:
rows: Final = read_rows('SELECT model_info FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,))
assert len(rows) == 1, f"Deployment {identity} has {len(rows)} rows"
stored: Final = rows[0]["model_info"]
return object_value(json.loads(stored) if isinstance(stored, str) else stored)
@pytest.mark.covers("pricing.model_update.echoed_cost_map_price_is_not_persisted_as_override")
def test_saving_echoed_model_info_does_not_freeze_cost_map_price_into_deployment(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model()
entries: Final = gateway.get("/model/info")["data"]
assert isinstance(entries, list)
target: Final = next(object_value(entry) for entry in entries if object_value(entry)["model_name"] == model)
displayed: Final = object_value(target["model_info"])
identity: Final = string_value(displayed["id"])
assert isinstance(displayed["input_cost_per_token"], float), displayed
assert isinstance(displayed["output_cost_per_token"], float), displayed
fresh: Final = persisted_model_info(identity)
assert {key: value for key, value in fresh.items() if key in COST_MAP_DISPLAY_PRICING_KEYS} == {}, fresh
saved: Final = gateway.request(
"PATCH", f"/model/{identity}/update", {"model_info": {**displayed, "description": "echoed ui save"}}
)
assert saved.status_code == 200, saved.text
stored: Final = persisted_model_info(identity)
assert stored["description"] == "echoed ui save", stored
assert {key: value for key, value in stored.items() if key in COST_MAP_DISPLAY_PRICING_KEYS} == {}, stored
@pytest.mark.covers("quota_management.spend_tracking.default_prices.loaded_router_preserves_cached_defaults")
def test_loaded_router_preserves_cached_defaults_during_real_requests(gateway: Gateway, tmp_path: Path) -> None:
from litellm import Router

89
tests/integration/pricing/test_databricks_cache_pricing.py Normal file
View file

@ -0,0 +1,89 @@
import json
import uuid
from typing import Final
import pytest
from tests.integration._support.client import Gateway, eventually, object_value, string_value
from tests.integration._support.database import read_rows
from tests.integration._support.upstream import delete_scenario, register_scenario
from tests.integration.cost_calculation.cost_tracking_case import JsonResponse
INPUT_RATE: Final = 0.001
OUTPUT_RATE: Final = 0.002
CACHE_CREATION_RATE: Final = 0.004
CACHE_READ_RATE: Final = 0.0001
UNCACHED_PROMPT_TOKENS: Final = 1000
CACHE_CREATION_TOKENS: Final = 2000
CACHE_READ_TOKENS: Final = 8000
PROMPT_TOKENS: Final = UNCACHED_PROMPT_TOKENS + CACHE_CREATION_TOKENS + CACHE_READ_TOKENS
COMPLETION_TOKENS: Final = 500
def databricks_cached_response() -> JsonResponse:
return JsonResponse(
content_type="application/json",
body={
"id": "chatcmpl-$REQUEST_ID",
"object": "chat.completion",
"created": 1700000000,
"model": "databricks-claude-integration",
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "cached reply"}, "finish_reason": "stop"}
],
"usage": {
"prompt_tokens": PROMPT_TOKENS,
"completion_tokens": COMPLETION_TOKENS,
"total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS,
"cache_creation_input_tokens": CACHE_CREATION_TOKENS,
"cache_read_input_tokens": CACHE_READ_TOKENS,
},
},
)
@pytest.mark.covers("pricing.databricks.cached_prompt_tokens_bill_at_cache_rates")
def test_databricks_cached_prompt_tokens_bill_at_cache_rates_not_input_rate(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
scenario_id: Final = f"databricks-cache-{uuid.uuid4().hex[:12]}"
handle: Final = register_scenario(scenario_id, databricks_cached_response())
scenario.cleanups.callback(delete_scenario, handle)
model: Final = scenario.model(
model="databricks/databricks-claude-integration",
api_base=handle.api_base(),
input_cost_per_token=INPUT_RATE,
output_cost_per_token=OUTPUT_RATE,
cache_creation_input_token_cost=CACHE_CREATION_RATE,
cache_read_input_token_cost=CACHE_READ_RATE,
)
response: Final = gateway.request(
"POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "cache control"}]}
)
assert response.status_code == 200, response.text
expected_prompt_cost: Final = (
UNCACHED_PROMPT_TOKENS * INPUT_RATE
+ CACHE_CREATION_TOKENS * CACHE_CREATION_RATE
+ CACHE_READ_TOKENS * CACHE_READ_RATE
)
expected_completion_cost: Final = COMPLETION_TOKENS * OUTPUT_RATE
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(
expected_prompt_cost + expected_completion_cost, rel=1e-6
), response.text
request_id: Final = string_value(object_value(response.json())["id"])
rows: Final = eventually(
lambda: read_rows(
'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" '
"WHERE request_id = %s",
(request_id,),
),
lambda values: len(values) == 1,
seconds=70,
)
assert rows[0]["prompt_tokens"] == PROMPT_TOKENS
assert rows[0]["completion_tokens"] == COMPLETION_TOKENS
assert float(rows[0]["spend"]) == pytest.approx(expected_prompt_cost + expected_completion_cost, rel=1e-6)
metadata: Final = rows[0]["metadata"]
parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata)
breakdown: Final = object_value(parsed["cost_breakdown"])
assert float(breakdown["input_cost"]) == pytest.approx(expected_prompt_cost, rel=1e-6)
assert float(breakdown["output_cost"]) == pytest.approx(expected_completion_cost, rel=1e-6)

59
tests/integration/pricing/test_ocr_page_pricing.py Normal file
View file

@ -0,0 +1,59 @@
import uuid
from typing import Final
import pytest
from tests.integration._support.client import Gateway, eventually, string_value
from tests.integration._support.database import read_rows
from tests.integration._support.upstream import delete_scenario, register_scenario
from tests.integration.cost_calculation.cost_tracking_case import JsonResponse
@pytest.mark.covers("pricing.ocr.annotation_pages_billed_at_annotation_rate")
def test_ocr_annotation_pages_are_billed_at_annotation_cost_per_page(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
scenario_id: Final = f"ocr-annotation-{uuid.uuid4().hex[:12]}"
handle: Final = register_scenario(
scenario_id,
JsonResponse(
content_type="application/json",
body={
"pages": [{"index": index, "markdown": f"page {index}"} for index in range(3)],
"model": "integration-ocr",
"document_annotation": '{"title": "annotated"}',
"usage_info": {"pages_processed": 3, "pages_processed_annotation": 2, "doc_size_bytes": 4096},
},
),
)
scenario.cleanups.callback(delete_scenario, handle)
model: Final = scenario.model(
model=f"mistral/integration-ocr-{scenario_id}",
api_base=f"{handle.api_base()}/v1",
ocr_cost_per_page=0.002,
annotation_cost_per_page=0.01,
)
response: Final = gateway.request(
"POST",
"/v1/ocr",
{
"model": model,
"document": {"type": "document_url", "document_url": "https://example.com/annotated.pdf"},
"document_annotation_format": {"type": "json_schema", "json_schema": {"name": "title"}},
},
)
assert response.status_code == 200, response.text
assert response.json()["usage_info"] == {
"pages_processed": 3,
"pages_processed_annotation": 2,
"credits": None,
"doc_size_bytes": 4096,
}, response.text
expected: Final = 3 * 0.002 + 2 * 0.01
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected), response.text
request_id: Final = string_value(response.headers["x-litellm-call-id"])
rows: Final = eventually(
lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)),
lambda values: len(values) == 1,
seconds=70,
)
assert float(rows[0]["spend"]) == pytest.approx(expected)

71
tests/integration/pricing/test_service_tier_pricing.py Normal file
View file

@ -0,0 +1,71 @@
import json
from typing import Final
import httpx
import pytest
from tests.integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value
from tests.integration._support.database import read_rows
STANDARD_INPUT_RATE: Final = 0.001
STANDARD_OUTPUT_RATE: Final = 0.002
ULTRAFAST_INPUT_RATE: Final = 0.01
ULTRAFAST_OUTPUT_RATE: Final = 0.02
def assert_chat_bills_rates(
gateway: Gateway, model: str, service_tier: str | None, input_rate: float, output_rate: float
) -> None:
with httpx.Client(base_url=gateway.upstream_url, trust_env=False) as upstream:
upstream.get("/__observations").raise_for_status()
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": f"service tier {service_tier} control"}],
**({} if service_tier is None else {"service_tier": service_tier}),
},
)
assert response.status_code == 200, response.text
expected: Final = 20 * input_rate + 20 * output_rate
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected, rel=1e-6), response.text
observations: Final = JSON_OBJECT.validate_json(upstream.get("/__observations").content)["requests"]
assert isinstance(observations, list)
assert len(observations) == 1
body: Final = object_value(object_value(observations[0])["body"])
assert body == {
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": f"service tier {service_tier} control"}],
**({} if service_tier is None else {"service_tier": service_tier}),
}, response.text
request_id: Final = string_value(object_value(response.json())["id"])
rows: Final = eventually(
lambda: read_rows(
'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id = %s',
(request_id,),
),
lambda values: len(values) == 1,
seconds=70,
)
assert rows[0]["prompt_tokens"] == 20
assert rows[0]["completion_tokens"] == 20
assert float(rows[0]["spend"]) == pytest.approx(expected, rel=1e-6)
metadata: Final = rows[0]["metadata"]
parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata)
breakdown: Final = object_value(parsed["cost_breakdown"])
assert float(breakdown["input_cost"]) == pytest.approx(20 * input_rate, rel=1e-6)
assert float(breakdown["output_cost"]) == pytest.approx(20 * output_rate, rel=1e-6)
@pytest.mark.covers("quota_management.spend_tracking.service_tier_pricing.ultrafast_bills_ultrafast_rates")
def test_ultrafast_service_tier_bills_ultrafast_rates_and_keeps_pricing_off_the_wire(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model(
input_cost_per_token=STANDARD_INPUT_RATE,
output_cost_per_token=STANDARD_OUTPUT_RATE,
input_cost_per_token_ultrafast=ULTRAFAST_INPUT_RATE,
output_cost_per_token_ultrafast=ULTRAFAST_OUTPUT_RATE,
)
assert_chat_bills_rates(gateway, model, "ultrafast", ULTRAFAST_INPUT_RATE, ULTRAFAST_OUTPUT_RATE)
assert_chat_bills_rates(gateway, model, None, STANDARD_INPUT_RATE, STANDARD_OUTPUT_RATE)

113
tests/integration/providers/test_anthropic_advisor_wire.py Normal file
View file

@ -0,0 +1,113 @@
import json
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
_ADVISOR_KEY: Final = "synthetic-advisor-key"
_QUESTION: Final = "which index should this query use"
_ADVICE: Final = "use the composite index on (tenant_id, created_at)"
_FINAL_ANSWER: Final = "done, the composite index is the right one"
_ADVISOR_CALL_MESSAGE: Final = {
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "advisor-call",
"type": "function",
"function": {"name": "advisor", "arguments": json.dumps({"question": _QUESTION})},
}
],
}
_FINAL_MESSAGE: Final = {"role": "assistant", "content": _FINAL_ANSWER}
def _chat_completion(identity: str, message: dict[str, object], finish_reason: str) -> Reply:
return Reply(
body=json.dumps(
{
"id": f"chatcmpl-{identity}",
"object": "chat.completion",
"created": 1,
"model": "llama-3.3-70b-versatile",
"choices": [{"index": 0, "message": message, "finish_reason": finish_reason}],
"usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14},
}
).encode()
)
def _executor_reply(body: dict[str, object], identity: str) -> Reply:
messages: Final = body["messages"]
assert isinstance(messages, list)
if any(message.get("role") == "tool" for message in messages):
assert messages[-1]["content"] == _ADVICE
return _chat_completion(identity, _FINAL_MESSAGE, "stop")
tools: Final = body["tools"]
assert isinstance(tools, list)
assert tools[0]["function"]["name"] == "advisor"
return _chat_completion(identity, _ADVISOR_CALL_MESSAGE, "tool_calls")
@pytest.mark.covers("providers.anthropic_messages_advisor.sub_call_uses_the_configured_advisor_deployment")
def test_advisor_sub_call_reaches_the_router_deployment_with_its_key_instead_of_anthropic_unauthenticated(
gateway: Gateway,
) -> None:
identity: Final = "advisor-wire-" + uuid.uuid4().hex
def respond(request: Request) -> Reply:
body: Final = json.loads(request.body)
if request.target == "/v1/chat/completions":
assert request.headers["authorization"] == "Bearer integration-provider-key"
return _executor_reply(body, identity)
assert request.target == "/v1/messages"
assert request.headers["x-api-key"] == _ADVISOR_KEY
assert body["model"] == "claude-opus-4-1-20250805"
assert body["messages"] == [
{"role": "user", "content": "please plan the migration"},
{"role": "user", "content": _QUESTION},
]
assert "tools" not in body
return Reply(
body=json.dumps(
{
"id": f"msg-{identity}",
"type": "message",
"role": "assistant",
"model": "claude-opus-4-1-20250805",
"content": [{"type": "text", "text": _ADVICE}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 12, "output_tokens": 6},
}
).encode()
)
with wire_server(respond) as wire, gateway.scenario() as scenario:
executor: Final = scenario.model(model="hosted_vllm/llama-3.3-70b", api_base=wire.url + "/v1")
advisor: Final = scenario.model(
model="anthropic/claude-opus-4-1-20250805", api_base=wire.url, api_key=_ADVISOR_KEY
)
response: Final = gateway.request(
"POST",
"/v1/messages",
{
"model": executor,
"max_tokens": 64,
"messages": [{"role": "user", "content": "please plan the migration"}],
"tools": [{"type": "advisor_20260301", "name": "advisor", "model": advisor}],
},
)
assert response.status_code == 200, response.text
body: Final = response.json()
assert body["content"] == [{"type": "text", "text": _FINAL_ANSWER}], response.text
assert body["stop_reason"] == "end_turn", response.text
assert [request.target for request in wire.drain()] == [
"/v1/chat/completions",
"/v1/messages",
"/v1/chat/completions",
]

85
tests/integration/providers/test_azure_ai_chat_wire.py Normal file
View file

@ -0,0 +1,85 @@
import json
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_BACKEND: Final = "kimi-k2-thinking"
_API_KEY: Final = "synthetic-azure-ai-key"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_THINKING_BLOCK: Final[JsonValue] = {
"type": "thinking",
"thinking": "The user wants the sum of 17 and 26.",
"signature": "synthetic-signature",
}
_HISTORY_WITH_ANTHROPIC_FIELDS: Final[JsonValue] = [
{
"role": "system",
"content": "You are a calculator.",
"cache_control": {"type": "ephemeral"},
},
{"role": "user", "content": "What is 17 + 26?"},
{
"role": "assistant",
"content": "43",
"thinking_blocks": [_THINKING_BLOCK],
"provider_specific_fields": {"citations": None},
},
{"role": "user", "content": "And doubled?"},
]
_HISTORY_AS_OPENAI_SPEC: Final[JsonValue] = [
{"role": "system", "content": "You are a calculator."},
{"role": "user", "content": "What is 17 + 26?"},
{"role": "assistant", "content": "43"},
{"role": "user", "content": "And doubled?"},
]
def _completion(identity: str) -> bytes:
return json.dumps(
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": _BACKEND,
"choices": [{"index": 0, "message": {"role": "assistant", "content": "86"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 31, "completion_tokens": 2, "total_tokens": 33},
}
).encode()
@pytest.mark.covers("providers.azure_ai.anthropic_message_fields_are_stripped_before_foundry")
def test_azure_ai_strips_thinking_blocks_and_cache_control_from_forwarded_messages(gateway: Gateway) -> None:
identity: Final = f"azure-ai-strip-{uuid.uuid4().hex}"
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == "/chat/completions"
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
body: Final = _JSON_OBJECT.validate_json(request.body)
assert body["model"] == _BACKEND
assert body["messages"] == _HISTORY_AS_OPENAI_SPEC
return Reply(body=_completion(identity))
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"azure_ai/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": _HISTORY_WITH_ANTHROPIC_FIELDS},
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["id"] == identity
assert payload["choices"] == [
{
"finish_reason": "stop",
"index": 0,
"message": {"role": "assistant", "content": "86"},
"provider_specific_fields": {},
}
]
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]

48
tests/integration/providers/test_azure_ai_flux2_image_wire.py Normal file
View file

@ -0,0 +1,48 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_FLEX_MODEL: Final = "azure_ai/FLUX.2-flex"
_PROMPT: Final = "a red fox in the snow"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
@pytest.mark.covers("other.provider_wire.azure_ai.flux2_flex_generation_targets_flex_path_with_bfl_body")
def test_azure_flux2_flex_generation_hits_flex_provider_path_not_pro(gateway: Gateway) -> None:
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == "/providers/blackforestlabs/v1/flux-2-flex?api-version=preview"
assert request.headers["api-key"] == "synthetic-azure-key"
assert _JSON_OBJECT.validate_json(request.body) == {
"model": "FLUX.2-flex",
"prompt": _PROMPT,
"num_images": 2,
"width": 1536,
"height": 1024,
"guidance": 4.5,
"steps": 32,
}
return Reply(body=json.dumps({"data": [{"b64_json": "aW1n"}, {"b64_json": "aW1n"}]}).encode())
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=_FLEX_MODEL, api_base=wire.url, api_key="synthetic-azure-key", api_version="preview"
)
response: Final = gateway.request(
"POST",
"/v1/images/generations",
{"model": model, "prompt": _PROMPT, "n": 2, "size": "1536x1024", "guidance": 4.5, "steps": 32},
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["data"] == [
{"url": None, "b64_json": "aW1n", "revised_prompt": None, "provider_specific_fields": None},
{"url": None, "b64_json": "aW1n", "revised_prompt": None, "provider_specific_fields": None},
]
assert [(request.method, request.target) for request in wire.drain()] == [
("POST", "/providers/blackforestlabs/v1/flux-2-flex?api-version=preview")
]

46
tests/integration/providers/test_azure_ai_rerank_auth_wire.py Normal file
View file

@ -0,0 +1,46 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
MODEL: Final = "azure_ai/Cohere-rerank-v4.0-fast"
ENTRA_TOKEN: Final = "synthetic-entra-access-token"
QUERY: Final = "which document mentions the gateway"
DOCUMENTS: Final = ("the gateway proxies rerank calls", "unrelated synthetic text")
RESPONSE: Final = json.dumps(
{
"id": "synthetic-rerank-id",
"results": [{"index": 0, "relevance_score": 0.91}, {"index": 1, "relevance_score": 0.03}],
"meta": {"api_version": {"version": "2"}, "billed_units": {"search_units": 1}},
}
).encode()
def entra_rerank_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/providers/cohere/v2/rerank"
assert request.headers["authorization"] == f"Bearer {ENTRA_TOKEN}"
assert "api-key" not in request.headers
body: Final = json.loads(request.body)
assert body == {"model": "Cohere-rerank-v4.0-fast", "query": QUERY, "documents": list(DOCUMENTS), "top_n": 2}
return Reply(body=RESPONSE)
@pytest.mark.covers("other.provider_wire.azure_ai.rerank_entra_token_without_api_key_reaches_provider")
def test_azure_ai_rerank_with_entra_token_and_no_api_key_sends_bearer_to_provider(gateway: Gateway) -> None:
with wire_server(entra_rerank_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=MODEL,
api_key=None,
api_base=f"{wire.url}/providers/cohere/v2",
azure_ad_token=ENTRA_TOKEN,
model_info={"mode": "rerank"},
)
response: Final = gateway.request(
"POST", "/v1/rerank", {"model": model, "query": QUERY, "documents": list(DOCUMENTS), "top_n": 2}
)
assert response.status_code == 200, response.text
body: Final = response.json()
assert [(result["index"], result["relevance_score"]) for result in body["results"]] == [(0, 0.91), (1, 0.03)]
assert len(wire.drain()) == 1, "Expected exactly one provider rerank call"

164
tests/integration/providers/test_bedrock_auth_wire.py
View file

@ -7,18 +7,22 @@ from typing import Final
import pytest
import yaml
from integration._support.client import Gateway
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
MODEL: Final = "bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0"
TOKEN: Final = "synthetic-bedrock-bearer"
RESPONSE: Final = json.dumps({
"output": {"message": {"role": "assistant", "content": [{"text": "bedrock wire control"}]}},
"stopReason": "end_turn", "usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15},
"metrics": {"latencyMs": 1},
}).encode()
ACCESS_KEY: Final = "AKIAINTEGRATION000002"
CLIENT_OAUTH_TOKEN: Final = "Bearer sk-ant-oat01-synthetic-client-subscription-token"
RESPONSE: Final = json.dumps(
{
"output": {"message": {"role": "assistant", "content": [{"text": "bedrock wire control"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15},
"metrics": {"latencyMs": 1},
}
).encode()
def bearer_peer(request: Request) -> Reply:
@ -34,31 +38,58 @@ def bearer_peer(request: Request) -> Reply:
@pytest.mark.covers("other.provider_wire.bedrock.bearer_sdk_skips_credential_chain")
async def test_bearer_only_sdk_sync_async_requests_do_not_require_aws_credentials(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
async def test_bearer_only_sdk_sync_async_requests_do_not_require_aws_credentials(
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
import litellm
empty: Final = tmp_path / "empty-aws-config"
empty.write_text("")
for name in tuple(name for name in os.environ if name.startswith("AWS_")):
monkeypatch.delenv(name, raising=False)
for name, value in {"AWS_CONFIG_FILE": str(empty), "AWS_SHARED_CREDENTIALS_FILE": str(empty), "AWS_EC2_METADATA_DISABLED": "true", "LITELLM_RUST": "false"}.items():
for name, value in {
"AWS_CONFIG_FILE": str(empty),
"AWS_SHARED_CREDENTIALS_FILE": str(empty),
"AWS_EC2_METADATA_DISABLED": "true",
"LITELLM_RUST": "false",
}.items():
monkeypatch.setenv(name, value)
with wire_server(bearer_peer) as wire:
with pytest.raises(litellm.APIConnectionError, match=r"config profile .* could not be found"):
await asyncio.to_thread(litellm.completion, model=MODEL, aws_profile_name="integration-profile-must-not-be-read", aws_region_name="us-east-1", aws_bedrock_runtime_endpoint=wire.url, messages=[{"role": "user", "content": "synthetic credential control"}], timeout=5, num_retries=0)
await asyncio.to_thread(
litellm.completion,
model=MODEL,
aws_profile_name="integration-profile-must-not-be-read",
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=wire.url,
messages=[{"role": "user", "content": "synthetic credential control"}],
timeout=5,
num_retries=0,
)
assert wire.drain() == ()
for source in ("argument", "environment"):
if source == "environment":
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", TOKEN)
parameters: Final = {
"model": MODEL, "api_key": TOKEN if source == "argument" else None,
"aws_region_name": "us-east-1", "aws_profile_name": "integration-profile-must-not-be-read",
"aws_bedrock_runtime_endpoint": wire.url, "timeout": 5, "num_retries": 0,
"messages": [{"role": "system", "content": "synthetic system"}, {"role": "user", "content": "synthetic bearer request"}],
"model": MODEL,
"api_key": TOKEN if source == "argument" else None,
"aws_region_name": "us-east-1",
"aws_profile_name": "integration-profile-must-not-be-read",
"aws_bedrock_runtime_endpoint": wire.url,
"timeout": 5,
"num_retries": 0,
"messages": [
{"role": "system", "content": "synthetic system"},
{"role": "user", "content": "synthetic bearer request"},
],
"max_tokens": 16,
}
for asynchronous in (False, True):
result: Final = await litellm.acompletion(**parameters) if asynchronous else await asyncio.to_thread(litellm.completion, **parameters)
result: Final = (
await litellm.acompletion(**parameters)
if asynchronous
else await asyncio.to_thread(litellm.completion, **parameters)
)
assert result.choices[0].message.content == "bedrock wire control"
assert result.choices[0].finish_reason == "stop"
assert result.usage.prompt_tokens == 11 and result.usage.completion_tokens == 4
@ -66,28 +97,57 @@ async def test_bearer_only_sdk_sync_async_requests_do_not_require_aws_credential
@pytest.mark.covers("other.provider_wire.bedrock.bearer_db_yaml_survives_reload")
def test_bearer_environment_reference_loads_from_db_and_yaml_and_survives_reload(gateway: Gateway, tmp_path: Path) -> None:
def test_bearer_environment_reference_loads_from_db_and_yaml_and_survives_reload(
gateway: Gateway, tmp_path: Path
) -> None:
empty: Final = tmp_path / "empty-aws-config"
empty.write_text("")
with wire_server(bearer_peer) as wire:
parameters: Final = {
"model": MODEL, "api_key": "os.environ/INTEGRATION_BEARER_TOKEN", "aws_region_name": "us-east-1",
"aws_profile_name": "integration-profile-must-not-be-read", "aws_bedrock_runtime_endpoint": wire.url,
"model": MODEL,
"api_key": "os.environ/INTEGRATION_BEARER_TOKEN",
"aws_region_name": "us-east-1",
"aws_profile_name": "integration-profile-must-not-be-read",
"aws_bedrock_runtime_endpoint": wire.url,
}
alias: Final = f"integration-yaml-{uuid.uuid4().hex}"
configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
configuration["model_list"] = [{"model_name": alias, "litellm_params": parameters, "model_info": {"id": alias}}]
path: Final = tmp_path / "bedrock.yaml"
path.write_text(yaml.safe_dump(configuration))
overrides: Final = {"INTEGRATION_BEARER_TOKEN": TOKEN, "AWS_CONFIG_FILE": str(empty), "AWS_SHARED_CREDENTIALS_FILE": str(empty), "AWS_EC2_METADATA_DISABLED": "true", "LITELLM_RUST": "false"}
with owned_proxy(gateway, tmp_path, overrides, config=path, remove_environment=tuple(name for name in os.environ if name.startswith("AWS_"))) as candidate, candidate.scenario() as scenario:
overrides: Final = {
"INTEGRATION_BEARER_TOKEN": TOKEN,
"AWS_CONFIG_FILE": str(empty),
"AWS_SHARED_CREDENTIALS_FILE": str(empty),
"AWS_EC2_METADATA_DISABLED": "true",
"LITELLM_RUST": "false",
}
with (
owned_proxy(
gateway,
tmp_path,
overrides,
config=path,
remove_environment=tuple(name for name in os.environ if name.startswith("AWS_")),
) as candidate,
candidate.scenario() as scenario,
):
database_model: Final = scenario.model(**parameters)
for generation in range(2):
for model in (alias, database_model):
response: Final = candidate.request("POST", "/v1/chat/completions", {
"model": model, "messages": [{"role": "system", "content": "synthetic system"}, {"role": "user", "content": "synthetic bearer request"}],
"max_tokens": 16, "cache": {"no-cache": True},
})
response: Final = candidate.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [
{"role": "system", "content": "synthetic system"},
{"role": "user", "content": "synthetic bearer request"},
],
"max_tokens": 16,
"cache": {"no-cache": True},
},
)
assert response.status_code == 200, response.text
assert response.json()["choices"][0]["message"]["content"] == "bedrock wire control"
assert response.json()["usage"]["total_tokens"] == 15
@ -95,5 +155,61 @@ def test_bearer_environment_reference_loads_from_db_and_yaml_and_survives_reload
if generation == 0:
entries: Final = candidate.get("/model/info")["data"]
target: Final = next(entry for entry in entries if entry["model_name"] == database_model)
response: Final = candidate.request("PATCH", f"/model/{target['model_info']['id']}/update", {"model_info": {"description": "bearer reload"}})
response: Final = candidate.request(
"PATCH",
f"/model/{target['model_info']['id']}/update",
{"model_info": {"description": "bearer reload"}},
)
assert response.status_code == 200, response.text
INVOKE_MODEL: Final = "bedrock/invoke/anthropic.claude-3-haiku-20240307-v1:0"
INVOKE_RESPONSE: Final = json.dumps(
{
"id": "msg_synthetic",
"type": "message",
"role": "assistant",
"model": "anthropic.claude-3-haiku-20240307-v1:0",
"content": [{"type": "text", "text": "bedrock invoke wire control"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 11, "output_tokens": 4},
}
).encode()
def sigv4_invoke_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/model/anthropic.claude-3-haiku-20240307-v1:0/invoke"
assert request.headers["authorization"].startswith(f"AWS4-HMAC-SHA256 Credential={ACCESS_KEY}/"), dict(
request.headers
)
assert CLIENT_OAUTH_TOKEN not in request.headers.values(), dict(request.headers)
assert json.loads(request.body)["messages"] == [{"role": "user", "content": "synthetic oauth isolation request"}]
return Reply(body=INVOKE_RESPONSE)
@pytest.mark.covers("providers.bedrock_auth.client_anthropic_oauth_token_never_replaces_sigv4_authorization")
def test_client_anthropic_oauth_authorization_header_does_not_replace_bedrock_sigv4_signature(gateway: Gateway) -> None:
with wire_server(sigv4_invoke_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=INVOKE_MODEL,
api_key=None,
aws_access_key_id=ACCESS_KEY,
aws_secret_access_key="synthetic-secret-key-for-testing",
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=wire.url,
api_base=wire.url,
)
response: Final = gateway.request(
"POST",
"/v1/messages",
{
"model": model,
"messages": [{"role": "user", "content": "synthetic oauth isolation request"}],
"max_tokens": 16,
},
headers={"Authorization": CLIENT_OAUTH_TOKEN, "x-litellm-api-key": f"Bearer {gateway.key}"},
)
assert response.status_code == 200, response.text
assert response.json()["content"] == [{"type": "text", "text": "bedrock invoke wire control"}], response.text
assert len(wire.drain()) == 1, response.text

78
tests/integration/providers/test_bedrock_batch_files_wire.py Normal file
View file

@ -0,0 +1,78 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
MODEL: Final = "bedrock/anthropic.claude-3-haiku-20240307-v1:0"
BUCKET: Final = "integration-batch-bucket"
PROMPT: Final = "synthetic completions prompt"
RESPONSES_INPUT: Final = "synthetic responses input"
INPUT_LINES: Final = (
{
"custom_id": "completions-record",
"method": "POST",
"url": "/v1/completions",
"body": {"model": MODEL, "prompt": PROMPT, "max_tokens": 64},
},
{
"custom_id": "responses-record",
"method": "POST",
"url": "/v1/responses",
"body": {"model": MODEL, "input": RESPONSES_INPUT, "max_output_tokens": 16},
},
)
EXPECTED_S3_OBJECT: Final = (
{
"recordId": "completions-record",
"modelInput": {
"messages": [{"role": "user", "content": [{"type": "text", "text": PROMPT}]}],
"max_tokens": 64,
"anthropic_version": "bedrock-2023-05-31",
},
},
{
"recordId": "responses-record",
"modelInput": {
"messages": [{"role": "user", "content": [{"type": "text", "text": RESPONSES_INPUT}]}],
"max_tokens": 16,
"anthropic_version": "bedrock-2023-05-31",
},
},
)
def s3_peer(request: Request) -> Reply:
assert request.method == "PUT" and request.target.startswith(f"/{BUCKET}/"), request.target
assert request.headers["authorization"].startswith("AWS4-HMAC-SHA256 ")
return Reply(body=b"")
@pytest.mark.covers(
"other.provider_wire.bedrock.batch_file_completions_and_responses_records_reach_s3_as_user_messages"
)
def test_completions_and_responses_batch_records_upload_as_anthropic_user_messages(gateway: Gateway) -> None:
with wire_server(s3_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=MODEL,
api_key=None,
api_base=None,
aws_access_key_id="AKIAIOSFODNN7EXAMPLE",
aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
aws_region_name="us-east-1",
s3_bucket_name=BUCKET,
s3_endpoint_url=wire.url,
)
jsonl: Final = "\n".join(json.dumps(line, separators=(",", ":")) for line in INPUT_LINES) + "\n"
response: Final = gateway.request_multipart(
"/v1/files",
{"purpose": "batch", "model": model},
{"file": ("in.jsonl", jsonl.encode(), "application/jsonl")},
)
assert response.status_code == 200, response.text
assert response.json()["object"] == "file" and response.json()["purpose"] == "batch", response.text
uploads: Final = wire.drain()
assert len(uploads) == 1, f"Expected exactly one S3 PUT, saw {[upload.target for upload in uploads]}"
stored: Final = tuple(json.loads(line) for line in uploads[0].body.decode().splitlines() if line.strip())
assert stored == EXPECTED_S3_OBJECT

60
tests/integration/providers/test_bedrock_claude_thinking_wire.py Normal file
View file

@ -0,0 +1,60 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
MODEL: Final = "bedrock/invoke/us.anthropic.claude-opus-4-8"
TOKEN: Final = "synthetic-bedrock-bearer"
RESPONSE: Final = json.dumps(
{
"id": "msg_adaptive_control",
"type": "message",
"role": "assistant",
"model": "us.anthropic.claude-opus-4-8",
"content": [{"type": "text", "text": "adaptive thinking control"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 12, "output_tokens": 5},
}
).encode()
def adaptive_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/model/us.anthropic.claude-opus-4-8/invoke"
assert request.headers["authorization"] == f"Bearer {TOKEN}"
body: Final = json.loads(request.body)
assert body["messages"] == [{"role": "user", "content": [{"type": "text", "text": "synthetic effort request"}]}]
assert body["thinking"]["type"] == "adaptive", body
assert body["output_config"] == {"effort": "high"}, body
assert "budget_tokens" not in json.dumps(body), body
return Reply(body=RESPONSE)
@pytest.mark.covers("other.provider_wire.bedrock.prefixed_opus_4_8_reasoning_effort_sends_adaptive_thinking")
def test_prefixed_opus_4_8_reasoning_effort_reaches_bedrock_as_adaptive_thinking_not_budget_tokens(
gateway: Gateway,
) -> None:
with wire_server(adaptive_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=MODEL,
api_key=TOKEN,
aws_region_name="us-east-1",
api_base=wire.url,
aws_bedrock_runtime_endpoint=wire.url,
)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": "synthetic effort request"}],
"max_tokens": 4096,
"reasoning_effort": "high",
},
)
assert response.status_code == 200, response.text
assert response.json()["choices"][0]["message"]["content"] == "adaptive thinking control"
assert response.json()["usage"]["prompt_tokens"] == 12 and response.json()["usage"]["completion_tokens"] == 5
assert len(wire.drain()) == 1

46
tests/integration/providers/test_bedrock_converse_config_blocks_wire.py Normal file
View file

@ -0,0 +1,46 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from integration.providers.test_bedrock_auth_wire import MODEL, RESPONSE, TOKEN
GUARDRAIL: Final = {"guardrailIdentifier": "integration-guardrail", "guardrailVersion": "DRAFT", "trace": "enabled"}
PERFORMANCE: Final = {"latency": "optimized"}
def converse_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/model/anthropic.claude-3-haiku-20240307-v1%3A0/converse"
body: Final = json.loads(request.body)
assert body["inferenceConfig"] == {"maxTokens": 16, "temperature": 0.2}, body
assert body["guardrailConfig"] == GUARDRAIL, body
assert body["performanceConfig"] == PERFORMANCE, body
return Reply(body=RESPONSE)
@pytest.mark.covers("other.provider_wire.bedrock.converse_config_blocks_sent_once_at_top_level")
def test_guardrail_and_performance_config_are_not_duplicated_inside_inference_config(gateway: Gateway) -> None:
with wire_server(converse_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=MODEL,
api_key=TOKEN,
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=wire.url,
guardrailConfig=GUARDRAIL,
)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": "synthetic guardrail request"}],
"max_tokens": 16,
"temperature": 0.2,
"performanceConfig": PERFORMANCE,
"cache": {"no-cache": True},
},
)
assert response.status_code == 200, response.text
assert response.json()["choices"][0]["message"]["content"] == "bedrock wire control"
assert len(wire.drain()) == 1, response.text

58
tests/integration/providers/test_bedrock_embedding_wire.py Normal file
View file

@ -0,0 +1,58 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
MODEL: Final = "bedrock/cohere.embed-english-v3"
TOKEN: Final = "synthetic-bedrock-bearer"
INPUT: Final = "hello world"
VECTOR: Final = [0.1, 0.2, 0.3]
RESPONSE: Final = json.dumps(
{
"embeddings": {"float": [VECTOR]},
"id": "synthetic-cohere-embed",
"response_type": "embeddings_by_type",
"texts": [INPUT],
}
).encode()
def cohere_english_v3_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/model/cohere.embed-english-v3/invoke"
assert request.headers["authorization"] == f"Bearer {TOKEN}"
assert json.loads(request.body) == {
"texts": [INPUT],
"input_type": "search_document",
"embedding_types": ["float"],
"output_dimension": 512,
}
return Reply(body=RESPONSE)
@pytest.mark.covers("other.provider_wire.bedrock.cohere_embed_english_v3_accepts_encoding_format")
def test_cohere_embed_english_v3_accepts_encoding_format_and_dimensions(gateway: Gateway) -> None:
with wire_server(cohere_english_v3_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=MODEL,
api_key=TOKEN,
api_base=wire.url,
aws_region_name="us-east-1",
)
for encoding_format in ("float", "base64"):
response: Final = gateway.request(
"POST",
"/v1/embeddings",
{
"model": model,
"input": INPUT,
"encoding_format": encoding_format,
"dimensions": 512,
},
)
assert response.status_code == 200, f"encoding_format={encoding_format}: {response.text}"
assert response.json()["data"] == [
{"object": "embedding", "index": 0, "embedding": VECTOR, "type": "float"},
], response.text
assert len(wire.drain()) == 1, f"encoding_format={encoding_format} never reached Bedrock"

50
tests/integration/providers/test_bedrock_gpt5_reasoning_wire.py Normal file
View file

@ -0,0 +1,50 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
MODEL: Final = "bedrock/converse/us.openai.gpt-5.6-sol"
TOKEN: Final = "synthetic-bedrock-bearer"
RESPONSE: Final = json.dumps(
{
"output": {"message": {"role": "assistant", "content": [{"text": "gpt-5 reasoning wire control"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 9, "outputTokens": 5, "totalTokens": 14},
"metrics": {"latencyMs": 1},
}
).encode()
def gpt5_converse_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/model/us.openai.gpt-5.6-sol/converse"
assert request.headers["authorization"] == f"Bearer {TOKEN}"
body: Final = json.loads(request.body)
assert body["messages"] == [{"role": "user", "content": [{"text": "synthetic reasoning request"}]}]
assert body["additionalModelRequestFields"] == {"reasoning": {"effort": "high"}}, body
assert body["inferenceConfig"] == {"maxTokens": 16}, body
return Reply(body=RESPONSE)
@pytest.mark.covers("providers.bedrock_converse.gpt5_reasoning_effort_reaches_provider_as_reasoning_effort")
def test_gpt5_reasoning_effort_is_accepted_and_sent_as_converse_reasoning_effort(gateway: Gateway) -> None:
with wire_server(gpt5_converse_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=MODEL, api_key=TOKEN, aws_region_name="us-east-1", aws_bedrock_runtime_endpoint=wire.url
)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": "synthetic reasoning request"}],
"reasoning_effort": "high",
"max_tokens": 16,
"cache": {"no-cache": True},
},
)
assert response.status_code == 200, response.text
assert response.json()["choices"][0]["message"]["content"] == "gpt-5 reasoning wire control", response.text
assert response.json()["usage"]["total_tokens"] == 14, response.text
assert len(wire.drain()) == 1

88
tests/integration/providers/test_bedrock_invoke_tool_search_wire.py Normal file
View file

@ -0,0 +1,88 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
MODEL_ID: Final = "us.anthropic.claude-sonnet-5"
TOKEN: Final = "synthetic-bedrock-bearer"
TOOL_SEARCH_TOOL: Final = {"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}
DEFERRED_TOOL: Final = {
"name": "get_weather",
"description": "Weather lookup",
"input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
"defer_loading": True,
}
RESPONSE: Final = json.dumps(
{
"id": "msg_tool_search_control",
"type": "message",
"role": "assistant",
"model": MODEL_ID,
"content": [
{
"type": "server_tool_use",
"id": "srvtoolu_control",
"name": "tool_search_tool_regex",
"input": {"pattern": "weather"},
},
{
"type": "tool_search_tool_result",
"tool_use_id": "srvtoolu_control",
"content": {
"type": "tool_search_tool_search_result",
"tool_references": [{"type": "tool_reference", "tool_name": "get_weather"}],
},
},
{"type": "text", "text": "tool search wire control"},
],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 12, "output_tokens": 6},
}
).encode()
def tool_search_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == f"/model/{MODEL_ID}/invoke", request.target
assert request.headers["authorization"] == f"Bearer {TOKEN}"
body: Final = json.loads(request.body)
assert body["anthropic_beta"] == ["tool-search-tool-2025-10-19"], body
assert body["messages"] == [{"role": "user", "content": "find the weather tool"}]
assert body["tools"] == [TOOL_SEARCH_TOOL, DEFERRED_TOOL], body["tools"]
assert body["max_tokens"] == 64
assert "model" not in body
return Reply(body=RESPONSE)
@pytest.mark.covers("providers.bedrock_invoke.tool_search_gen5_claude_sends_bedrock_beta_and_reports_support")
def test_gen5_claude_bedrock_invoke_messages_tool_search_sends_bedrock_beta_field(gateway: Gateway) -> None:
with wire_server(tool_search_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"bedrock/invoke/{MODEL_ID}",
api_key=TOKEN,
aws_region_name="us-east-1",
api_base=wire.url,
)
response: Final = gateway.request(
"POST",
"/v1/messages",
{
"model": model,
"max_tokens": 64,
"messages": [{"role": "user", "content": "find the weather tool"}],
"tools": [TOOL_SEARCH_TOOL, DEFERRED_TOOL],
},
)
assert response.status_code == 200, response.text
body: Final = response.json()
assert body["content"][2] == {"type": "text", "text": "tool search wire control"}, response.text
assert body["stop_reason"] == "end_turn"
assert body["usage"]["input_tokens"] == 12 and body["usage"]["output_tokens"] == 6
assert len(wire.drain()) == 1
entries: Final = gateway.get("/v1/model/info")["data"]
assert isinstance(entries, list)
info: Final = next(entry for entry in entries if isinstance(entry, dict) and entry["model_name"] == model)
assert isinstance(info["model_info"], dict)
assert info["model_info"]["supports_tool_search"] is True, info["model_info"]

88
tests/integration/providers/test_bedrock_mantle_codex_input_wire.py Normal file
View file

@ -0,0 +1,88 @@
import json
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
MODEL: Final = "bedrock_mantle/openai.gpt-5.6-sol"
TOKEN: Final = "synthetic-mantle-bearer"
CIPHERTEXT: Final = "synthetic-compaction-ciphertext"
CALL_ID: Final = "call_synthetic_shell"
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
ACTION: Final[dict[str, JsonValue]] = {"type": "exec", "command": ["ls", "-la"], "timeout_ms": 1000}
RESPONSE: Final = json.dumps(
{
"id": "resp_synthetic_mantle",
"object": "response",
"created_at": 1789788253,
"status": "completed",
"model": "openai.gpt-5.6-sol",
"output": [
{
"type": "message",
"id": "msg_synthetic_mantle",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "mantle wire control", "annotations": []}],
}
],
"usage": {"input_tokens": 21, "output_tokens": 4, "total_tokens": 25},
}
).encode()
def user_turn(text: str) -> JsonValue:
return {"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]}
def codex_history(marker: str) -> tuple[JsonValue, ...]:
return (
user_turn(f"first turn {marker}"),
{"type": "agent_message", "role": "assistant", "content": [{"type": "output_text", "text": "sub-agent reply"}]},
{"type": "context_compaction", "encrypted_content": CIPHERTEXT},
{"type": "local_shell_call", "call_id": CALL_ID, "status": "completed", "action": ACTION},
{"type": "function_call_output", "call_id": CALL_ID, "output": "synthetic shell output"},
user_turn(f"next turn {marker}"),
)
def mantle_history(marker: str) -> tuple[JsonValue, ...]:
return (
user_turn(f"first turn {marker}"),
{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "sub-agent reply"}]},
{"type": "compaction", "encrypted_content": CIPHERTEXT},
{"type": "function_call", "call_id": CALL_ID, "name": "local_shell", "arguments": json.dumps(ACTION)},
{"type": "function_call_output", "call_id": CALL_ID, "output": "synthetic shell output"},
user_turn(f"next turn {marker}"),
)
@pytest.mark.covers("other.provider_wire.bedrock_mantle.codex_history_items_reach_mantle_as_supported_types")
def test_codex_agent_message_context_compaction_and_local_shell_call_reach_mantle_as_supported_items(
gateway: Gateway,
) -> None:
marker: Final = uuid.uuid4().hex
expected_input: Final = list(mantle_history(marker))
def mantle_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/openai/v1/responses", request.target
assert request.headers["authorization"] == f"Bearer {TOKEN}"
body: Final = JSON_OBJECT.validate_json(request.body)
assert body["model"] == "openai.gpt-5.6-sol", body
assert body["input"] == expected_input, body["input"]
return Reply(body=RESPONSE)
with wire_server(mantle_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=MODEL, api_key=TOKEN, api_base=wire.url, aws_region_name="us-east-2")
response: Final = gateway.request(
"POST", "/v1/responses", {"model": model, "input": list(codex_history(marker)), "store": False}
)
assert response.status_code == 200, response.text
assert response.json()["output"][0]["content"][0]["text"] == "mantle wire control", response.text
assert response.json()["usage"]["total_tokens"] == 25, response.text
forwarded: Final = wire.drain()
assert len(forwarded) == 1, forwarded
assert JSON_OBJECT.validate_json(forwarded[0].body)["input"] == expected_input, forwarded[0].body

106
tests/integration/providers/test_bedrock_mantle_responses_wire.py Normal file
View file

@ -0,0 +1,106 @@
import json
from collections.abc import Callable
from typing import Final
from uuid import uuid4
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_MODEL: Final = "bedrock_mantle/openai.gpt-5.6-sol"
_TOKEN: Final = "synthetic-mantle-bearer"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_SHELL_ACTION: Final[dict[str, JsonValue]] = {"type": "exec", "command": ["ls", "-la"], "timeout_ms": 1000}
_OUTPUT_MESSAGE: Final[dict[str, JsonValue]] = {
"type": "message",
"id": "msg_mantle",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "mantle wire control", "annotations": []}],
}
_RESPONSE: Final = json.dumps(
{
"id": "resp_mantle",
"object": "response",
"status": "completed",
"created_at": 1700000000,
"model": "gpt-5.6-sol",
"output": [_OUTPUT_MESSAGE],
"usage": {
"input_tokens": 11,
"output_tokens": 4,
"total_tokens": 15,
"input_tokens_details": {"cached_tokens": 0},
"output_tokens_details": {"reasoning_tokens": 0},
},
}
).encode()
def _codex_history(marker: str) -> list[JsonValue]:
return [
{"type": "message", "role": "user", "content": f"delegate to a subagent {marker}"},
{
"type": "agent_message",
"id": "msg_agent",
"content": [{"type": "text", "text": "sub-agent said "}, {"type": "text", "encrypted_content": "hello"}],
},
{"type": "context_compaction", "id": "cmp_1", "encrypted_content": "compacted-history"},
{
"type": "local_shell_call",
"id": "lsc_1",
"call_id": "call_shell",
"status": "completed",
"action": _SHELL_ACTION,
},
{"type": "function_call_output", "call_id": "call_shell", "output": "total 0"},
]
def _mantle_history(marker: str) -> list[JsonValue]:
return [
{"type": "message", "role": "user", "content": f"delegate to a subagent {marker}"},
{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "sub-agent said hello"}]},
{"type": "compaction", "encrypted_content": "compacted-history"},
{
"type": "function_call",
"call_id": "call_shell",
"name": "local_shell",
"arguments": json.dumps(_SHELL_ACTION),
},
{"type": "function_call_output", "call_id": "call_shell", "output": "total 0"},
]
def _mantle_peer(marker: str) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/openai/v1/responses", request.target
assert request.headers["authorization"] == f"Bearer {_TOKEN}"
body: Final = _JSON_OBJECT.validate_json(request.body)
assert body["input"] == _mantle_history(marker), json.dumps(body["input"])
return Reply(body=_RESPONSE)
return respond
@pytest.mark.covers("providers.bedrock_mantle.codex_history_items_reach_the_wire_as_supported_input_items")
def test_codex_agent_message_compaction_and_local_shell_items_are_rewritten_for_mantle(gateway: Gateway) -> None:
marker: Final = uuid4().hex
with wire_server(_mantle_peer(marker)) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_TOKEN, aws_region_name="us-east-1")
response: Final = gateway.request(
"POST", "/v1/responses", {"model": model, "input": _codex_history(marker), "stream": False}
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["output"] == [
{
**_OUTPUT_MESSAGE,
"phase": None,
"content": [
{"type": "output_text", "text": "mantle wire control", "annotations": [], "logprobs": None}
],
}
], response.text
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/openai/v1/responses")]

53
tests/integration/providers/test_bedrock_mantle_wire.py Normal file
View file

@ -0,0 +1,53 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_BACKEND: Final = "openai.gpt-5.6-sol"
_API_KEY: Final = "synthetic-mantle-bearer"
_PROMPT: Final = "synthetic long conversation control"
_PROMPT_TOKENS: Final = 1055489
_MODEL_MAXIMUM: Final = 1050000
_RESPONSES_PATH: Final = "/openai/v1/responses"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_OVERFLOW_BODY: Final = json.dumps(
{
"error": {
"code": "validation_error",
"message": f"prompt tokens ({_PROMPT_TOKENS}) exceed model maximum ({_MODEL_MAXIMUM}) for {_BACKEND}",
"type": "invalid_request_error",
}
}
).encode()
def _overflow_peer(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == _RESPONSES_PATH
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
body: Final = _JSON_OBJECT.validate_json(request.body)
assert body["model"] == _BACKEND
assert _PROMPT in json.dumps(body["input"]), body
return Reply(status=400, body=_OVERFLOW_BODY)
@pytest.mark.covers("other.provider_wire.bedrock_mantle.context_overflow_is_reported_as_prompt_too_long")
def test_bedrock_mantle_context_overflow_returns_400_saying_prompt_is_too_long(gateway: Gateway) -> None:
with wire_server(_overflow_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"bedrock_mantle/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": _PROMPT}]},
)
assert response.status_code == 400, response.text
error: Final = _JSON_OBJECT.validate_json(response.content)["error"]
assert isinstance(error, dict), response.text
assert error["code"] == "400", response.text
message: Final = error["message"]
assert isinstance(message, str), response.text
assert f"prompt is too long: {_PROMPT_TOKENS} tokens > {_MODEL_MAXIMUM} maximum" in message, response.text
assert [(request.method, request.target) for request in wire.drain()] == [("POST", _RESPONSES_PATH)]

88
tests/integration/providers/test_bedrock_messages_web_search_replay_wire.py Normal file
View file

@ -0,0 +1,88 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
BEDROCK_MODEL: Final = "us.anthropic.claude-opus-5-v1:0"
TOKEN: Final = "synthetic-bedrock-bearer"
SNIPPET: Final = "synthetic snippet about the integration harness"
INTERCEPTED_TURN: Final = (
{"type": "server_tool_use", "id": "srvtoolu_synthetic", "name": "web_search", "input": {"query": "harness docs"}},
{
"type": "web_search_tool_result",
"tool_use_id": "srvtoolu_synthetic",
"content": [
{
"type": "web_search_result",
"url": "https://example.test/harness",
"title": "Harness",
"page_age": None,
"encrypted_content": "",
"snippet": SNIPPET,
},
],
},
{"type": "text", "text": "The harness is documented at example.test"},
)
FLATTENED_TURN: Final = (
{
"type": "text",
"text": f"Web search results for 'harness docs':\n\nTitle: Harness\nURL: https://example.test/harness\nSnippet: {SNIPPET}",
},
{"type": "text", "text": "The harness is documented at example.test"},
)
REPLY: Final = json.dumps(
{
"id": "msg_synthetic_replay",
"type": "message",
"role": "assistant",
"model": BEDROCK_MODEL,
"content": [{"type": "text", "text": "replay accepted"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 30, "output_tokens": 3},
}
).encode()
def bedrock_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == f"/model/{BEDROCK_MODEL}/invoke"
assert request.headers["authorization"] == f"Bearer {TOKEN}"
body: Final = json.loads(request.body)
assert body["messages"] == [
{"role": "user", "content": "where is the harness documented"},
{"role": "assistant", "content": list(FLATTENED_TURN)},
{"role": "user", "content": "and what does it say"},
], request.body.decode()
assert "tools" not in body, request.body.decode()
return Reply(body=REPLY)
@pytest.mark.covers("providers.bedrock_messages.replayed_intercepted_web_search_turn_is_flattened_to_text")
def test_replayed_intercepted_web_search_turn_reaches_bedrock_as_text_and_answers(gateway: Gateway) -> None:
with wire_server(bedrock_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"bedrock/{BEDROCK_MODEL}",
api_key=TOKEN,
api_base=wire.url,
aws_region_name="us-east-1",
)
response: Final = gateway.request(
"POST",
"/v1/messages",
{
"model": model,
"max_tokens": 64,
"messages": [
{"role": "user", "content": "where is the harness documented"},
{"role": "assistant", "content": list(INTERCEPTED_TURN)},
{"role": "user", "content": "and what does it say"},
],
},
headers={"x-api-key": gateway.key, "anthropic-version": "2023-06-01"},
)
assert response.status_code == 200, response.text
assert response.json()["content"] == [{"type": "text", "text": "replay accepted"}], response.text
assert len(wire.drain()) == 1

42
tests/integration/providers/test_bedrock_passthrough_stream_wire.py Normal file
View file

@ -0,0 +1,42 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.upstream import _aws_event_frame
from integration._support.wire import Reply, Request, wire_server
_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0"
_EVENT_STREAM: Final = "application/vnd.amazon.eventstream"
_REQUEST_BODY: Final = {"messages": [{"role": "user", "content": [{"text": "synthetic passthrough stream"}]}]}
_EVENTS: Final = (
("messageStart", {"role": "assistant"}),
("contentBlockDelta", {"delta": {"text": "bedrock stream control"}, "contentBlockIndex": 0}),
("messageStop", {"stopReason": "end_turn"}),
("metadata", {"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15}}),
)
_STREAM_BYTES: Final = b"".join(_aws_event_frame(kind, payload, "sc", "u") for kind, payload in _EVENTS)
def event_stream_peer(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == f"/model/{_MODEL_ID}/converse-stream"
assert json.loads(request.body)["messages"] == _REQUEST_BODY["messages"]
return Reply(body=_STREAM_BYTES, content_type=_EVENT_STREAM)
@pytest.mark.covers("other.provider_wire.bedrock.passthrough_stream_keeps_event_stream_content_type")
def test_bedrock_passthrough_converse_stream_response_carries_event_stream_content_type(gateway: Gateway) -> None:
with wire_server(event_stream_peer) as wire, gateway.scenario() as scenario:
deployment: Final = scenario.model(
model=f"bedrock/{_MODEL_ID}",
api_base=wire.url,
aws_access_key_id="AKIASCRIPTEDPROVIDER",
aws_secret_access_key="scripted-secret",
aws_region_name="us-east-1",
)
response: Final = gateway.request("POST", f"/bedrock/model/{deployment}/converse-stream", _REQUEST_BODY)
assert response.status_code == 200, response.text
assert len(wire.drain()) == 1, response.text
assert response.headers.get("content-type") == _EVENT_STREAM, dict(response.headers)
assert response.content == _STREAM_BYTES, response.text

92
tests/integration/providers/test_bedrock_rerank_wire.py Normal file
View file

@ -0,0 +1,92 @@
import json
import os
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import Gateway
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
MODEL: Final = "bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0"
ACCESS_KEY: Final = "AKIAINTEGRATION000002"
FORWARDED_FOR: Final = "203.0.113.5"
RESPONSE: Final = json.dumps(
{"results": [{"index": 1, "relevanceScore": 0.9}, {"index": 0, "relevanceScore": 0.1}]}
).encode()
def signed_headers(authorization: str) -> tuple[str, ...]:
return tuple(authorization.split("SignedHeaders=")[1].split(",")[0].split(";"))
def rerank_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/rerank"
assert request.headers["authorization"].startswith(f"AWS4-HMAC-SHA256 Credential={ACCESS_KEY}/")
assert signed_headers(request.headers["authorization"]) == ("content-type", "host", "x-amz-date"), request.headers[
"authorization"
]
assert request.headers["x-forwarded-for"] == FORWARDED_FOR
body: Final = json.loads(request.body)
assert body["queries"] == [{"textQuery": {"text": "synthetic rerank query"}, "type": "TEXT"}]
assert body["rerankingConfiguration"]["bedrockRerankingConfiguration"]["modelConfiguration"] == {
"modelArn": "arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0"
}
assert body["rerankingConfiguration"]["bedrockRerankingConfiguration"]["numberOfResults"] == 2
return Reply(body=RESPONSE)
@pytest.mark.covers("providers.bedrock_rerank.forwarded_client_headers_are_sent_unsigned")
def test_forwarded_client_header_on_rerank_is_excluded_from_the_sigv4_signature(
gateway: Gateway, tmp_path: Path
) -> None:
empty: Final = tmp_path / "empty-aws-config"
empty.write_text("")
configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
configuration["general_settings"]["forward_client_headers_to_llm_api"] = True
path: Final = tmp_path / "forwarding.yaml"
path.write_text(yaml.safe_dump(configuration))
overrides: Final = {
"AWS_CONFIG_FILE": str(empty),
"AWS_SHARED_CREDENTIALS_FILE": str(empty),
"AWS_EC2_METADATA_DISABLED": "true",
"LITELLM_RUST": "false",
}
with wire_server(rerank_peer) as wire:
with (
owned_proxy(
gateway,
tmp_path,
overrides,
config=path,
remove_environment=tuple(name for name in os.environ if name.startswith("AWS_")),
) as candidate,
candidate.scenario() as scenario,
):
model: Final = scenario.model(
model=MODEL,
api_key=None,
api_base=None,
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=wire.url,
aws_access_key_id=ACCESS_KEY,
aws_secret_access_key="synthetic-rerank-secret-key-for-testing",
)
response: Final = candidate.request(
"POST",
"/v1/rerank",
{
"model": model,
"query": "synthetic rerank query",
"documents": ["first synthetic document", "second synthetic document"],
"top_n": 2,
},
headers={"x-forwarded-for": FORWARDED_FOR},
)
assert response.status_code == 200, response.text
assert response.json()["results"] == [
{"index": 1, "relevance_score": 0.9},
{"index": 0, "relevance_score": 0.1},
], response.text
assert len(wire.drain()) == 1

94
tests/integration/providers/test_bedrock_thinking_tokens_wire.py Normal file
View file

@ -0,0 +1,94 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
MODEL: Final = "bedrock/converse/global.anthropic.claude-opus-4-8"
TOKEN: Final = "synthetic-bedrock-bearer"
PROMPT: Final = "How many prime numbers are less than 30? Think it through, then answer with just the number."
RESPONSES_PROMPT: Final = "How many prime numbers are less than 30? Answer with just the number."
REDACTED_DATA: Final = "RWRhY3RlZC1ieS1CZWRyb2Nr"
INPUT_TOKENS: Final = 31
OUTPUT_TOKENS: Final = 257
RESPONSE: Final = json.dumps(
{
"output": {
"message": {
"role": "assistant",
"content": [{"reasoningContent": {"redactedContent": REDACTED_DATA}}, {"text": "10"}],
}
},
"stopReason": "end_turn",
"usage": {
"inputTokens": INPUT_TOKENS,
"outputTokens": OUTPUT_TOKENS,
"totalTokens": INPUT_TOKENS + OUTPUT_TOKENS,
},
"metrics": {"latencyMs": 1},
}
).encode()
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_JSON_LIST: Final = TypeAdapter(list[dict[str, JsonValue]])
def redacted_thinking_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/model/global.anthropic.claude-opus-4-8/converse"
assert request.headers["authorization"] == f"Bearer {TOKEN}"
body: Final = json.loads(request.body)
assert body["messages"] in (
[{"role": "user", "content": [{"text": PROMPT}]}],
[{"role": "user", "content": [{"text": RESPONSES_PROMPT}]}],
), body
assert body["additionalModelRequestFields"]["thinking"]["type"] == "adaptive", body
return Reply(body=RESPONSE)
@pytest.mark.covers("other.provider_wire.bedrock.hidden_thinking_tokens_are_not_reported_as_text")
def test_bedrock_redacted_thinking_is_not_reported_as_zero_reasoning_tokens(gateway: Gateway) -> None:
with wire_server(redacted_thinking_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=MODEL, api_key=TOKEN, aws_region_name="us-east-1", aws_bedrock_runtime_endpoint=wire.url
)
chat: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": PROMPT}],
"max_tokens": 4000,
"reasoning_effort": "max",
},
)
assert chat.status_code == 200, chat.text
chat_body: Final = _JSON_OBJECT.validate_json(chat.content)
message: Final = _JSON_OBJECT.validate_python(_JSON_LIST.validate_python(chat_body["choices"])[0]["message"])
assert message["content"] == "10", chat.text
assert message["thinking_blocks"] == [{"type": "redacted_thinking", "data": REDACTED_DATA}], chat.text
usage: Final = _JSON_OBJECT.validate_python(chat_body["usage"])
assert usage["completion_tokens"] == OUTPUT_TOKENS, chat.text
details: Final = _JSON_OBJECT.validate_python(usage["completion_tokens_details"])
assert details == {}, chat.text
assert len(wire.drain()) == 1
responses: Final = gateway.request(
"POST",
"/v1/responses",
{"model": model, "input": RESPONSES_PROMPT, "max_output_tokens": 4000, "reasoning": {"effort": "max"}},
)
assert responses.status_code == 200, responses.text
responses_body: Final = _JSON_OBJECT.validate_json(responses.content)
output: Final = _JSON_LIST.validate_python(responses_body["output"])
reasoning_items: Final = tuple(item for item in output if item["type"] == "reasoning")
assert len(reasoning_items) == 1, responses.text
assert reasoning_items[0]["encrypted_content"] == json.dumps(
[{"type": "redacted_thinking", "data": REDACTED_DATA}], separators=(",", ":")
), responses.text
responses_usage: Final = _JSON_OBJECT.validate_python(responses_body["usage"])
assert responses_usage["output_tokens"] == OUTPUT_TOKENS, responses.text
assert _JSON_OBJECT.validate_python(responses_usage["output_tokens_details"])["reasoning_tokens"] == 0, (
responses.text
)
assert len(wire.drain()) == 1

62
tests/integration/providers/test_dashscope_chat_wire.py Normal file
View file

@ -0,0 +1,62 @@
import json
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_BACKEND: Final = "qwen3.7-plus"
_API_KEY: Final = "synthetic-dashscope-key"
_PROMPT: Final = "What is 3^3?"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def _completion(identity: str) -> bytes:
return json.dumps(
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": _BACKEND,
"choices": [{"index": 0, "message": {"role": "assistant", "content": "27"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 17, "completion_tokens": 5, "total_tokens": 22},
}
).encode()
@pytest.mark.covers("other.provider_wire.dashscope.reasoning_effort_reaches_provider")
def test_dashscope_chat_forwards_reasoning_effort_none_to_the_provider(gateway: Gateway) -> None:
identity: Final = f"dashscope-reasoning-{uuid.uuid4().hex}"
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == "/chat/completions"
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
assert _JSON_OBJECT.validate_json(request.body) == {
"model": _BACKEND,
"messages": [{"role": "user", "content": _PROMPT}],
"reasoning_effort": "none",
}
return Reply(body=_completion(identity))
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"dashscope/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": _PROMPT}], "reasoning_effort": "none"},
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["id"] == identity
assert payload["choices"] == [
{
"finish_reason": "stop",
"index": 0,
"message": {"role": "assistant", "content": "27", "provider_specific_fields": {"refusal": None}},
"provider_specific_fields": {},
}
]
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]

129
tests/integration/providers/test_databricks_chat_wire.py Normal file
View file

@ -0,0 +1,129 @@
import json
import uuid
from collections.abc import Mapping
from typing import Final
import pytest
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.wire import Reply, Request, wire_server
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
_BACKEND: Final = "databricks-glm-5-2"
_API_KEY: Final = "synthetic-databricks-key"
_PROMPT: Final = "Summarise the cached briefing in one sentence."
_PROVIDER_USAGE: Final[Mapping[str, JsonValue]] = {
"prompt_tokens": 12011,
"completion_tokens": 8,
"total_tokens": 12019,
"cache_read_input_tokens": 12002,
"cache_creation_input_tokens": 0,
}
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
class _PromptTokensDetails(BaseModel):
model_config = ConfigDict(extra="ignore")
cached_tokens: int | None = None
class _Usage(BaseModel):
model_config = ConfigDict(extra="ignore")
prompt_tokens: int
completion_tokens: int
total_tokens: int
prompt_tokens_details: _PromptTokensDetails | None = None
class _Delta(BaseModel):
model_config = ConfigDict(extra="ignore")
content: str | None = None
class _Choice(BaseModel):
model_config = ConfigDict(extra="ignore")
delta: _Delta
class _Chunk(BaseModel):
model_config = ConfigDict(extra="ignore")
id: str
choices: tuple[_Choice, ...]
usage: _Usage | None = None
def _frame(identity: str, choices: list[Mapping[str, object]], usage: Mapping[str, JsonValue] | None = None) -> bytes:
value: Final = {
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": _BACKEND,
"choices": choices,
**({} if usage is None else {"usage": usage}),
}
return b"data: " + json.dumps(value).encode() + b"\n\n"
@pytest.mark.covers("other.provider_wire.databricks.stream_usage_and_cache_reads_reach_client_and_spend_log")
def test_databricks_stream_final_usage_chunk_reaches_client_and_spend_log(gateway: Gateway) -> None:
identity: Final = f"databricks-stream-{uuid.uuid4().hex}"
frames: Final = (
_frame(
identity, [{"index": 0, "delta": {"role": "assistant", "content": "The briefing "}, "finish_reason": None}]
),
_frame(identity, [{"index": 0, "delta": {"content": "is short."}, "finish_reason": None}]),
_frame(identity, [{"index": 0, "delta": {}, "finish_reason": "stop"}]),
_frame(identity, [], usage=_PROVIDER_USAGE),
b"data: [DONE]\n\n",
)
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == "/chat/completions"
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
body: Final = _JSON_OBJECT.validate_json(request.body)
assert body["model"] == _BACKEND
assert body["messages"] == [{"role": "user", "content": _PROMPT}]
assert body["stream"] is True
return Reply(content_type="text/event-stream", chunks=frames)
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"databricks/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
with gateway.client.stream(
"POST",
"/v1/chat/completions",
json={
"model": model,
"messages": [{"role": "user", "content": _PROMPT}],
"stream": True,
"stream_options": {"include_usage": True},
},
headers={"Authorization": f"Bearer {gateway.key}"},
) as response:
assert response.status_code == 200, response.read()
lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: "))
assert lines[-1] == "data: [DONE]", lines
chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1])
assert {chunk.id for chunk in chunks} == {identity}
assert (
"".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices)
== "The briefing is short."
)
usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None)
assert len(usages) == 1, lines
assert (
usages[0].prompt_tokens,
usages[0].completion_tokens,
usages[0].total_tokens,
usages[0].prompt_tokens_details.cached_tokens if usages[0].prompt_tokens_details is not None else None,
) == (12011, 8, 12019, 12002), lines
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
rows: Final = eventually(
lambda: read_rows(
'SELECT prompt_tokens, completion_tokens, total_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
(identity,),
),
lambda values: len(values) == 1,
seconds=70,
)
assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"], rows[0]["total_tokens"]) == (12011, 8, 12019)

92
tests/integration/providers/test_databricks_oauth_wire.py Normal file
View file

@ -0,0 +1,92 @@
import base64
import json
import uuid
from pathlib import Path
from typing import Final
from urllib.parse import parse_qs
import pytest
from integration._support.client import Gateway
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_MODEL: Final = "databricks/synthetic-vendor.chat-model.v1"
_CLIENT_ID: Final = "synthetic-databricks-client-id"
_CLIENT_SECRET: Final = "synthetic-databricks-client-secret"
_ACCESS_TOKEN: Final = "synthetic-databricks-oauth-token"
_PROMPT: Final = "Which workspace issued this token?"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def _basic_credentials(client_id: str, client_secret: str) -> str:
return "Basic " + base64.b64encode(f"{client_id}:{client_secret}".encode()).decode()
def _completion(identity: str) -> bytes:
return json.dumps(
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": _MODEL.removeprefix("databricks/"),
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "the workspace origin"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 9, "completion_tokens": 4, "total_tokens": 13},
}
).encode()
@pytest.mark.covers("other.provider_wire.databricks.oauth_token_url_uses_workspace_origin_for_ai_gateway_api_base")
def test_databricks_ai_gateway_api_base_requests_oauth_token_from_workspace_origin(
gateway: Gateway, tmp_path: Path
) -> None:
identity: Final = f"databricks-oauth-{uuid.uuid4().hex}"
def respond(request: Request) -> Reply:
if request.target == "/oidc/v1/token":
assert request.method == "POST"
assert request.headers["authorization"] == _basic_credentials(_CLIENT_ID, _CLIENT_SECRET)
assert request.headers["content-type"] == "application/x-www-form-urlencoded"
assert parse_qs(request.body.decode()) == {"grant_type": ["client_credentials"], "scope": ["all-apis"]}
return Reply(
body=json.dumps({"access_token": _ACCESS_TOKEN, "token_type": "Bearer", "expires_in": 3600}).encode()
)
if request.target == "/ai-gateway/mlflow/v1/chat/completions":
assert request.method == "POST"
assert request.headers["authorization"] == f"Bearer {_ACCESS_TOKEN}"
body: Final = _JSON_OBJECT.validate_json(request.body)
assert body["model"] == _MODEL.removeprefix("databricks/")
assert body["messages"] == [{"role": "user", "content": _PROMPT}]
return Reply(body=_completion(identity))
return Reply(status=401, body=json.dumps({"error": f"unauthenticated path {request.target}"}).encode())
overrides: Final = {"DATABRICKS_CLIENT_ID": _CLIENT_ID, "DATABRICKS_CLIENT_SECRET": _CLIENT_SECRET}
with wire_server(respond) as wire, owned_proxy(gateway, tmp_path, overrides) as candidate:
with candidate.scenario() as scenario:
model: Final = scenario.model(model=_MODEL, api_base=f"{wire.url}/ai-gateway/mlflow/v1", api_key=None)
response: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": _PROMPT}]},
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["id"] == identity
assert payload["choices"] == [
{
"finish_reason": "stop",
"index": 0,
"message": {"content": "the workspace origin", "role": "assistant"},
}
]
assert payload["usage"] == {"prompt_tokens": 9, "completion_tokens": 4, "total_tokens": 13}
assert [(request.method, request.target) for request in wire.drain()] == [
("POST", "/oidc/v1/token"),
("POST", "/ai-gateway/mlflow/v1/chat/completions"),
]

40
tests/integration/providers/test_deepseek_vision_wire.py Normal file
View file

@ -0,0 +1,40 @@
from typing import Final
import httpx
import pytest
from pydantic import JsonValue
from tests.integration._support.client import JSON_OBJECT, Gateway, object_value
_VISION_MODEL: Final = "deepseek-v4-flash-vision-exp"
_API_KEY: Final = "synthetic-deepseek-key"
_VISION_CONTENT: Final[JsonValue] = [
{"type": "text", "text": "what is in this image?"},
{"type": "image_url", "image_url": {"url": "https://example.com/pic.png"}},
]
@pytest.mark.covers("other.provider_wire.deepseek.vision_image_content_list_reaches_provider")
def test_deepseek_vision_forwards_image_url_content_list_instead_of_collapsing_to_text(gateway: Gateway) -> None:
with gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, trust_env=False) as upstream:
upstream.get("/__observations").raise_for_status()
model: Final = scenario.model(
model=f"deepseek/{_VISION_MODEL}",
api_key=_API_KEY,
model_info={"mode": "chat", "supports_vision": True},
)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": _VISION_CONTENT}]},
)
assert response.status_code == 200, response.text
observations: Final = JSON_OBJECT.validate_json(upstream.get("/__observations").content)["requests"]
assert isinstance(observations, list)
assert len(observations) == 1, response.text
observed: Final = object_value(observations[0])
assert observed["path"] == "/v1/chat/completions", response.text
assert observed["authorization"] == f"Bearer {_API_KEY}", response.text
body: Final = object_value(observed["body"])
assert body["model"] == _VISION_MODEL, response.text
assert body["messages"] == [{"role": "user", "content": _VISION_CONTENT}], response.text

84
tests/integration/providers/test_fireworks_ai_router_slug_wire.py Normal file
View file

@ -0,0 +1,84 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_ROUTER_SLUG: Final = "routers/glm-latest"
_ROUTER_RESOURCE: Final = "accounts/fireworks/routers/glm-latest"
_API_KEY: Final = "synthetic-fireworks-key"
_PROMPT: Final = "route me through the router"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def _provider_body(request: Request, target: str) -> dict[str, JsonValue]:
assert request.method == "POST"
assert request.target == target
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
return _JSON_OBJECT.validate_json(request.body)
@pytest.mark.covers("other.provider_wire.fireworks_ai.router_slug_chat_sends_router_resource_name")
def test_fireworks_router_slug_chat_sends_router_resource_not_models_path(gateway: Gateway) -> None:
def respond(request: Request) -> Reply:
body: Final = _provider_body(request, "/chat/completions")
assert body["model"] == _ROUTER_RESOURCE, body
assert body["messages"] == [{"role": "user", "content": _PROMPT}]
return Reply(
body=json.dumps(
{
"id": "fw-router-chat",
"object": "chat.completion",
"created": 1,
"model": _ROUTER_RESOURCE,
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "routed"}, "finish_reason": "stop"}
],
"usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6},
}
).encode()
)
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"fireworks_ai/{_ROUTER_SLUG}", api_base=wire.url, api_key=_API_KEY)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": _PROMPT}]},
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["choices"] == [
{"finish_reason": "stop", "index": 0, "message": {"role": "assistant", "content": "routed"}}
]
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
@pytest.mark.covers("other.provider_wire.fireworks_ai.router_slug_text_completion_sends_router_resource_name")
def test_fireworks_router_slug_text_completion_sends_router_resource_not_models_path(gateway: Gateway) -> None:
def respond(request: Request) -> Reply:
body: Final = _provider_body(request, "/completions")
assert body["model"] == _ROUTER_RESOURCE, body
assert body["prompt"] == _PROMPT
return Reply(
body=json.dumps(
{
"id": "fw-router-text",
"object": "text_completion",
"created": 1,
"model": _ROUTER_RESOURCE,
"choices": [{"index": 0, "text": "routed", "finish_reason": "stop", "logprobs": None}],
"usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6},
}
).encode()
)
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"fireworks_ai/{_ROUTER_SLUG}", api_base=wire.url, api_key=_API_KEY)
response: Final = gateway.request("POST", "/v1/completions", {"model": model, "prompt": _PROMPT})
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["choices"] == [{"index": 0, "text": "routed", "finish_reason": "stop", "logprobs": None}]
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/completions")]

66
tests/integration/providers/test_openai_chat_wire.py Normal file
View file

@ -0,0 +1,66 @@
import json
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_BACKEND: Final = "gpt-5.4-mini"
_API_KEY: Final = "synthetic-openai-key"
_PROMPT: Final = "Summarize this conversation in one sentence."
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def _completion(identity: str, content: str) -> bytes:
return json.dumps(
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": _BACKEND,
"choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 19, "completion_tokens": 7, "total_tokens": 26},
}
).encode()
@pytest.mark.covers("providers.openai_chat_wire.tool_choice_without_tools_is_dropped_before_the_wire")
def test_openai_chat_tool_choice_without_tools_is_not_forwarded(gateway: Gateway) -> None:
identity: Final = f"openai-toolless-{uuid.uuid4().hex}"
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == "/chat/completions"
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
body: Final = _JSON_OBJECT.validate_json(request.body)
assert body["model"] == _BACKEND
assert body["messages"] == [{"role": "user", "content": _PROMPT}]
assert "tool_choice" not in body, body
assert "tools" not in body, body
return Reply(body=_completion(identity, "One sentence."))
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"openai/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": _PROMPT}], "tool_choice": "none"},
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["id"] == identity
assert payload["choices"] == [
{
"finish_reason": "stop",
"index": 0,
"message": {
"role": "assistant",
"content": "One sentence.",
"provider_specific_fields": {"refusal": None},
},
"provider_specific_fields": {},
}
]
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]

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