diff --git a/.circleci/scripts/classify_changes.sh b/.circleci/scripts/classify_changes.sh index 21a85c1d914..01bc8290199 100755 --- a/.circleci/scripts/classify_changes.sh +++ b/.circleci/scripts/classify_changes.sh @@ -1,16 +1,22 @@ #!/usr/bin/env bash set -uo pipefail -category="${1:?usage: classify_changes.sh }" +category="${1:?usage: classify_changes.sh }" has_client=false has_backend=false has_ci=false has_provider_harness=false has_cost_map=false +has_mcp_dependencies=false outside_cost_map_set=false while IFS= read -r file || [ -n "$file" ]; do [ -n "$file" ] || continue + case "$file" in + *.md | *.mdx) : ;; + pyproject.toml | */pyproject.toml | uv.lock | uv.toml | .python-version | rust-toolchain.toml | litellm-rust/* | litellm/__init__.py | litellm/proxy/proxy_server.py | litellm/*mcp* | tests/*mcp* | litellm/integrations/arize/* | tests/base_sdk_tests/* | scripts/check_mcp_sdk_install.py | .github/workflows/test-mcp-dependency-resolution.yml | .github/actions/detect-changes/* | .github/actions/setup-uv-with-retries/* | .github/actions/cache-cargo-build/* | .github/scripts/detect_changes.sh | .github/scripts/uv_sync_with_retries.sh | .circleci/scripts/classify_changes.sh | tests/test_litellm/test_circleci_path_filter.py | tests/test_litellm/test_detect_changes.py) + has_mcp_dependencies=true ;; + esac case "$file" in tests/e2e/*/*.py) : ;; tests/e2e/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/test_litellm/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock) @@ -31,6 +37,9 @@ while IFS= read -r file || [ -n "$file" ]; do done case "$category" in + mcp-dependencies) + [ "$has_mcp_dependencies" = true ] && echo run || echo skip + ;; cost-map-only) { [ "$has_cost_map" = true ] && [ "$outside_cost_map_set" = false ]; } && echo run || echo skip ;; diff --git a/.github/actions/detect-changes/action.yml b/.github/actions/detect-changes/action.yml index 9b22d2c23a8..b0f9b72f0ee 100644 --- a/.github/actions/detect-changes/action.yml +++ b/.github/actions/detect-changes/action.yml @@ -14,7 +14,7 @@ description: >- inputs: category: - description: "Which classification to apply: backend, client or ui" + description: "Which classification to apply: backend, client, ui, provider-harness, cost-map-only or mcp-dependencies" required: false default: backend github-token: diff --git a/.github/scripts/auto_merge_price_sync.py b/.github/scripts/auto_merge_price_sync.py deleted file mode 100644 index 2cb1b79d867..00000000000 --- a/.github/scripts/auto_merge_price_sync.py +++ /dev/null @@ -1,393 +0,0 @@ -"""Auto-merge the provider-info-sync bot's cost-map pull requests. - -Evaluates every gate (author allowlist, cost-map-only diff, required and -non-required checks, human reviews) and merges with a merge commit when -all of them hold. Every hold reason is logged; the process exits 0 on hold -and 1 only on API or programming errors. -``DRY_RUN=1`` prints the verdict without calling the merge endpoint. -""" - -from __future__ import annotations - -import json -import os -import subprocess -import sys -import time -import urllib.error -import urllib.request -from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass -from datetime import datetime, timezone -from typing import Final - -REPO_ROOT: Final = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -CLASSIFY_SCRIPT: Final = os.path.join(REPO_ROOT, ".circleci", "scripts", "classify_changes.sh") -API_ROOT: Final = "https://api.github.com" -CHANGED_FILE_CEILING: Final = 3000 -OK_CHECK_CONCLUSIONS: Final = frozenset({"success", "skipped", "neutral"}) - - -@dataclass(frozen=True, slots=True) -class PullRequest: - number: int - title: str - author_login: str - state: str - draft: bool - mergeable: bool | None - mergeable_state: str - head_sha: str - - -@dataclass(frozen=True, slots=True) -class CheckRun: - name: str - status: str - conclusion: str | None - - -@dataclass(frozen=True, slots=True) -class CommitStatus: - context: str - state: str - - -@dataclass(frozen=True, slots=True) -class Review: - author_login: str - state: str - body: str - commit_id: str - submitted_at: datetime - - -@dataclass(frozen=True, slots=True) -class Verdict: - merge: bool - reasons: tuple[str, ...] - - -@dataclass(frozen=True, slots=True) -class EvaluationInputs: - pr: PullRequest - changed_files: tuple[str, ...] - required_contexts: frozenset[str] - check_runs: tuple[CheckRun, ...] - statuses: tuple[CommitStatus, ...] - reviews: tuple[Review, ...] - self_check_name: str - author_allowlist: frozenset[str] - - -def _is_bot_login(login: str) -> bool: - return login.lower().endswith("[bot]") - - -def _classify(changed_files: Sequence[str]) -> str: - result: Final = subprocess.run( - ["bash", CLASSIFY_SCRIPT, "cost-map-only"], - input="\n".join(changed_files), - capture_output=True, - text=True, - check=False, - ) - if result.returncode != 0: - return "error" - return result.stdout.strip() - - -def evaluate( - inputs: EvaluationInputs, - *, - classify: Callable[[Sequence[str]], str] = _classify, -) -> Verdict: - pr: Final = inputs.pr - reasons: list[str] = [] - - if pr.author_login.lower() not in {login.lower() for login in inputs.author_allowlist}: - reasons.append(f"author {pr.author_login!r} not in allowlist") - if pr.state != "open": - reasons.append("pr not open") - if pr.draft: - reasons.append("pr is a draft") - if pr.mergeable is None: - reasons.append("mergeability unknown") - elif not pr.mergeable: - reasons.append("pr not mergeable") - if pr.mergeable_state == "dirty": - reasons.append("pr has merge conflicts") - - if len(inputs.changed_files) > CHANGED_FILE_CEILING: - reasons.append(f"changed file count {len(inputs.changed_files)} over {CHANGED_FILE_CEILING} ceiling") - else: - decision: Final = classify(inputs.changed_files) - if decision != "run": - reasons.append("changed files outside the cost-map-only set") - - green_runs: Final = frozenset(run.name for run in inputs.check_runs if run.conclusion in OK_CHECK_CONCLUSIONS) - green_statuses: Final = frozenset(status.context for status in inputs.statuses if status.state == "success") - for context in sorted(inputs.required_contexts): - if context not in green_runs and context not in green_statuses: - reasons.append(f"required check {context!r} not green") - for run in inputs.check_runs: - if run.name == inputs.self_check_name: - continue - if run.status != "completed" or run.conclusion not in OK_CHECK_CONCLUSIONS: - reasons.append(f"check run {run.name!r} is {run.status}/{run.conclusion}") - for status in inputs.statuses: - if status.state != "success": - reasons.append(f"commit status {status.context!r} is {status.state}") - - latest_state_by_reviewer: Final[dict[str, str]] = {} - for review in sorted(inputs.reviews, key=lambda review: review.submitted_at): - if _is_bot_login(review.author_login): - continue - latest_state_by_reviewer[review.author_login] = review.state - for reviewer, state in latest_state_by_reviewer.items(): - if state == "CHANGES_REQUESTED": - reasons.append(f"changes requested by {reviewer}") - - return Verdict(merge=not reasons, reasons=tuple(reasons)) - - -def _request(token: str, method: str, path: str, body: Mapping[str, object] | None = None) -> object: - url: Final = path if path.startswith("http") else f"{API_ROOT}{path}" - data: Final = None if body is None else json.dumps(body).encode("utf-8") - request: Final = urllib.request.Request( - url, - data=data, - method=method, - headers={ - "Accept": "application/vnd.github+json", - "Authorization": f"Bearer {token}", - "X-GitHub-Api-Version": "2022-11-28", - }, - ) - with urllib.request.urlopen(request) as response: - return json.loads(response.read().decode("utf-8")) - - -def _request_allow_fail( - token: str, method: str, path: str, body: Mapping[str, object] | None = None -) -> tuple[int, object | None]: - url: Final = path if path.startswith("http") else f"{API_ROOT}{path}" - data: Final = None if body is None else json.dumps(body).encode("utf-8") - request: Final = urllib.request.Request( - url, - data=data, - method=method, - headers={ - "Accept": "application/vnd.github+json", - "Authorization": f"Bearer {token}", - "X-GitHub-Api-Version": "2022-11-28", - }, - ) - try: - with urllib.request.urlopen(request) as response: - return response.status, json.loads(response.read().decode("utf-8")) - except urllib.error.HTTPError as exc: - return exc.code, None - - -def _items(payload: object, key: str | None = None) -> tuple[object, ...]: - source: Final = payload.get(key) if key and isinstance(payload, Mapping) else payload - if not isinstance(source, list): - return () - return tuple(source) - - -def _paginate(token: str, path: str, key: str | None = None) -> list[object]: - separator: Final = "&" if "?" in path else "?" - results: list[object] = [] - for page in range(1, 10_000): - batch: Final = _items(_request(token, "GET", f"{path}{separator}per_page=100&page={page}"), key) - results.extend(batch) - if len(batch) < 100: - return results - return results - - -def _text(value: object) -> str: - return value if isinstance(value, str) else "" - - -def _int(value: object) -> int: - return value if isinstance(value, int) else 0 - - -def _bool(value: object) -> bool: - return value is True - - -def _nested(value: object, *keys: str) -> object: - current: object = value - for key in keys: - if not isinstance(current, Mapping): - return None - current = current.get(key) - return current - - -def _parse_time(value: object) -> datetime: - text: Final = _text(value) - if not text: - return datetime.min.replace(tzinfo=timezone.utc) - return datetime.fromisoformat(text.replace("Z", "+00:00")) - - -def _load_pr(token: str, repo: str, number: int) -> PullRequest: - data: Final = _request(token, "GET", f"/repos/{repo}/pulls/{number}") - if not isinstance(data, Mapping): - raise RuntimeError(f"unexpected pull payload for #{number}") - return PullRequest( - number=number, - title=_text(data.get("title")), - author_login=_text(_nested(data, "user", "login")), - state=_text(data.get("state")), - draft=_bool(data.get("draft")), - mergeable=data.get("mergeable") if isinstance(data.get("mergeable"), bool) else None, - mergeable_state=_text(data.get("mergeable_state")), - head_sha=_text(_nested(data, "head", "sha")), - ) - - -def _list_candidate_prs(token: str, repo: str, base: str, allowlist: frozenset[str]) -> list[int]: - candidates: Final = _paginate(token, f"/repos/{repo}/pulls?state=open&base={base}") - return [ - _int(item.get("number")) - for item in candidates - if isinstance(item, Mapping) and _text(_nested(item, "user", "login")).lower() in allowlist - ] - - -def _changed_files(token: str, repo: str, number: int) -> tuple[str, ...]: - files: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/files") - return tuple(_text(item.get("filename")) for item in files if isinstance(item, Mapping)) - - -def _required_contexts(token: str, repo: str, base: str) -> frozenset[str]: - payload: Final = _request(token, "GET", f"/repos/{repo}/rules/branches/{base}") - contexts: set[str] = set() - for rule in _items(payload): - if not isinstance(rule, Mapping) or rule.get("type") != "required_status_checks": - continue - checks: Final = _nested(rule, "parameters", "required_status_checks") - for check in _items(checks): - if isinstance(check, Mapping): - context: Final = _text(check.get("context")) - if context: - contexts.add(context) - return frozenset(contexts) - - -def _check_runs(token: str, repo: str, sha: str) -> tuple[CheckRun, ...]: - runs: Final = _paginate(token, f"/repos/{repo}/commits/{sha}/check-runs", key="check_runs") - return tuple( - CheckRun( - name=_text(item.get("name")), - status=_text(item.get("status")), - conclusion=item.get("conclusion") if isinstance(item.get("conclusion"), str) else None, - ) - for item in runs - if isinstance(item, Mapping) - ) - - -def _statuses(token: str, repo: str, sha: str) -> tuple[CommitStatus, ...]: - payload: Final = _request(token, "GET", f"/repos/{repo}/commits/{sha}/status") - return tuple( - CommitStatus(context=_text(item.get("context")), state=_text(item.get("state"))) - for item in _items(payload, "statuses") - if isinstance(item, Mapping) - ) - - -def _reviews(token: str, repo: str, number: int) -> tuple[Review, ...]: - reviews: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/reviews") - return tuple( - Review( - author_login=_text(_nested(item, "user", "login")), - state=_text(item.get("state")), - body=_text(item.get("body")), - commit_id=_text(item.get("commit_id")), - submitted_at=_parse_time(item.get("submitted_at")), - ) - for item in reviews - if isinstance(item, Mapping) - ) - - -def _mergeable_or_refetch(token: str, repo: str, pr: PullRequest) -> PullRequest: - if pr.mergeable is not None: - return pr - time.sleep(5) - return _load_pr(token, repo, pr.number) - - -def _gather_inputs( - token: str, - repo: str, - number: int, - base: str, - self_check_name: str, - allowlist: frozenset[str], -) -> EvaluationInputs: - pr: Final = _mergeable_or_refetch(token, repo, _load_pr(token, repo, number)) - return EvaluationInputs( - pr=pr, - changed_files=_changed_files(token, repo, number), - required_contexts=_required_contexts(token, repo, base), - check_runs=_check_runs(token, repo, pr.head_sha), - statuses=_statuses(token, repo, pr.head_sha), - reviews=_reviews(token, repo, number), - self_check_name=self_check_name, - author_allowlist=allowlist, - ) - - -def merge_request_body(pr: PullRequest) -> dict[str, str]: - return {"merge_method": "merge", "commit_title": f"{pr.title} (#{pr.number})", "sha": pr.head_sha} - - -def _merge(token: str, repo: str, pr: PullRequest) -> None: - status, _ = _request_allow_fail(token, "PUT", f"/repos/{repo}/pulls/{pr.number}/merge", merge_request_body(pr)) - if status in (200, 405, 409): - print(f"auto-merge-price-sync: PR #{pr.number} merge call returned {status}") - return - raise RuntimeError(f"merge call for PR #{pr.number} returned {status}") - - -def main() -> int: - token: Final = os.environ.get("GH_TOKEN", "") - repo: Final = os.environ.get("REPO", "") - base: Final = os.environ.get("BASE_BRANCH", "main") - dry_run: Final = os.environ.get("DRY_RUN", "") != "" - self_check_name: Final = os.environ.get("SELF_CHECK_NAME", "auto-merge-price-sync") - allowlist: Final = frozenset(login.lower() for login in os.environ.get("PR_AUTHOR_ALLOWLIST", "").split() if login) - if not token: - print("auto-merge-price-sync: app credentials not configured") - return 0 - if not repo: - print("auto-merge-price-sync: REPO not set", file=sys.stderr) - return 1 - - pr_number_env: Final = os.environ.get("PR_NUMBER", "") - candidates: Final = [int(pr_number_env)] if pr_number_env else _list_candidate_prs(token, repo, base, allowlist) - for number in candidates: - inputs: Final = _gather_inputs(token, repo, number, base, self_check_name, allowlist) - verdict: Final = evaluate(inputs) - for reason in verdict.reasons: - print(f"auto-merge-price-sync: PR #{number} hold: {reason}") - if not verdict.merge: - continue - print(f"auto-merge-price-sync: PR #{number} all gates green") - if dry_run: - print(f"auto-merge-price-sync: DRY_RUN merge suppressed for PR #{number}") - continue - _merge(token, repo, inputs.pr) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index 617b09a8075..d4e9a65e7c0 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -63,6 +63,11 @@ on: description: "Unique name for the coverage artifact (must be unique per run)" required: true type: string + legacy-mcp-peer: + description: "Install the isolated SDK1 peer for MCP compatibility tests" + required: false + type: boolean + default: false permissions: contents: read @@ -125,10 +130,17 @@ jobs: - name: Install dependencies if: steps.changes.outputs.decision != 'skip' timeout-minutes: 8 + env: + LEGACY_MCP_PEER: ${{ inputs.legacy-mcp-peer }} run: | diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]' + if [ "$LEGACY_MCP_PEER" = "true" ]; then + uv venv --python "${UV_PYTHON}" .venv-mcp-peer + uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1' + echo "MCP_TEST_PEER_PYTHON=$GITHUB_WORKSPACE/.venv-mcp-peer/bin/python" >> "$GITHUB_ENV" + fi - name: Cache Prisma binaries if: steps.changes.outputs.decision != 'skip' diff --git a/.github/workflows/auto-merge-price-sync.yml b/.github/workflows/auto-merge-price-sync.yml deleted file mode 100644 index e14fc3f955b..00000000000 --- a/.github/workflows/auto-merge-price-sync.yml +++ /dev/null @@ -1,61 +0,0 @@ -name: auto-merge-price-sync - -on: - issue_comment: - types: [created, edited] - check_suite: - types: [completed] - status: {} - schedule: - - cron: "*/30 * * * *" - workflow_dispatch: - inputs: - pr-number: - description: "Evaluate only this PR number (empty = scan all open sync-bot PRs)" - required: false - default: "" - -permissions: - contents: read - pull-requests: read - checks: read - statuses: read - -concurrency: - group: auto-merge-price-sync - cancel-in-progress: false - -jobs: - auto-merge-price-sync: - runs-on: ubuntu-latest - timeout-minutes: 15 - env: - PROVIDER_INFO_SYNC_APP_ID: ${{ secrets.PROVIDER_INFO_SYNC_APP_ID }} - PROVIDER_INFO_SYNC_APP_PRIVATE_KEY: ${{ secrets.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY }} - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Mint app token - id: app-token - if: ${{ env.PROVIDER_INFO_SYNC_APP_ID != '' && env.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY != '' }} - uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0 - with: - app-id: ${{ secrets.PROVIDER_INFO_SYNC_APP_ID }} - private-key: ${{ secrets.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY }} - - - name: Auto-merge eligible sync PRs - env: - GH_TOKEN: ${{ steps.app-token.outputs.token }} - REPO: ${{ github.repository }} - PR_NUMBER: ${{ (github.event.issue.pull_request && github.event.issue.number) || github.event.inputs.pr-number || '' }} - BASE_BRANCH: main - PR_AUTHOR_ALLOWLIST: "berriai-litellm-provider-info-sync[bot]" - SELF_CHECK_NAME: auto-merge-price-sync - run: python3 .github/scripts/auto_merge_price_sync.py diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index 7e013b7bb0b..fd7513a3937 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -69,7 +69,7 @@ jobs: uv run --frozen --no-default-groups --with pytest==8.3.5 --with pytest-codspeed==4.3.0 - --with "mcp>=1.26.0,<2.0" + --with "mcp>=2.2.0,<3.0" --with "a2a-sdk>=1.1.0,<2.0" pytest -p pytest_codspeed.plugin @@ -86,7 +86,7 @@ jobs: uv run --frozen --no-default-groups --with pytest==8.3.5 --with pytest-codspeed==4.3.0 - --with "mcp>=1.26.0,<2.0" + --with "mcp>=2.2.0,<3.0" --with "a2a-sdk>=1.1.0,<2.0" pytest -p pytest_codspeed.plugin diff --git a/.github/workflows/test-mcp-dependency-resolution.yml b/.github/workflows/test-mcp-dependency-resolution.yml new file mode 100644 index 00000000000..b5dc573c1c1 --- /dev/null +++ b/.github/workflows/test-mcp-dependency-resolution.yml @@ -0,0 +1,98 @@ +name: LiteLLM MCP Dependency Resolution + +on: + pull_request: + branches: + - main + - litellm_internal_staging + - litellm_oss_staging + - "litellm_**" + +permissions: + contents: read + pull-requests: read + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +jobs: + resolve: + runs-on: ubuntu-latest + timeout-minutes: 15 + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] + + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Detect relevant changes + id: changes + uses: ./.github/actions/detect-changes + with: + category: mcp-dependencies + + - name: Set up Python + if: steps.changes.outputs.decision != 'skip' + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: ${{ matrix.python-version }} + + - name: Set up uv + if: steps.changes.outputs.decision != 'skip' + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Cache the Rust build + if: steps.changes.outputs.decision != 'skip' + uses: ./.github/actions/cache-cargo-build + + - name: Verify lockfile + if: steps.changes.outputs.decision != 'skip' + run: | + uv lock --check + + - name: Check locked runtime installations + if: steps.changes.outputs.decision != 'skip' + run: | + for extra in core mcp proxy; do + args=() + if [ "$extra" != core ]; then args=(--extra "$extra"); fi + UV_PROJECT_ENVIRONMENT=".venv-$extra" .github/scripts/uv_sync_with_retries.sh --frozen --no-dev --no-editable --python ${{ matrix.python-version }} "${args[@]}" + uv pip check --python ".venv-$extra" + if [ "$extra" = core ]; then + checker=("$GITHUB_WORKSPACE/tests/base_sdk_tests/check_base_sdk_install.py") + else + checker=("$GITHUB_WORKSPACE/scripts/check_mcp_sdk_install.py" --extra "$extra") + fi + (cd "$RUNNER_TEMP" && "$GITHUB_WORKSPACE/.venv-$extra/bin/python" "${checker[@]}") + done + + - name: Build the public wheel + if: steps.changes.outputs.decision != 'skip' + run: uv build --all-packages --wheel --out-dir dist/mcp-check + + - name: Check lowest direct runtime installations + if: steps.changes.outputs.decision != 'skip' + run: | + wheel=$(realpath dist/mcp-check/litellm-[0-9]*.whl) + for extra in core mcp proxy; do + args=() + if [ "$extra" != core ]; then args=(--extra "$extra"); fi + uv pip compile pyproject.toml --no-sources --find-links dist/mcp-check "${args[@]}" --python-version ${{ matrix.python-version }} --resolution lowest-direct -o "lowest-$extra.txt" + uv venv --python ${{ matrix.python-version }} ".venv-lowest-$extra" + uv pip sync --find-links dist/mcp-check --python ".venv-lowest-$extra" "lowest-$extra.txt" + uv pip install --python ".venv-lowest-$extra" --no-deps "$wheel" + uv pip check --python ".venv-lowest-$extra" + if [ "$extra" = core ]; then + checker=("$GITHUB_WORKSPACE/tests/base_sdk_tests/check_base_sdk_install.py") + else + checker=("$GITHUB_WORKSPACE/scripts/check_mcp_sdk_install.py" --extra "$extra") + fi + (cd "$RUNNER_TEMP" && "$GITHUB_WORKSPACE/.venv-lowest-$extra/bin/python" "${checker[@]}") + done diff --git a/.github/workflows/test-mcp.yml b/.github/workflows/test-mcp.yml deleted file mode 100644 index 93ffcbe0586..00000000000 --- a/.github/workflows/test-mcp.yml +++ /dev/null @@ -1,63 +0,0 @@ -name: LiteLLM MCP Tests (folder - tests/mcp_tests) - -on: - pull_request: - branches: - - main - - litellm_internal_staging - - litellm_oss_staging - - "litellm_**" - -permissions: - contents: read - pull-requests: read - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} - cancel-in-progress: ${{ github.event_name == 'pull_request' }} - -jobs: - test: - runs-on: ubuntu-latest - timeout-minutes: 25 - - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Detect relevant changes - id: changes - uses: ./.github/actions/detect-changes - - - name: Thank You Message - run: | - echo "### 🙏 Thank you for contributing to LiteLLM!" >> $GITHUB_STEP_SUMMARY - echo "Your PR is being tested now. We appreciate your help in making LiteLLM better!" >> $GITHUB_STEP_SUMMARY - - - name: Set up Python - if: steps.changes.outputs.decision != 'skip' - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Set up uv - if: steps.changes.outputs.decision != 'skip' - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - - name: Cache the Rust build - if: steps.changes.outputs.decision != 'skip' - uses: ./.github/actions/cache-cargo-build - - - name: Install dependencies - if: steps.changes.outputs.decision != 'skip' - run: | - uv lock --check - .github/scripts/uv_sync_with_retries.sh --frozen --group proxy-dev --extra proxy --extra semantic-router - - - name: Run MCP tests - if: steps.changes.outputs.decision != 'skip' - run: | - uv run --no-sync pytest tests/mcp_tests -x -vv -n 4 --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml --durations=5 diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index a32b5ebb2a8..aa82a0bf3ee 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -49,6 +49,14 @@ jobs: fail-fast: false matrix: include: + - shard: mcp-integration + artifact-name: mcp-integration + test-path: "tests/mcp_tests" + workers: 2 + reruns: 0 + timeout-minutes: 20 + job-timeout-minutes: 60 + - shard: core-utils artifact-name: core-utils test-path: "tests/test_litellm/litellm_core_utils" @@ -254,3 +262,4 @@ jobs: timeout-minutes: ${{ matrix.timeout-minutes }} job-timeout-minutes: ${{ matrix.job-timeout-minutes }} artifact-name: ${{ matrix.artifact-name }} + legacy-mcp-peer: ${{ matrix.shard == 'mcp-integration' }} diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py index d10b5a2ab09..3422e8969b0 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py @@ -8,7 +8,7 @@ ## This provides an LLM Guard Integration for content moderation on the proxy import asyncio -from typing import Optional +from typing import Final, Optional import aiohttp from fastapi import HTTPException @@ -137,15 +137,20 @@ class _ENTERPRISE_LLMGuard(CustomLogger): return self.print_verbose("Makes LLM Guard Check") - if call_type not in [ + accepted_call_types: Final = ( "completion", + "acompletion", + "text_completion", + "atext_completion", "embeddings", + "embedding", + "aembedding", "image_generation", - "moderation", - "audio_transcription", - ]: + "aimage_generation", + ) + if call_type not in accepted_call_types: self.print_verbose( - f"Call Type - {call_type}, not in accepted list - ['completion','embeddings','image_generation','moderation','audio_transcription']" + f"Call Type - {call_type}, not in accepted list - {accepted_call_types}" ) return data @@ -163,16 +168,14 @@ class _ENTERPRISE_LLMGuard(CustomLogger): *(self._moderate_message(message) for message in messages) ) ) - return data input_ = data.get("input") if input_ is not None: - data["input"] = await self._moderate_input(input_) - return data + data["input"] = await self._moderate_text_or_list(input_) prompt = data.get("prompt") - if isinstance(prompt, str): - data["prompt"] = await self.moderation_check(text=prompt) + if prompt is not None: + data["prompt"] = await self._moderate_text_or_list(prompt) return data async def _moderate_message(self, message: dict) -> dict: @@ -195,17 +198,17 @@ class _ENTERPRISE_LLMGuard(CustomLogger): return {**part, "text": await self.moderation_check(text=part["text"])} return part - async def _moderate_input(self, input_: object) -> object: - if isinstance(input_, str): - return await self.moderation_check(text=input_) - if isinstance(input_, list): + async def _moderate_text_or_list(self, value: object) -> object: + if isinstance(value, str): + return await self.moderation_check(text=value) + if isinstance(value, list): return [ await self.moderation_check(text=item) if isinstance(item, str) else item - for item in input_ + for item in value ] - return input_ + return value async def async_post_call_streaming_hook( self, user_api_key_dict: UserAPIKeyAuth, response: str diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 726d2f484da..860f01c4ad1 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1960,6 +1960,7 @@ dependencies = [ "aws-smithy-runtime-api", "aws-types", "litellm-auth", + "litellm-http", "moka", "reqwest 0.12.28", "serde_json", @@ -2142,6 +2143,7 @@ dependencies = [ "reqwest 0.12.28", "rstest", "rustls 0.23.42", + "serde", "serde_json", "thiserror 2.0.19", "tokio", @@ -2174,6 +2176,7 @@ dependencies = [ "serde_json", "serde_path_to_error", "serde_with", + "strum", "thiserror 2.0.19", "time", "tokio", diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index d998b647960..1f27c7bc990 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] litellm-auth.workspace = true +litellm-http.workspace = true moka = { workspace = true, features = ["sync"] } serde_json.workspace = true diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index 3b6b73bc6a9..bbcb0f016c8 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -20,8 +20,7 @@ use super::constants::{ AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION, AWS_REGION_NAME, AWS_ROLE_ARN, AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN, AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE, - BEDROCK_SERVICE, DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX, - SIGV4_COMPUTED_HEADER_NAMES, + DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX, SIGV4_COMPUTED_HEADER_NAMES, }; const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60); @@ -451,11 +450,12 @@ pub fn is_sigv4_computed_header(name: &str) -> bool { SIGV4_COMPUTED_HEADER_NAMES.contains(&name.to_ascii_lowercase().as_str()) } -pub fn sign_bedrock_post( +pub fn sign_post( url: &str, body: &[u8], headers: &BTreeMap, region: &str, + service: &str, credentials: &Credentials, signing_time: SystemTime, ) -> Result, Error> { @@ -463,7 +463,7 @@ pub fn sign_bedrock_post( let params = v4::SigningParams::builder() .identity(&identity) .region(region) - .name(BEDROCK_SERVICE) + .name(service) .time(signing_time) .settings(SigningSettings::default()) .build() @@ -534,22 +534,28 @@ fn is_bedrock_region(value: &str) -> bool { .all(|char| char.is_ascii_alphanumeric() || char == '-') } +/// The region a caller configured: `aws_region_name`, then the model's own +/// region, then the environment. Each service decides what a missing one means. +pub fn resolve_aws_region( + model_region: Option<&str>, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option { + optional_params + .get("aws_region_name") + .and_then(Value::as_str) + .or(model_region) + .map(str::to_string) + .or_else(|| env_lookup(AWS_REGION_NAME)) + .or_else(|| env_lookup(AWS_REGION)) +} + pub fn resolve_bedrock_region( model_region: Option<&str>, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, ) -> String { - if let Some(region) = optional_params - .get("aws_region_name") - .and_then(Value::as_str) - { - return region.to_string(); - } - if let Some(region) = model_region { - return region.to_string(); - } - env_lookup(AWS_REGION_NAME) - .or_else(|| env_lookup(AWS_REGION)) + resolve_aws_region(model_region, optional_params, env_lookup) .unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string()) } @@ -609,11 +615,36 @@ pub fn host_supplied_credentials(optional_params: &Map) -> Option #[cfg(test)] mod tests { use super::*; + use crate::constants::BEDROCK_SERVICE; fn no_env(_: &str) -> Option { None } + #[test] + fn a_region_comes_from_the_call_then_the_model_then_the_environment() { + let params = Map::from_iter([("aws_region_name".to_string(), Value::from("eu-west-1"))]); + let region_name = |key: &str| (key == AWS_REGION_NAME).then(|| "ap-south-1".to_string()); + let region = |key: &str| (key == AWS_REGION).then(|| "sa-east-1".to_string()); + + let resolved = [ + resolve_aws_region(Some("us-east-2"), ¶ms, ®ion_name), + resolve_aws_region(Some("us-east-2"), &Map::new(), ®ion_name), + resolve_aws_region(None, &Map::new(), ®ion_name), + resolve_aws_region(None, &Map::new(), ®ion), + resolve_aws_region(None, &Map::new(), &no_env), + ]; + + assert_eq!( + resolved.map(|region| region.unwrap_or_else(|| "none".into())), + ["eu-west-1", "us-east-2", "ap-south-1", "sa-east-1", "none"] + ); + assert_eq!( + resolve_bedrock_region(None, &Map::new(), &no_env), + DEFAULT_BEDROCK_REGION + ); + } + fn parity_inputs() -> (String, Vec, BTreeMap) { ( "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke" @@ -811,11 +842,12 @@ mod tests { None, "test", ); - let signed = sign_bedrock_post( + let signed = sign_post( &url, &body, &signable, "us-east-1", + BEDROCK_SERVICE, &credentials, SystemTime::UNIX_EPOCH, ) @@ -843,11 +875,12 @@ mod tests { None, "test", ); - let signed = sign_bedrock_post( + let signed = sign_post( &url, &body, &headers, "us-east-1", + BEDROCK_SERVICE, &credentials, UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), ) @@ -878,11 +911,12 @@ mod tests { None, "test", ); - let signed = sign_bedrock_post( + let signed = sign_post( &url, &body, &headers, "us-east-1", + BEDROCK_SERVICE, &credentials, UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), ) @@ -915,11 +949,12 @@ mod tests { let url = format!( "https://bedrock-runtime.{region}.amazonaws.com/model/us.anthropic.claude-opus-4-8/invoke" ); - let signed_headers = sign_bedrock_post( + let signed_headers = sign_post( &url, &body, &headers, region, + BEDROCK_SERVICE, &credentials, SystemTime::now(), )?; diff --git a/litellm-rust/crates/auth-aws/src/lib.rs b/litellm-rust/crates/auth-aws/src/lib.rs index 264592ccb2e..0fe0b390110 100644 --- a/litellm-rust/crates/auth-aws/src/lib.rs +++ b/litellm-rust/crates/auth-aws/src/lib.rs @@ -1,6 +1,9 @@ mod aws; pub mod constants; mod error; +mod signer; pub use aws::*; +pub use aws_credential_types::Credentials; pub use error::Error; +pub use signer::SigV4Signer; diff --git a/litellm-rust/crates/auth-aws/src/signer.rs b/litellm-rust/crates/auth-aws/src/signer.rs new file mode 100644 index 00000000000..49a3910c1d5 --- /dev/null +++ b/litellm-rust/crates/auth-aws/src/signer.rs @@ -0,0 +1,178 @@ +use std::{collections::BTreeMap, time::SystemTime}; + +use aws_credential_types::Credentials; +use litellm_http::outbound::{RequestSigner, UnsignedRequest}; +use serde_json::{Map, Value}; + +use crate::{ + Error, aws_auth_config, aws_signature_headers, host_supplied_credentials, + is_sigv4_computed_header, resolve_credentials, sign_post, +}; + +#[derive(Clone, Debug)] +pub struct SigV4Signer { + region: String, + service: &'static str, + credentials: Credentials, + clock: fn() -> SystemTime, +} + +impl SigV4Signer { + pub fn new(region: String, service: &'static str, credentials: Credentials) -> Self { + Self { + region, + service, + credentials, + clock: SystemTime::now, + } + } + + pub fn with_clock(self, clock: fn() -> SystemTime) -> Self { + Self { clock, ..self } + } + + pub async fn resolve( + region: String, + service: &'static str, + optional_params: &Map, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result { + let credentials = match host_supplied_credentials(optional_params) { + Some(credentials) => credentials, + None => { + resolve_credentials(aws_auth_config(optional_params, env_lookup), env_lookup) + .await? + } + }; + Ok(Self::new(region, service, credentials)) + } +} + +impl RequestSigner for SigV4Signer { + fn sign( + &self, + request: UnsignedRequest<'_>, + ) -> Result, litellm_http::Error> { + if let Some((name, _)) = request + .headers + .iter() + .find(|(name, _)| is_sigv4_computed_header(name)) + { + return Err(litellm_http::Error::ComputedHeader(name.clone())); + } + let headers: BTreeMap = request.headers.iter().cloned().collect(); + sign_post( + request.url, + request.body, + &aws_signature_headers(&headers), + &self.region, + self.service, + &self.credentials, + (self.clock)(), + ) + .map(|signature| signature.into_iter().collect()) + .map_err(|error| litellm_http::Error::Signature(error.to_string())) + } +} + +#[cfg(test)] +mod tests { + use std::time::{Duration, UNIX_EPOCH}; + + use litellm_http::outbound::OutboundRequest; + use serde_json::json; + + use super::*; + + fn fixed_clock() -> SystemTime { + UNIX_EPOCH + Duration::from_secs(1_700_000_000) + } + + fn signer(service: &'static str) -> SigV4Signer { + SigV4Signer::new( + "us-east-1".into(), + service, + Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"), + ) + .with_clock(fixed_clock) + } + + fn authorization(body: &Value, service: &'static str) -> String { + OutboundRequest::signed_json( + "https://textract.us-east-1.amazonaws.com/".into(), + vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())], + body, + None, + &signer(service), + ) + .unwrap() + .header("Authorization") + .unwrap() + .to_string() + } + + #[test] + fn the_signature_verifies_against_the_bytes_that_are_sent() { + let sent = OutboundRequest::signed_json( + "https://textract.us-east-1.amazonaws.com/".into(), + vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())], + &json!({"Document": {"Bytes": "aGk="}}), + None, + &signer("textract"), + ) + .unwrap(); + let unsigned: BTreeMap = sent + .headers() + .iter() + .filter(|(name, _)| !is_sigv4_computed_header(name)) + .cloned() + .collect(); + let recomputed = sign_post( + sent.url(), + sent.body(), + &aws_signature_headers(&unsigned), + "us-east-1", + "textract", + &Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"), + fixed_clock(), + ) + .unwrap(); + + assert_eq!( + sent.header("Authorization"), + Some(recomputed["Authorization"].as_str()) + ); + } + + #[test] + fn the_signature_depends_on_the_body_and_the_service() { + let original = authorization(&json!({"text": "card 4111"}), "textract"); + + assert_ne!( + original, + authorization(&json!({"text": "card [REDACTED]"}), "textract") + ); + assert_ne!( + original, + authorization(&json!({"text": "card 4111"}), "bedrock") + ); + assert!(original.contains("/us-east-1/textract/aws4_request")); + } + + #[test] + fn a_forwarded_computed_header_is_refused_instead_of_sent_twice() { + let error = OutboundRequest::signed_json( + "https://textract.us-east-1.amazonaws.com/".into(), + vec![("authorization".into(), "Bearer caller".into())], + &json!({}), + None, + &signer("textract"), + ) + .unwrap_err(); + + assert_eq!( + error, + litellm_http::Error::ComputedHeader("authorization".into()) + ); + } +} diff --git a/litellm-rust/crates/auth/src/http.rs b/litellm-rust/crates/auth/src/http.rs index 7d20991d838..dd87d00e70f 100644 --- a/litellm-rust/crates/auth/src/http.rs +++ b/litellm-rust/crates/auth/src/http.rs @@ -40,13 +40,22 @@ pub fn apply_credential( ) } -/// How the upstream call is authenticated. API-key strategies are resolved in -/// `prepare`; SigV4 needs the serialized body, so the handler signs it. +/// How the upstream call is authenticated. API-key strategies become headers +/// in `prepare`; SigV4 covers the serialized body, so it is applied where the +/// outbound request is built. #[derive(Clone, Debug, PartialEq, Eq)] pub enum RequestAuth { - Header { name: &'static str, value: String }, - Bearer { token: String }, - AwsSigV4 { region: String }, + Header { + name: &'static str, + value: String, + }, + Bearer { + token: String, + }, + AwsSigV4 { + region: String, + service: &'static str, + }, } #[cfg(test)] diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy/src/adapter.rs index 6c013cd1ea5..883a35f0df5 100644 --- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy/src/adapter.rs @@ -12,7 +12,7 @@ use pyo3::{ exceptions::{PyBaseException, PyException}, gc::{PyTraverseError, PyVisit}, prelude::*, - types::{PyDict, PyList}, + types::{PyDateTime, PyDict, PyList}, }; use serde_json::Value; @@ -73,10 +73,7 @@ pub struct LegacyLogging { } fn datetime(py: Python<'_>, epoch_seconds: f64) -> PyResult> { - py.import("datetime")? - .getattr("datetime")? - .call_method1("fromtimestamp", (epoch_seconds,)) - .map(Bound::unbind) + PyDateTime::from_timestamp(py, epoch_seconds, None).map(|value| value.into_any().unbind()) } fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool { diff --git a/litellm-rust/crates/core/src/audio_transcription/error.rs b/litellm-rust/crates/core/src/audio_transcription/error.rs index 122cbab358f..81b57af2c6c 100644 --- a/litellm-rust/crates/core/src/audio_transcription/error.rs +++ b/litellm-rust/crates/core/src/audio_transcription/error.rs @@ -24,6 +24,8 @@ pub enum Error { #[error(transparent)] Headers(#[from] litellm_http::request::HeaderError), #[error(transparent)] + Http(#[from] litellm_http::Error), + #[error(transparent)] Aws(#[from] litellm_auth_aws::Error), } diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs index 503cc922966..a1862f341a5 100644 --- a/litellm-rust/crates/core/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -1,4 +1,4 @@ -use litellm_http::request::{http_request, truncate_error_body}; +use litellm_http::request::truncate_error_body; use serde_json::Value; use super::{Error, client::http_client}; @@ -7,17 +7,18 @@ use crate::audio_transcription::types::ProviderAudioTranscriptionRequest; pub async fn execute_audio_transcription_provider_call( request: ProviderAudioTranscriptionRequest, ) -> Result { - let body = serde_json::to_vec(&request.body) - .map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?; - let headers = signed_headers(&request, &body).await?; - let mut request_builder = http_client().post(&request.url).body(body); - for (key, value) in headers { - request_builder = request_builder.header(key, value); - } - if let Some(duration) = request.timeout { - request_builder = request_builder.timeout(duration); - } - let response = http_request(request_builder).await.map_err(|error| { + let response = crate::outbound::outbound_request::( + &request.auth, + request.url.clone(), + request.upstream_headers.clone(), + &request.body, + request.timeout, + &request.optional_params, + ) + .await? + .send(http_client()) + .await + .map_err(|error| { Error::Transport(litellm_http::transport::Error::Network(error.to_string())) })?; let status = response.status(); @@ -37,33 +38,3 @@ pub async fn execute_audio_transcription_provider_call( .transform_audio_transcription_response(&request.model, response_json)? .into_json()) } - -async fn signed_headers( - request: &ProviderAudioTranscriptionRequest, - body: &[u8], -) -> Result, Error> { - use std::{collections::BTreeMap, time::SystemTime}; - - use litellm_auth_aws::{aws_auth_config, resolve_credentials, sign_bedrock_post}; - use litellm_llms::base_llm::audio_transcription::transformation::AudioTranscriptionAuth; - - let AudioTranscriptionAuth::AwsSigV4 { region, .. } = &request.auth else { - return Ok(request.upstream_headers.clone()); - }; - let env_lookup = |key: &str| std::env::var(key).ok(); - let credentials = resolve_credentials( - aws_auth_config(&request.optional_params, &env_lookup), - &env_lookup, - ) - .await?; - let unsigned: BTreeMap = request.upstream_headers.iter().cloned().collect(); - let signature = sign_bedrock_post( - &request.url, - body, - &unsigned, - region, - &credentials, - SystemTime::now(), - )?; - Ok(unsigned.into_iter().chain(signature).collect()) -} diff --git a/litellm-rust/crates/core/src/audio_transcription/prepare.rs b/litellm-rust/crates/core/src/audio_transcription/prepare.rs index 829617d26bd..807993c38b7 100644 --- a/litellm-rust/crates/core/src/audio_transcription/prepare.rs +++ b/litellm-rust/crates/core/src/audio_transcription/prepare.rs @@ -1,9 +1,7 @@ use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider}; use litellm_http::request::{has_header, string_headers}; use litellm_llms::{ - base_llm::audio_transcription::transformation::{ - AudioTranscriptionAuth, BaseAudioTranscriptionConfig, - }, + base_llm::audio_transcription::transformation::{BaseAudioTranscriptionConfig, RequestAuth}, bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG, }; @@ -43,11 +41,14 @@ pub fn prepare_audio_transcription_provider_call( let env_lookup = |key: &str| std::env::var(key).ok(); let mut headers = string_headers("audio transcription", request.extra_headers)?; let auth = config.auth_strategy(&model, &request.optional_params, &env_lookup)?; - if matches!(auth, AudioTranscriptionAuth::Bearer) - && !has_header(&headers, "authorization") - && let Some(api_key) = request.api_key - { - headers.push(("Authorization".to_string(), format!("Bearer {api_key}"))); + match &auth { + RequestAuth::Bearer { token } if !has_header(&headers, "authorization") => { + headers.push(("Authorization".to_string(), format!("Bearer {token}"))); + } + RequestAuth::Header { name, value } if !has_header(&headers, name) => { + headers.push(((*name).to_string(), value.clone())); + } + RequestAuth::Bearer { .. } | RequestAuth::Header { .. } | RequestAuth::AwsSigV4 { .. } => {} } if !has_header(&headers, "content-type") { headers.push(("Content-Type".to_string(), "application/json".to_string())); diff --git a/litellm-rust/crates/core/src/audio_transcription/types.rs b/litellm-rust/crates/core/src/audio_transcription/types.rs index ca09dd945be..0d87483c9bf 100644 --- a/litellm-rust/crates/core/src/audio_transcription/types.rs +++ b/litellm-rust/crates/core/src/audio_transcription/types.rs @@ -1,7 +1,7 @@ use std::time::Duration; use litellm_llms::base_llm::audio_transcription::transformation::{ - AudioTranscriptionAuth, BaseAudioTranscriptionConfig, + BaseAudioTranscriptionConfig, RequestAuth, }; use serde_json::{Map, Value}; @@ -24,7 +24,7 @@ pub struct ProviderAudioTranscriptionRequest { pub url: String, pub body: Value, pub upstream_headers: Vec<(String, String)>, - pub auth: AudioTranscriptionAuth, + pub auth: RequestAuth, pub optional_params: Map, pub timeout: Option, } diff --git a/litellm-rust/crates/core/src/chat_completions/error.rs b/litellm-rust/crates/core/src/chat_completions/error.rs index 122cbab358f..81b57af2c6c 100644 --- a/litellm-rust/crates/core/src/chat_completions/error.rs +++ b/litellm-rust/crates/core/src/chat_completions/error.rs @@ -24,6 +24,8 @@ pub enum Error { #[error(transparent)] Headers(#[from] litellm_http::request::HeaderError), #[error(transparent)] + Http(#[from] litellm_http::Error), + #[error(transparent)] Aws(#[from] litellm_auth_aws::Error), } diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index b73d4838760..de926c715d5 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -1,5 +1,5 @@ -use litellm_http::request::{http_request, truncate_error_body}; -use litellm_llms::base_llm::chat::transformation::{ChatCompletionsAuth, ProviderChatResponseData}; +use litellm_http::{outbound::OutboundRequest, request::truncate_error_body}; +use litellm_llms::base_llm::chat::transformation::ProviderChatResponseData; use litellm_types::utils::ChatCompletionsResponse; use serde_json::Value; @@ -12,22 +12,9 @@ pub(super) async fn execute_chat_completions_provider_call( request: ResolvedChatCompletionsRequest<'_>, ) -> Result { let request = prepare_provider_request(request)?; - let body = serde_json::to_vec(&request.body).map_err(|err| { - Error::InvalidRequest(format!( - "failed to serialize chat completions request: {err}" - )) - })?; - let headers = signed_headers(&request, &body).await?; + let outbound = outbound_request(&request).await?; - let mut request_builder = http_client().post(&request.url).body(body); - for (key, value) in &headers { - request_builder = request_builder.header(key, value); - } - if let Some(duration) = request.timeout { - request_builder = request_builder.timeout(duration); - } - - let response = http_request(request_builder).await.map_err(|err| { + let response = outbound.send(http_client()).await.map_err(|err| { // Failing to establish the connection means the request never went out, // so the host can still serve it. Everything else here, a timeout // above all, may have reached the provider and been answered. @@ -77,57 +64,24 @@ pub(super) fn as_response_error(err: Error) -> Error { } } -pub(super) async fn signed_headers( +pub(super) async fn outbound_request( request: &ProviderChatCompletionsRequest, - body: &[u8], -) -> Result, Error> { - use std::{collections::BTreeMap, time::SystemTime}; - - use litellm_auth_aws::{ - aws_auth_config, aws_signature_headers, host_supplied_credentials, - is_sigv4_computed_header, resolve_credentials, sign_bedrock_post, - }; - - let ChatCompletionsAuth::AwsSigV4 { region } = &request.auth else { - return Ok(request.upstream_headers.clone()); - }; - // Reattaching a header the signer also emits would put both copies on the - // wire, and Bedrock rejects that pair. Python instead drops the caller's - // copy and prefers a forwarded Authorization over the signature, so leave - // the request to Python rather than serving it a different way here. - if request - .upstream_headers - .iter() - .any(|(name, _)| is_sigv4_computed_header(name)) - { - return Err(Error::Unsupported( - "request forwards a header AWS SigV4 computes", - )); - } - let env_lookup = |key: &str| std::env::var(key).ok(); - let unsigned: BTreeMap = request.upstream_headers.iter().cloned().collect(); - // A host with its own resolution chain hands the result down; only fall - // back to deriving credentials here when it supplied none. - let credentials = match host_supplied_credentials(&request.optional_params) { - Some(credentials) => credentials, - None => { - resolve_credentials( - aws_auth_config(&request.optional_params, &env_lookup), - &env_lookup, - ) - .await? +) -> Result { + crate::outbound::outbound_request( + &request.auth, + request.url.clone(), + request.upstream_headers.clone(), + &request.body, + request.timeout, + &request.optional_params, + ) + .await + .map_err(|error| match error { + // Python drops the caller's copy and prefers a forwarded Authorization + // over the signature, so leave the request to it. + Error::Http(litellm_http::Error::ComputedHeader(_)) => { + Error::Unsupported("request forwards a header AWS SigV4 computes") } - }; - let signature = sign_bedrock_post( - &request.url, - body, - &aws_signature_headers(&unsigned), - region, - &credentials, - SystemTime::now(), - )?; - // Every original header goes back on the wire alongside the computed ones, - // as Python reattaches them. The guard above already rejected the names - // that would collide, so no name appears twice. - Ok(unsigned.into_iter().chain(signature).collect()) + other => other, + }) } diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index d0aa1e88011..c8e6365121e 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -1,6 +1,6 @@ use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider}; use litellm_http::request::has_header; -use litellm_llms::base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth}; +use litellm_llms::base_llm::chat::transformation::{BaseConfig, RequestAuth}; use litellm_types::llms::openai::ChatMessage; use serde_json::Value; @@ -67,7 +67,7 @@ fn validate_environment( request: &ResolvedChatCompletionsRequest<'_>, model: &str, config: &dyn BaseConfig, -) -> Result<(Vec<(String, String)>, ChatCompletionsAuth), Error> { +) -> Result<(Vec<(String, String)>, RequestAuth), Error> { let env_lookup = |key: &str| std::env::var(key).ok(); let mut headers = string_headers(request.extra_headers.clone())?; let auth = config.auth( @@ -77,7 +77,7 @@ fn validate_environment( &env_lookup, )?; match &auth { - ChatCompletionsAuth::Header { name, value } => { + RequestAuth::Header { name, value } => { // The deployment's credential replaces whatever the caller forwarded // under the same name, mirroring Python's // `{**headers, **anthropic_headers}`: letting a request header win @@ -92,7 +92,7 @@ fn validate_environment( headers.push(((*name).to_string(), value.clone())); } } - ChatCompletionsAuth::Bearer { token } => { + RequestAuth::Bearer { token } => { // Bedrock's `get_request_headers` assigns `headers["Authorization"]` // unconditionally once a bearer token resolves, so the deployment's // identity outranks whatever the caller forwarded. Keeping the @@ -105,7 +105,7 @@ fn validate_environment( headers.push(("authorization".to_string(), format!("Bearer {token}"))); } // SigV4 signs the serialized body, so the handler adds its headers. - ChatCompletionsAuth::AwsSigV4 { .. } => {} + RequestAuth::AwsSigV4 { .. } => {} } for (name, value) in config.default_headers() { diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs index dcaa3397add..dd5938cf168 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -1,4 +1,4 @@ -use litellm_llms::base_llm::chat::transformation::ChatCompletionsAuth; +use litellm_llms::base_llm::chat::transformation::RequestAuth; use serde_json::{Map, Value, json}; use super::{ @@ -90,7 +90,7 @@ fn adds_the_auth_and_default_headers() { ); assert!(matches!( prepared.auth, - ChatCompletionsAuth::Header { + RequestAuth::Header { name: "x-api-key", .. } @@ -289,8 +289,9 @@ fn prepares_a_bedrock_call_without_resolving_credentials() { ); assert_eq!( prepared.auth, - ChatCompletionsAuth::AwsSigV4 { - region: "us-east-1".to_string() + RequestAuth::AwsSigV4 { + region: "us-east-1".to_string(), + service: "bedrock", } ); // SigV4 signs the serialized body, so prepare must not have added an @@ -326,15 +327,14 @@ async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() { json!("abc-123"), )])); let prepared = prepare_chat_completions_call(call).expect("prepares"); - let signed = super::handler::signed_headers(&prepared, br#"{"a":1}"#) + let signed = super::handler::outbound_request(&prepared) .await .expect("signs"); let authorization = signed - .iter() - .find(|(name, _)| name.eq_ignore_ascii_case("authorization")) - .map(|(_, value)| value.clone()) - .expect("carries an authorization header"); + .header("authorization") + .expect("carries an authorization header") + .to_string(); assert!( authorization.starts_with("AWS4-HMAC-SHA256"), "expected a SigV4 signature, got {authorization}" @@ -346,6 +346,7 @@ async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() { // It still goes on the wire, it is just not part of the signature. assert!( signed + .headers() .iter() .any(|(name, value)| name == "x-request-id" && value == "abc-123"), "forwarded header was dropped instead of reattached" @@ -376,7 +377,7 @@ async fn a_forwarded_header_the_signer_computes_declines_to_python() { call.api_key = None; call.extra_headers = Some(Map::from_iter([(forwarded.to_string(), json!("forged"))])); let prepared = prepare_chat_completions_call(call).expect("prepares"); - let error = super::handler::signed_headers(&prepared, br#"{"a":1}"#) + let error = super::handler::outbound_request(&prepared) .await .expect_err("{forwarded} should decline instead of being signed"); assert!( @@ -466,7 +467,7 @@ fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() { .expect("prepares"); assert_eq!( prepared.auth, - ChatCompletionsAuth::Bearer { + RequestAuth::Bearer { token: "sk-test".to_string() } ); diff --git a/litellm-rust/crates/core/src/chat_completions/types.rs b/litellm-rust/crates/core/src/chat_completions/types.rs index 882611d5862..3b74cf5dace 100644 --- a/litellm-rust/crates/core/src/chat_completions/types.rs +++ b/litellm-rust/crates/core/src/chat_completions/types.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use litellm_llms::base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth}; +use litellm_llms::base_llm::chat::transformation::{BaseConfig, RequestAuth}; use litellm_types::llms::openai::ChatMessage; use serde_json::{Map, Value}; @@ -38,7 +38,7 @@ pub struct ProviderChatCompletionsRequest { pub url: String, pub body: Value, pub upstream_headers: Vec<(String, String)>, - pub auth: ChatCompletionsAuth, + pub auth: RequestAuth, pub optional_params: Map, pub timeout: Option, } diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index e3e2fb48721..afe5ea595aa 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -4,6 +4,7 @@ pub mod constants; pub mod error; pub mod messages; pub mod ocr; +mod outbound; pub mod responses; pub use error::Error; diff --git a/litellm-rust/crates/core/src/ocr/arguments.rs b/litellm-rust/crates/core/src/ocr/arguments.rs index 43f1c6d6d43..05aa345f01d 100644 --- a/litellm-rust/crates/core/src/ocr/arguments.rs +++ b/litellm-rust/crates/core/src/ocr/arguments.rs @@ -15,6 +15,18 @@ const AZURE_AUTH_OPTION_FIELDS: &[&str] = &[ "azure_federated_token_file", "enable_azure_ad_token_refresh", ]; +const AWS_AUTH_OPTION_FIELDS: &[&str] = &[ + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "aws_region_name", + "aws_session_name", + "aws_profile_name", + "aws_role_name", + "aws_web_identity_token", + "aws_sts_endpoint", + "aws_external_id", +]; const VERTEX_AUTH_OPTION_FIELDS: &[&str] = &[ "vertex_credentials", "vertex_ai_credentials", @@ -35,6 +47,7 @@ pub fn consumed_optional_param_names( let (model, config) = resolve_provider_config(model, custom_llm_provider)?; let provider_fields = config.get_supported_ocr_params(&model); let auth_fields: &[&str] = match config { + OcrConfigKind::AwsTextract | OcrConfigKind::AwsTextractAnalyze => AWS_AUTH_OPTION_FIELDS, OcrConfigKind::AzureAi | OcrConfigKind::AzureDocumentIntelligence | OcrConfigKind::AzureCohere => AZURE_AUTH_OPTION_FIELDS, @@ -57,6 +70,9 @@ pub(crate) fn is_secret_param(name: &str) -> bool { | "azure_federated_token_file" | "vertex_credentials" | "vertex_ai_credentials" + | "aws_secret_access_key" + | "aws_session_token" + | "aws_web_identity_token" ) } diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index c977f721a70..f298f106a5f 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -8,6 +8,10 @@ pub mod route; pub mod types; pub mod wire; +#[cfg(test)] +#[path = "../../tests/aws_textract_ocr.rs"] +mod aws_textract_tests; + #[cfg(test)] #[path = "../../tests/azure_ai_ocr.rs"] mod azure_ai_tests; diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index f1e1dcaaa6d..715aedc69df 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -19,7 +19,10 @@ pub(crate) fn prepare_request( Some("MISTRAL_AZURE_API_BASE"), ), OcrProvider::AzureAi => (None, Some("AZURE_AI_API_BASE")), - OcrProvider::Cohere | OcrProvider::Reducto | OcrProvider::VertexAi => (None, None), + OcrProvider::AwsTextract + | OcrProvider::Cohere + | OcrProvider::Reducto + | OcrProvider::VertexAi => (None, None), }; let secret = |name: &str| client.secrets().truthy(name); let dynamic_api_key = credentials.dynamic_api_key.or_else(|| { diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs index ee9ba76928d..d38d87b92cc 100644 --- a/litellm-rust/crates/core/src/ocr/provider_config.rs +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -1,5 +1,9 @@ use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider}; use litellm_llms::{ + aws_textract::ocr::{ + analyze_transformation::TextractAnalyzeDocumentConfig, common_utils::TextractOperation, + transformation::TextractDetectTextConfig, + }, azure_ai::ocr::{ cohere_parse_transformation::AzureAICohereParseConfig, document_intelligence::transformation::AzureDocumentIntelligenceOcrConfig, @@ -25,6 +29,14 @@ use strum::{EnumString, IntoStaticStr}; macro_rules! with_config { ($kind:expr, $config:ident => $body:expr) => { match $kind { + OcrConfigKind::AwsTextract => { + let $config = TextractDetectTextConfig; + $body + } + OcrConfigKind::AwsTextractAnalyze => { + let $config = TextractAnalyzeDocumentConfig; + $body + } OcrConfigKind::Cohere => { let $config = CohereParseConfig; $body @@ -67,6 +79,8 @@ macro_rules! with_config { #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum OcrConfigKind { + AwsTextract, + AwsTextractAnalyze, Cohere, Mistral, AzureAi, @@ -81,6 +95,7 @@ pub(crate) enum OcrConfigKind { impl OcrConfigKind { pub(crate) const fn provider(self) -> OcrProvider { match self { + Self::AwsTextract | Self::AwsTextractAnalyze => OcrProvider::AwsTextract, Self::Cohere => OcrProvider::Cohere, Self::Mistral => OcrProvider::Mistral, Self::AzureAi | Self::AzureCohere | Self::AzureDocumentIntelligence => { @@ -141,6 +156,7 @@ pub fn get_health_check_document( #[derive(Clone, Copy, Debug, EnumString, IntoStaticStr, PartialEq, Eq)] #[strum(serialize_all = "snake_case")] pub(crate) enum OcrProvider { + AwsTextract, Cohere, Mistral, AzureAi, @@ -162,6 +178,10 @@ pub(crate) fn resolve_provider_config( .parse::() .map_err(|_| Error::InvalidProvider(provider.custom_llm_provider.to_string()))?; let config = match ocr_provider { + OcrProvider::AwsTextract => match TextractOperation::from_model(provider.model)? { + TextractOperation::DetectDocumentText => OcrConfigKind::AwsTextract, + TextractOperation::AnalyzeDocument => OcrConfigKind::AwsTextractAnalyze, + }, OcrProvider::Cohere => OcrConfigKind::Cohere, OcrProvider::Mistral => OcrConfigKind::Mistral, OcrProvider::AzureAi if is_document_intelligence_model(provider.model) => { @@ -419,6 +439,22 @@ mod tests { } #[rstest] + #[case::misspelled_operation("aws_textract/analyse-document")] + #[case::operation_name_from_the_api("aws_textract/AnalyzeDocument")] + fn textract_models_outside_its_two_operations_are_refused(#[case] model: &str) { + assert!(matches!( + resolve_provider_config(model, None), + Err(Error::InvalidModel { + provider: "aws_textract", + .. + }) + )); + } + + #[rstest] + #[case("aws_textract/detect-document-text", OcrConfigKind::AwsTextract)] + #[case("aws_textract/analyze-document", OcrConfigKind::AwsTextractAnalyze)] + #[case("aws_textract/Analyze-Document", OcrConfigKind::AwsTextractAnalyze)] #[case("reducto/parse-legacy", OcrConfigKind::ReductoLegacy)] #[case("reducto/future-parse-model", OcrConfigKind::ReductoV3)] #[case("azure_ai/Cohere-parse-v5", OcrConfigKind::AzureCohere)] diff --git a/litellm-rust/crates/core/src/outbound.rs b/litellm-rust/crates/core/src/outbound.rs new file mode 100644 index 00000000000..7fc90084e6f --- /dev/null +++ b/litellm-rust/crates/core/src/outbound.rs @@ -0,0 +1,30 @@ +use std::time::Duration; + +use litellm_auth::RequestAuth; +use litellm_auth_aws::SigV4Signer; +use litellm_http::outbound::OutboundRequest; +use serde_json::{Map, Value}; + +/// Header credentials are already in `headers`; SigV4 is applied here, over the +/// bytes that are sent. +pub(crate) async fn outbound_request( + auth: &RequestAuth, + url: String, + headers: Vec<(String, String)>, + body: &Value, + timeout: Option, + optional_params: &Map, +) -> Result +where + E: From + From, +{ + let RequestAuth::AwsSigV4 { region, service } = auth else { + return Ok(OutboundRequest::json(url, headers, body, timeout)?); + }; + let env_lookup = |key: &str| std::env::var(key).ok(); + let signer = + SigV4Signer::resolve(region.clone(), service, optional_params, &env_lookup).await?; + Ok(OutboundRequest::signed_json( + url, headers, body, timeout, &signer, + )?) +} diff --git a/litellm-rust/crates/core/tests/aws_textract_ocr.rs b/litellm-rust/crates/core/tests/aws_textract_ocr.rs new file mode 100644 index 00000000000..c536317ad5c --- /dev/null +++ b/litellm-rust/crates/core/tests/aws_textract_ocr.rs @@ -0,0 +1,193 @@ +use std::{collections::BTreeMap, time::SystemTime}; + +use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post}; +use litellm_llms::base_llm::ocr::error::Error; +use serde_json::{Value, json}; +use time::{PrimitiveDateTime, format_description}; + +use crate::ocr::{ + route::LocalOcrHost, + test_support::{ + MockResponse, header, mock_server, perform_ocr_with, request_body, + wire_request_with_document, + }, + types::LiteLLMOcrRequest, +}; + +const ACCESS_KEY_ID: &str = "AKIDEXAMPLE"; +const SECRET_ACCESS_KEY: &str = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"; + +fn textract_request(base: &str) -> LiteLLMOcrRequest { + textract_request_for("aws_textract/detect-document-text", base) +} + +fn textract_request_for(model: &str, base: &str) -> LiteLLMOcrRequest { + wire_request_with_document( + model, + &format!("{base}/"), + json!({"type": "image_url", "image_url": "data:image/png;base64,b3JpZ2luYWw="}), + json!({ + "aws_access_key_id": ACCESS_KEY_ID, + "aws_secret_access_key": SECRET_ACCESS_KEY, + "aws_region_name": "eu-west-1" + }), + ) +} + +fn textract_response() -> MockResponse { + MockResponse::json(json!({ + "DocumentMetadata": {"Pages": 1}, + "Blocks": [{"BlockType": "PAGE"}, {"BlockType": "LINE", "Text": "Invoice 12345"}] + })) +} + +/// Recomputes SigV4 over the bytes the server received, at the time the client claimed. +fn expected_authorization(url: &str, raw_request: &str) -> String { + let format = + format_description::parse_borrowed::<2>("[year][month][day]T[hour][minute][second]Z") + .unwrap(); + let signed_at: SystemTime = + PrimitiveDateTime::parse(header(raw_request, "x-amz-date").unwrap(), &format) + .unwrap() + .assume_utc() + .into(); + let headers: BTreeMap = ["content-type", "x-amz-target"] + .into_iter() + .map(|name| { + ( + name.to_string(), + header(raw_request, name).unwrap().to_string(), + ) + }) + .collect(); + let body = raw_request.split_once("\r\n\r\n").unwrap().1; + sign_post( + url, + body.as_bytes(), + &aws_signature_headers(&headers), + "eu-west-1", + "textract", + &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), + signed_at, + ) + .unwrap()["Authorization"] + .clone() +} + +#[tokio::test] +async fn the_request_is_signed_for_textract_and_lines_become_the_page() { + let (base, seen, server) = mock_server(vec![textract_response()]).await; + + let response = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) + .await + .unwrap(); + server.await.unwrap(); + + let raw = seen.lock().unwrap()[0].clone(); + assert_eq!( + header(&raw, "x-amz-target"), + Some("Textract.DetectDocumentText") + ); + assert_eq!( + header(&raw, "content-type"), + Some("application/x-amz-json-1.1") + ); + assert_eq!( + request_body(&raw), + json!({"Document": {"Bytes": "b3JpZ2luYWw="}}) + ); + assert_eq!( + header(&raw, "authorization"), + Some(expected_authorization(&format!("{base}/"), &raw).as_str()) + ); + assert_eq!(response.pages[0].markdown, "Invoice 12345"); + assert_eq!(response.usage_info.unwrap().pages_processed, Some(1)); +} + +#[tokio::test] +async fn a_body_rewritten_by_before_send_is_what_gets_signed_and_sent() { + let (base, seen, server) = mock_server(vec![textract_response()]).await; + let host = LocalOcrHost::new(textract_request(&base)).with_before_send(|mut wire, _| { + assert!( + !wire + .headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case("authorization")), + "the hook ran after signing" + ); + wire.body["Document"]["Bytes"] = Value::from("cmVkYWN0ZWQ="); + Ok(wire) + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + + let raw = seen.lock().unwrap()[0].clone(); + assert_eq!( + request_body(&raw), + json!({"Document": {"Bytes": "cmVkYWN0ZWQ="}}) + ); + assert_eq!( + header(&raw, "authorization"), + Some(expected_authorization(&format!("{base}/"), &raw).as_str()) + ); +} + +#[tokio::test] +async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit() { + let (base, _, server) = mock_server(vec![MockResponse { + status: 400, + headers: vec![], + body: json!({ + "__type": "UnsupportedDocumentException", + "Message": "Request has unsupported document format" + }), + }]) + .await; + + let error = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) + .await + .unwrap_err(); + server.await.unwrap(); + + let Error::Provider { status, body, .. } = error else { + panic!("expected a provider error, got {error:?}"); + }; + assert_eq!(status, 400); + assert!( + body.contains("multi-page documents are not supported"), + "{body}" + ); +} + +#[tokio::test] +async fn analyze_document_asks_for_layout_and_tables_and_returns_markdown() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "DocumentMetadata": {"Pages": 1}, + "Blocks": [ + {"Id": "l1", "BlockType": "LINE", "Text": "Quarterly Report"}, + {"Id": "t", "BlockType": "LAYOUT_TITLE", + "Relationships": [{"Type": "CHILD", "Ids": ["l1"]}]} + ] + }))]) + .await; + let request = textract_request_for("aws_textract/analyze-document", &base); + + let response = perform_ocr_with(LocalOcrHost::new(request)).await.unwrap(); + server.await.unwrap(); + + let raw = seen.lock().unwrap()[0].clone(); + assert_eq!( + header(&raw, "x-amz-target"), + Some("Textract.AnalyzeDocument") + ); + assert_eq!( + request_body(&raw)["FeatureTypes"], + json!(["LAYOUT", "TABLES"]) + ); + assert_eq!( + header(&raw, "authorization"), + Some(expected_authorization(&format!("{base}/"), &raw).as_str()) + ); + assert_eq!(response.pages[0].markdown, "# Quarterly Report"); +} diff --git a/litellm-rust/crates/core/tests/cohere_ocr.rs b/litellm-rust/crates/core/tests/cohere_ocr.rs index fc1203f0980..12824f58b1d 100644 --- a/litellm-rust/crates/core/tests/cohere_ocr.rs +++ b/litellm-rust/crates/core/tests/cohere_ocr.rs @@ -38,7 +38,7 @@ mod transformation { ) .await .unwrap(); - let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap(); + let body: Value = serde_json::from_slice(http.body()).unwrap(); assert_eq!( body, json!({ @@ -75,7 +75,7 @@ mod transformation { ) .await .unwrap(); - let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap(); + let body: Value = serde_json::from_slice(http.body()).unwrap(); assert_eq!(body["output_format"], "markdown"); assert!(body.get("req_format").is_none()); } diff --git a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs index 399b7cac39a..035f3fe944d 100644 --- a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs +++ b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs @@ -160,17 +160,16 @@ async fn adapters_build_complete_requests_and_share_mistral_normalization() { .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) .await .unwrap(); - assert_eq!(direct_http.url().as_str(), "https://mistral.test/v1/ocr"); + assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); assert_eq!( - vertex_http.url().as_str(), + vertex_http.url(), "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" ); for http in [&direct_http, &vertex_http] { - assert_eq!(http.method(), reqwest::Method::POST); - assert_eq!(http.headers()["authorization"], "Bearer test-key"); - assert_eq!(http.headers()["content-type"], "application/json"); - assert_eq!(http.timeout(), Some(&Duration::from_secs(2))); - let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); + assert_eq!(http.header("content-type").unwrap(), "application/json"); + assert_eq!(http.timeout(), Some(Duration::from_secs(2))); + let body: Value = serde_json::from_slice(http.body()).unwrap(); assert_eq!( body, json!({ @@ -250,9 +249,9 @@ mod transformation { .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) .await .unwrap(); - assert_eq!(direct_http.url().as_str(), "https://mistral.test/v1/ocr"); + assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); assert_eq!( - vertex_http.url().as_str(), + vertex_http.url(), "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" ); let http = if use_vertex { @@ -260,11 +259,10 @@ mod transformation { } else { &direct_http }; - assert_eq!(http.method(), reqwest::Method::POST); - assert_eq!(http.headers()["authorization"], "Bearer test-key"); - assert_eq!(http.headers()["content-type"], "application/json"); - assert_eq!(http.timeout(), Some(&Duration::from_secs(2))); - let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); + assert_eq!(http.header("content-type").unwrap(), "application/json"); + assert_eq!(http.timeout(), Some(Duration::from_secs(2))); + let body: Value = serde_json::from_slice(http.body()).unwrap(); assert_eq!( body, json!({ diff --git a/litellm-rust/crates/host-python/src/callable.rs b/litellm-rust/crates/host-python/src/callable.rs index 424db002b0a..2e454422e95 100644 --- a/litellm-rust/crates/host-python/src/callable.rs +++ b/litellm-rust/crates/host-python/src/callable.rs @@ -73,13 +73,7 @@ abort = KeyboardInterrupt('cancelled') let wrapped = wrap_failure(py, TEMPLATE, failure(&original)).unwrap_err(); assert!(wrapped.is_instance_of::(py)); assert!(wrapped.cause(py).unwrap().value(py).is(&original)); - assert!( - wrapped - .value(py) - .getattr("__context__") - .unwrap() - .is(&original) - ); + assert!(wrapped.context(py).unwrap().value(py).is(&original)); assert_eq!( wrapped.value(py).str().unwrap().to_str().unwrap(), "Failed to reach the caller: unavailable" @@ -115,13 +109,7 @@ original = Unformattable('cannot render') let original = raised(&locals, "original"); let error = wrap_failure(py, TEMPLATE, failure(&original)).unwrap_err(); assert!(error.is_instance_of::(py)); - assert!( - error - .value(py) - .getattr("__context__") - .unwrap() - .is(&original) - ); + assert!(error.context(py).unwrap().value(py).is(&original)); }); } diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index 392d36e10f4..77a294d274b 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -445,14 +445,8 @@ where Ok(failure) => return failure.into(), Err(classifier_error) => classifier_error, }; - let attached = classifier_error.value(py).setattr( - "__context__", - PyRuntimeError::new_err(native).into_value(py), - ); - match attached { - Ok(()) => classifier_error, - Err(error) => error, - } + classifier_error.set_context(py, Some(PyRuntimeError::new_err(native))); + classifier_error } fn succeeded(&mut self, py: Python<'_>, response: Py) -> PyResult { @@ -1071,9 +1065,9 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let error = result.unwrap_err(); assert!(error.is_instance_of::(py)); assert_eq!(error.value(py).to_string(), "classifier failed"); - let context = error.value(py).getattr("__context__").unwrap(); - assert!(context.is_instance_of::()); - assert_eq!(context.str().unwrap().to_string(), "provider exploded"); + let context = error.context(py).unwrap(); + assert!(context.is_instance_of::(py)); + assert_eq!(context.value(py).to_string(), "provider exploded"); assert_eq!( log, [ @@ -1186,20 +1180,12 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri type Failure = Classified; fn invoke( &mut self, - py: Python<'_>, + _: Python<'_>, _: &Bound<'_, PyDict>, _: &'static str, ) -> Result> { self.0.push("route"); - Err(PyErr::from_value( - py.import("asyncio") - .unwrap() - .getattr("CancelledError") - .unwrap() - .call0() - .unwrap(), - ) - .into()) + Err(pyo3::exceptions::asyncio::CancelledError::new_err(()).into()) } fn chunk( &mut self, diff --git a/litellm-rust/crates/http/Cargo.toml b/litellm-rust/crates/http/Cargo.toml index d4457f5685c..cad5aa87e49 100644 --- a/litellm-rust/crates/http/Cargo.toml +++ b/litellm-rust/crates/http/Cargo.toml @@ -14,6 +14,7 @@ litellm-core-utils.workspace = true hyper-util.workspace = true reqwest.workspace = true rustls.workspace = true +serde.workspace = true serde_json.workspace = true thiserror.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs index 697d0cf59c8..e06f7c00cf5 100644 --- a/litellm-rust/crates/http/src/error.rs +++ b/litellm-rust/crates/http/src/error.rs @@ -8,6 +8,12 @@ pub enum Error { InvalidPem { path: PathBuf, message: String }, #[error("could not build the HTTP client: {0}")] Client(String), + #[error("request body could not be serialized: {0}")] + RequestBody(String), + #[error("request forwards a header the signer computes: {0}")] + ComputedHeader(String), + #[error("request signing failed: {0}")] + Signature(String), } impl From for Error { diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index c6d9959348d..6f62a00175c 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -1,6 +1,7 @@ mod config; mod error; pub mod media; +pub mod outbound; mod pool; mod proxy; pub mod request; diff --git a/litellm-rust/crates/http/src/outbound.rs b/litellm-rust/crates/http/src/outbound.rs new file mode 100644 index 00000000000..d100bdf624b --- /dev/null +++ b/litellm-rust/crates/http/src/outbound.rs @@ -0,0 +1,210 @@ +//! The request a route hands to the transport. The body is serialized once, +//! when the request is built, and a [`RequestSigner`] sees those exact bytes. +//! +//! Host hooks may rewrite the wire request (redaction, guardrails) and a +//! signature such as AWS SigV4 covers the body, so a route builds this after +//! its hooks ran and cannot change or re-serialize it afterwards. + +use std::time::Duration; + +use serde::Serialize; + +use crate::{ + Error, + request::{HeaderPolicy, has_header, with_headers}, +}; + +#[derive(Clone, Copy, Debug)] +pub struct UnsignedRequest<'a> { + pub url: &'a str, + pub headers: &'a [(String, String)], + pub body: &'a [u8], +} + +/// Returns the headers to add to the request; it never sees a mutable request. +pub trait RequestSigner: Send + Sync { + fn sign(&self, request: UnsignedRequest<'_>) -> Result, Error>; +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct OutboundRequest { + url: String, + headers: Vec<(String, String)>, + body: Vec, + timeout: Option, +} + +impl OutboundRequest { + pub fn json( + url: String, + headers: Vec<(String, String)>, + body: &impl Serialize, + timeout: Option, + ) -> Result { + Self::build(url, headers, body, timeout, None) + } + + pub fn signed_json( + url: String, + headers: Vec<(String, String)>, + body: &impl Serialize, + timeout: Option, + signer: &dyn RequestSigner, + ) -> Result { + Self::build(url, headers, body, timeout, Some(signer)) + } + + fn build( + url: String, + headers: Vec<(String, String)>, + body: &impl Serialize, + timeout: Option, + signer: Option<&dyn RequestSigner>, + ) -> Result { + let body = + serde_json::to_vec(body).map_err(|error| Error::RequestBody(error.to_string()))?; + let content_type = (!has_header(&headers, "content-type")) + .then(|| ("content-type".to_string(), "application/json".to_string())); + let unsigned: Vec<(String, String)> = headers.into_iter().chain(content_type).collect(); + let signature = signer + .map(|signer| { + signer.sign(UnsignedRequest { + url: &url, + headers: &unsigned, + body: &body, + }) + }) + .transpose()? + .unwrap_or_default(); + Ok(Self { + url, + headers: unsigned.into_iter().chain(signature).collect(), + body, + timeout, + }) + } + + pub fn url(&self) -> &str { + &self.url + } + + pub fn headers(&self) -> &[(String, String)] { + &self.headers + } + + pub fn header(&self, name: &str) -> Option<&str> { + self.headers + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.as_str()) + } + + pub fn body(&self) -> &[u8] { + &self.body + } + + pub fn timeout(&self) -> Option { + self.timeout + } + + pub async fn send(self, client: &reqwest::Client) -> Result { + let builder = with_headers( + client.post(&self.url).body(self.body), + &self.headers, + HeaderPolicy::All, + ); + match self.timeout { + Some(timeout) => builder.timeout(timeout), + None => builder, + } + .send() + .await + } +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use serde_json::json; + + use super::*; + + #[derive(Default)] + struct Recording(Mutex>); + + impl RequestSigner for Recording { + fn sign(&self, request: UnsignedRequest<'_>) -> Result, Error> { + *self.0.lock().unwrap() = request.body.to_vec(); + Ok(vec![("authorization".into(), "signed".into())]) + } + } + + #[test] + fn the_signer_sees_exactly_the_bytes_that_are_sent() { + let signer = Recording::default(); + let request = OutboundRequest::signed_json( + "https://provider.test/".into(), + vec![("x-caller".into(), "kept".into())], + &json!({"b": 1, "a": [true, null]}), + None, + &signer, + ) + .unwrap(); + + assert_eq!(request.body(), signer.0.lock().unwrap().as_slice()); + assert_eq!(request.header("authorization"), Some("signed")); + assert_eq!(request.header("x-caller"), Some("kept")); + } + + #[test] + fn the_content_type_is_part_of_what_the_signer_sees() { + struct RequiresContentType; + impl RequestSigner for RequiresContentType { + fn sign(&self, request: UnsignedRequest<'_>) -> Result, Error> { + has_header(request.headers, "content-type") + .then(Vec::new) + .ok_or_else(|| Error::Signature("content-type was not signed".into())) + } + } + + let defaulted = OutboundRequest::signed_json( + "u".into(), + Vec::new(), + &json!({}), + None, + &RequiresContentType, + ) + .unwrap(); + assert_eq!(defaulted.header("content-type"), Some("application/json")); + + let provider = OutboundRequest::signed_json( + "u".into(), + vec![("Content-Type".into(), "application/x-amz-json-1.1".into())], + &json!({}), + None, + &RequiresContentType, + ) + .unwrap(); + assert_eq!( + provider.header("content-type"), + Some("application/x-amz-json-1.1") + ); + assert_eq!(provider.headers().len(), 1); + } + + #[test] + fn a_signer_failure_produces_no_request() { + struct Refuses; + impl RequestSigner for Refuses { + fn sign(&self, _request: UnsignedRequest<'_>) -> Result, Error> { + Err(Error::ComputedHeader("authorization".into())) + } + } + + assert_eq!( + OutboundRequest::signed_json("u".into(), Vec::new(), &json!({}), None, &Refuses), + Err(Error::ComputedHeader("authorization".into())) + ); + } +} diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 7afc4171ca8..0cc7af1836f 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -27,6 +27,7 @@ serde.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } serde_path_to_error = "0.1" serde_with.workspace = true +strum.workspace = true thiserror.workspace = true time.workspace = true tokio = { workspace = true, features = ["sync"] } diff --git a/litellm-rust/crates/llms/src/anthropic/chat/tests.rs b/litellm-rust/crates/llms/src/anthropic/chat/tests.rs index 40e89c52c3c..3777347d240 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/tests.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/tests.rs @@ -428,7 +428,7 @@ fn resolves_the_messages_url_and_x_api_key_auth() { config .auth(Some("sk-x"), "claude-sonnet-4-5", &Map::new(), &|_| None) .expect("auth resolves"), - ChatCompletionsAuth::Header { + RequestAuth::Header { name: "x-api-key", value: "sk-x".to_string() } diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index 21fa4e9f82e..6fc4f00b981 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -16,7 +16,7 @@ use crate::{ }, }, base_llm::chat::transformation::{ - BaseConfig, ChatCompletionsAuth, Error, ProviderChatRequestData, ProviderChatResponseData, + BaseConfig, Error, ProviderChatRequestData, ProviderChatResponseData, RequestAuth, Unsupported, unsupported_message, unsupported_param, }, }; @@ -137,8 +137,8 @@ impl BaseConfig for AnthropicConfig { _model: &str, _optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - Ok(ChatCompletionsAuth::Header { + ) -> Result { + Ok(RequestAuth::Header { name: "x-api-key", value: resolve_anthropic_api_key(api_key, env_lookup)?, }) diff --git a/litellm-rust/crates/llms/src/aws_textract/mod.rs b/litellm-rust/crates/llms/src/aws_textract/mod.rs new file mode 100644 index 00000000000..3621ff6a2fd --- /dev/null +++ b/litellm-rust/crates/llms/src/aws_textract/mod.rs @@ -0,0 +1 @@ +pub mod ocr; diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/AGENTS.md b/litellm-rust/crates/llms/src/aws_textract/ocr/AGENTS.md new file mode 100644 index 00000000000..4913f89924a --- /dev/null +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/AGENTS.md @@ -0,0 +1,12 @@ +- https://docs.aws.amazon.com/textract/latest/APIReference/Welcome.md +- https://docs.aws.amazon.com/textract/latest/APIReference/API_Operations.md +- https://docs.aws.amazon.com/textract/latest/APIReference/API_DetectDocumentText.md +- https://docs.aws.amazon.com/textract/latest/APIReference/API_AnalyzeDocument.md +- https://docs.aws.amazon.com/textract/latest/APIReference/API_StartDocumentTextDetection.md +- https://docs.aws.amazon.com/textract/latest/APIReference/API_Document.md +- https://docs.aws.amazon.com/textract/latest/APIReference/API_Block.md +- https://docs.aws.amazon.com/textract/latest/dg/what-is.md +- https://docs.aws.amazon.com/textract/latest/dg/sync.md +- https://docs.aws.amazon.com/textract/latest/dg/async.md +- https://docs.aws.amazon.com/textract/latest/dg/how-it-works-document-layout.md +- https://docs.aws.amazon.com/textract/latest/dg/limits.md diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs new file mode 100644 index 00000000000..d476861e6e1 --- /dev/null +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs @@ -0,0 +1,479 @@ +use std::collections::{BTreeMap, BTreeSet, HashMap}; + +use litellm_core_utils::call_arguments::{CallArguments, parse_options}; +use serde::{Deserialize, Serialize}; + +use super::common_utils::{ + Block, BlockType, FeatureType, LayoutType, TextractDocument, TextractEnvironment, + TextractOperation, TextractResponse, document_bytes, endpoint, environment, error_class, + health_check_document, inline_document, lines_by_page, ocr_response, +}; +use crate::base_llm::ocr::{ + error::Error, + handler::OcrClient, + transformation::{ + BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrRequestContext, OcrResponseFormat, + PreparedOcrRequest, decode_and_normalize_response, + }, +}; + +const DEFAULT_FEATURE_TYPES: [FeatureType; 2] = [FeatureType::Layout, FeatureType::Tables]; + +#[derive(Default, Deserialize)] +pub struct AnalyzeDocumentOptions { + pub feature_types: Option>, +} + +#[derive(Debug, Deserialize, Serialize)] +pub struct AnalyzeDocumentRequest { + #[serde(rename = "Document")] + pub document: TextractDocument, + #[serde(rename = "FeatureTypes")] + pub feature_types: Vec, +} + +#[derive(Clone, Copy, Debug, Default)] +pub struct TextractAnalyzeDocumentConfig; + +impl BaseOcrConfig for TextractAnalyzeDocumentConfig { + type OcrParams = AnalyzeDocumentOptions; + type ProviderRequest = AnalyzeDocumentRequest; + type Environment = TextractEnvironment; + + fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { + &["feature_types"] + } + + fn get_health_check_document(&self) -> OcrDocument { + health_check_document() + } + + fn map_ocr_params( + &self, + non_default_params: &CallArguments, + _model: &str, + ) -> Result { + Ok(parse_options(non_default_params)?) + } + + async fn validate_environment( + &self, + request: &PreparedOcrRequest, + _client: &OcrClient, + ) -> Result { + environment(request, TextractOperation::AnalyzeDocument).await + } + + fn get_complete_url( + &self, + request: &PreparedOcrRequest, + _optional_params: &AnalyzeDocumentOptions, + environment: &TextractEnvironment, + ) -> Result { + Ok(endpoint(request, environment)) + } + + fn transform_ocr_request( + &self, + _model: &str, + document: OcrDocument, + optional_params: &AnalyzeDocumentOptions, + _headers: &[(String, String)], + ) -> Result { + Ok(AnalyzeDocumentRequest { + document: document_bytes(&document)?, + feature_types: optional_params + .feature_types + .clone() + .unwrap_or_else(|| DEFAULT_FEATURE_TYPES.to_vec()), + }) + } + + async fn async_transform_ocr_request( + &self, + model: &str, + document: OcrDocument, + optional_params: &AnalyzeDocumentOptions, + headers: &[(String, String)], + context: OcrRequestContext<'_>, + ) -> Result { + let document = inline_document(document, context).await?; + self.transform_ocr_request(model, document, optional_params, headers) + } + + fn transform_ocr_response( + &self, + model: &str, + raw_response: &[u8], + request_format: OcrResponseFormat, + ) -> Result { + decode_and_normalize_response(model, raw_response, request_format, normalize_response) + } + + fn get_error_class( + &self, + error_message: String, + status_code: u16, + headers: Vec<(String, String)>, + ) -> Error { + error_class(error_message, status_code, headers) + } +} + +fn normalize_response( + model: &str, + response: TextractResponse, +) -> Result { + let blocks = &response.blocks; + let has_layout = blocks + .iter() + .any(|block| block.block_type.layout().is_some()); + let page_markdown: Vec<(i64, String)> = if has_layout { + let by_id: HashMap<&str, &Block> = blocks + .iter() + .map(|block| (block.id.as_str(), block)) + .collect(); + let pages: BTreeSet = blocks.iter().map(Block::page).collect(); + pages + .into_iter() + .map(|page| (page, layout_markdown(blocks, page, &by_id))) + .filter(|(_, markdown)| !markdown.is_empty()) + .collect() + } else { + lines_by_page(blocks) + }; + Ok(ocr_response( + model, + page_markdown, + response.document_metadata, + )) +} + +/// Layout blocks arrive in reading order. A list's items are repeated as +/// top-level `LAYOUT_TEXT` blocks. A `LAYOUT_TABLE` that links to its `TABLE` +/// renders it; one that only links to the table's lines takes the `TABLE` at +/// the same position on the page. +fn layout_markdown(blocks: &[Block], page: i64, by_id: &HashMap<&str, &Block>) -> String { + let on_page = || blocks.iter().filter(move |block| block.page() == page); + let list_items: BTreeSet<&str> = on_page() + .filter(|block| block.block_type == BlockType::LayoutList) + .flat_map(Block::children) + .collect(); + let tables: Vec<&Block> = on_page() + .filter(|block| block.block_type == BlockType::Table) + .collect(); + let table_ordinal: HashMap<&str, usize> = on_page() + .filter(|block| block.block_type == BlockType::LayoutTable) + .enumerate() + .map(|(ordinal, block)| (block.id.as_str(), ordinal)) + .collect(); + let table_of = |layout_table: &Block| { + layout_table + .children() + .filter_map(|id| by_id.get(id).copied()) + .find(|child| child.block_type == BlockType::Table) + .or_else(|| { + table_ordinal + .get(layout_table.id.as_str()) + .and_then(|ordinal| tables.get(*ordinal).copied()) + }) + }; + let sections: Vec = on_page() + .filter(|block| !list_items.contains(block.id.as_str())) + .filter_map(|block| Some((block, block.block_type.layout()?))) + .map(|(block, layout)| match layout { + LayoutType::Title => format!("# {}", text_of(block, by_id, " ")), + LayoutType::SectionHeader => format!("## {}", text_of(block, by_id, " ")), + LayoutType::List => block + .children() + .filter_map(|id| by_id.get(id)) + .map(|item| format!("- {}", strip_bullet(&text_of(item, by_id, " ")))) + .collect::>() + .join("\n"), + LayoutType::Table => match table_of(block) { + Some(table) => table_markdown(table, by_id), + None => text_of(block, by_id, "\n"), + }, + LayoutType::KeyValue => text_of(block, by_id, "\n"), + LayoutType::Figure => String::new(), + LayoutType::Text | LayoutType::Header | LayoutType::Footer | LayoutType::PageNumber => { + text_of(block, by_id, " ") + } + }) + .filter(|section| !section.trim().is_empty()) + .collect(); + sections.join("\n\n") +} + +fn text_of(block: &Block, by_id: &HashMap<&str, &Block>, separator: &str) -> String { + match &block.text { + Some(text) => text.clone(), + None => block + .children() + .filter_map(|id| by_id.get(id)) + .map(|child| text_of(child, by_id, separator)) + .filter(|text| !text.is_empty()) + .collect::>() + .join(separator), + } +} + +fn strip_bullet(item: &str) -> &str { + item.trim_start_matches(['-', '*', '\u{2022}', '\u{00b7}']) + .trim_start() +} + +fn table_markdown(table: &Block, by_id: &HashMap<&str, &Block>) -> String { + let cells: BTreeMap<(usize, usize), String> = table + .children() + .filter_map(|id| by_id.get(id)) + .filter(|cell| cell.block_type == BlockType::Cell) + .filter_map(|cell| { + Some(( + (cell.row_index?, cell.column_index?), + text_of(cell, by_id, " ").replace('|', "\\|"), + )) + }) + .collect(); + let columns = cells.keys().map(|(_, column)| *column).max().unwrap_or(0); + let rows: BTreeSet = cells.keys().map(|(row, _)| *row).collect(); + let render = |row: usize| { + let values: Vec<&str> = (1..=columns) + .map(|column| cells.get(&(row, column)).map_or("", String::as_str)) + .collect(); + format!("| {} |", values.join(" | ")) + }; + let divider = format!("|{}", " --- |".repeat(columns)); + rows.iter() + .enumerate() + .flat_map(|(position, row)| { + std::iter::once(render(*row)).chain((position == 0).then(|| divider.clone())) + }) + .collect::>() + .join("\n") +} + +#[cfg(test)] +mod tests { + use rstest::{fixture, rstest}; + use serde_json::{Value, json}; + + use super::*; + + const MODEL: &str = "analyze-document"; + + #[fixture] + fn document() -> OcrDocument { + OcrDocument::ImageUrl { + image_url: "data:image/png;base64,aGk=".into(), + extra_fields: Default::default(), + } + } + + fn child(ids: &[&str]) -> Value { + json!([{"Type": "CHILD", "Ids": ids}]) + } + + fn line(id: &str, text: &str) -> Value { + json!({"Id": id, "BlockType": "LINE", "Text": text}) + } + + fn word(id: &str, text: &str) -> Value { + json!({"Id": id, "BlockType": "WORD", "Text": text}) + } + + fn layout(id: &str, block_type: &str, children: &[&str]) -> Value { + json!({"Id": id, "BlockType": block_type, "Relationships": child(children)}) + } + + fn table(id: &str, cells: &[&str]) -> Value { + json!({"Id": id, "BlockType": "TABLE", "Relationships": child(cells)}) + } + + fn cell(id: &str, row: usize, column: usize, words: &[&str]) -> Value { + json!({"Id": id, "BlockType": "CELL", "RowIndex": row, "ColumnIndex": column, + "Relationships": child(words)}) + } + + fn on_page(page: i64, mut block: Value) -> Value { + block["Page"] = json!(page); + block + } + + #[rstest] + #[case::headings_paragraphs_and_a_list_without_repeating_its_items( + json!([ + line("l1", "Quarterly Report"), + line("l2", "This report lists"), + line("l3", "the invoices."), + line("l4", "Line items"), + line("l5", "- Pay within 30 days"), + line("l6", "\u{2022} Quote the number"), + layout("t", "LAYOUT_TITLE", &["l1"]), + layout("p", "LAYOUT_TEXT", &["l2", "l3"]), + layout("h", "LAYOUT_SECTION_HEADER", &["l4"]), + layout("ul", "LAYOUT_LIST", &["i1", "i2"]), + layout("i1", "LAYOUT_TEXT", &["l5"]), + layout("i2", "LAYOUT_TEXT", &["l6"]) + ]), + vec![( + 0, + "# Quarterly Report\n\nThis report lists the invoices.\n\n## Line items\n\n- Pay within 30 days\n- Quote the number" + )] + )] + #[case::header_footer_and_page_number_stay_in_reading_order( + json!([ + line("l1", "ACME Corp"), line("l2", "Body"), line("l3", "Confidential"), line("l4", "3"), + layout("hd", "LAYOUT_HEADER", &["l1"]), + layout("p", "LAYOUT_TEXT", &["l2"]), + layout("ft", "LAYOUT_FOOTER", &["l3"]), + layout("pn", "LAYOUT_PAGE_NUMBER", &["l4"]) + ]), + vec![(0, "ACME Corp\n\nBody\n\nConfidential\n\n3")] + )] + #[case::a_table_is_rendered_from_its_cells_in_row_and_column_order( + json!([ + line("l1", "Invoice"), line("l2", "Total"), line("l3", "12345"), line("l4", "a|b"), + word("w1", "Invoice"), word("w2", "Total"), word("w3", "12345"), word("w4", "a|b"), + {"Id": "tb", "BlockType": "TABLE", "Relationships": [ + {"Type": "CHILD", "Ids": ["c4", "c1", "c3", "c2"]}, + {"Type": "TABLE_TITLE", "Ids": ["title"]} + ]}, + cell("c1", 1, 1, &["w1"]), cell("c2", 1, 2, &["w2"]), + cell("c3", 2, 1, &["w3"]), cell("c4", 2, 2, &["w4"]), + layout("lt", "LAYOUT_TABLE", &["l1", "l2", "l3", "l4"]) + ]), + vec![(0, "| Invoice | Total |\n| --- | --- |\n| 12345 | a\\|b |")] + )] + #[case::a_layout_table_that_links_its_table_renders_that_one( + json!([ + word("w1", "first"), word("w2", "second"), + table("tb1", &["c1"]), cell("c1", 1, 1, &["w1"]), + table("tb2", &["c2"]), cell("c2", 1, 1, &["w2"]), + layout("lt", "LAYOUT_TABLE", &["tb2"]) + ]), + vec![(0, "| second |\n| --- |")] + )] + #[case::a_missing_cell_leaves_an_empty_column( + json!([ + word("w1", "a"), word("w2", "b"), word("w3", "c"), + table("tb", &["c1", "c2", "c3"]), + cell("c1", 1, 1, &["w1"]), cell("c2", 1, 2, &["w2"]), cell("c3", 2, 2, &["w3"]), + layout("lt", "LAYOUT_TABLE", &[]) + ]), + vec![(0, "| a | b |\n| --- | --- |\n| | c |")] + )] + #[case::a_layout_table_without_table_blocks_keeps_its_lines( + json!([ + line("l1", "Invoice Total"), + line("l2", "12345 67.89"), + layout("lt", "LAYOUT_TABLE", &["l1", "l2"]) + ]), + vec![(0, "Invoice Total\n12345 67.89")] + )] + #[case::key_values_keep_one_line_each( + json!([ + line("l1", "Name: Ana"), + line("l2", "Date: 2024-01-01"), + layout("kv", "LAYOUT_KEY_VALUE", &["l1", "l2"]) + ]), + vec![(0, "Name: Ana\nDate: 2024-01-01")] + )] + #[case::a_figure_has_no_markdown( + json!([ + line("l1", "Caption"), + layout("f", "LAYOUT_FIGURE", &[]), + layout("p", "LAYOUT_TEXT", &["l1"]) + ]), + vec![(0, "Caption")] + )] + #[case::a_block_type_added_later_is_ignored( + json!([ + line("l1", "Body"), + layout("new", "LAYOUT_SIDEBAR", &["l1"]), + layout("p", "LAYOUT_TEXT", &["l1"]) + ]), + vec![(0, "Body")] + )] + #[case::without_layout_blocks_lines_are_used( + json!([line("l1", "first"), word("w1", "first"), line("l2", "second")]), + vec![(0, "first\nsecond")] + )] + #[case::each_page_gets_its_own_markdown_and_its_own_tables( + json!([ + on_page(1, line("a", "one")), + on_page(2, line("b", "two")), + on_page(2, word("w", "cell")), + on_page(1, layout("t1", "LAYOUT_TEXT", &["a"])), + on_page(2, table("tb", &["c"])), + on_page(2, cell("c", 1, 1, &["w"])), + on_page(2, layout("lt", "LAYOUT_TABLE", &["b"])) + ]), + vec![(0, "one"), (1, "| cell |\n| --- |")] + )] + fn blocks_become_markdown_pages(#[case] blocks: Value, #[case] expected: Vec<(i64, &str)>) { + let response = TextractAnalyzeDocumentConfig + .transform_ocr_response( + MODEL, + &serde_json::to_vec(&json!({"DocumentMetadata": {"Pages": 1}, "Blocks": blocks})) + .unwrap(), + OcrResponseFormat::Litellm, + ) + .unwrap(); + + let pages: Vec<(i64, &str)> = response + .pages + .iter() + .map(|page| (page.index, page.markdown.as_str())) + .collect(); + assert_eq!(pages, expected); + } + + #[rstest] + #[case::hyphen("- item", "item")] + #[case::asterisk("* item", "item")] + #[case::bullet("\u{2022} item", "item")] + #[case::middle_dot("\u{00b7}item", "item")] + #[case::no_bullet("item - with a dash", "item - with a dash")] + fn list_items_lose_their_own_bullet(#[case] item: &str, #[case] expected: &str) { + assert_eq!(strip_bullet(item), expected); + } + + #[rstest] + #[case::defaults_to_layout_and_tables(json!({}), json!(["LAYOUT", "TABLES"]))] + #[case::overridden(json!({"feature_types": ["FORMS", "SIGNATURES"]}), json!(["FORMS", "SIGNATURES"]))] + #[case::explicit_null_uses_the_default(json!({"feature_types": null}), json!(["LAYOUT", "TABLES"]))] + fn feature_types_reach_the_request( + document: OcrDocument, + #[case] arguments: Value, + #[case] expected: Value, + ) { + let arguments: CallArguments = serde_json::from_value(arguments).unwrap(); + let params = TextractAnalyzeDocumentConfig + .map_ocr_params(&arguments, MODEL) + .unwrap(); + + let request = TextractAnalyzeDocumentConfig + .transform_ocr_request(MODEL, document, ¶ms, &[]) + .unwrap(); + + assert_eq!( + serde_json::to_value(request).unwrap(), + json!({"Document": {"Bytes": "aGk="}, "FeatureTypes": expected}) + ); + } + + #[rstest] + #[case::undocumented_feature(json!({"feature_types": ["HANDWRITING"]}))] + #[case::lowercase_feature(json!({"feature_types": ["layout"]}))] + #[case::not_a_list(json!({"feature_types": "LAYOUT"}))] + fn feature_types_outside_the_documented_values_are_refused(#[case] arguments: Value) { + let arguments: CallArguments = serde_json::from_value(arguments).unwrap(); + + assert!( + TextractAnalyzeDocumentConfig + .map_ocr_params(&arguments, MODEL) + .is_err() + ); + } +} diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs new file mode 100644 index 00000000000..8268ad066a1 --- /dev/null +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs @@ -0,0 +1,678 @@ +use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_auth_aws::{SigV4Signer, resolve_aws_region}; +use litellm_http::outbound::RequestSigner; +use serde::{Deserialize, Serialize}; +use strum::{EnumString, IntoStaticStr, VariantNames}; + +use crate::base_llm::ocr::{ + document::{InlineDocument, inline_remote_document}, + error::Error, + transformation::{ + LiteLLMOcrResponse, OcrDocument, OcrEnvironment, OcrPage, OcrRequestContext, OcrUsageInfo, + PreparedOcrRequest, + }, +}; + +const TEXTRACT_SERVICE: &str = "textract"; +const AWS_JSON_CONTENT_TYPE: &str = "application/x-amz-json-1.1"; +const TARGET_HEADER: &str = "X-Amz-Target"; +const CONTENT_TYPE_HEADER: &str = "Content-Type"; +const UNSUPPORTED_DOCUMENT: &str = "UnsupportedDocumentException"; +const SYNC_DOCUMENT_MAX_BYTES: usize = 10 * 1024 * 1024; + +const HEALTH_CHECK_IMAGE_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC"; + +/// Textract has operations rather than models; the model slot of +/// `aws_textract/` names the one to call. +#[derive(Clone, Copy, Debug, EnumString, IntoStaticStr, VariantNames, PartialEq, Eq)] +#[strum(serialize_all = "kebab-case", ascii_case_insensitive)] +pub enum TextractOperation { + DetectDocumentText, + AnalyzeDocument, +} + +impl TextractOperation { + pub const PROVIDER: &'static str = "aws_textract"; + + pub fn from_model(model: &str) -> Result { + model.parse().map_err(|_| Error::InvalidModel { + provider: Self::PROVIDER, + model: model.to_string(), + supported: Self::VARIANTS, + }) + } + + fn target(self) -> &'static str { + match self { + Self::DetectDocumentText => "Textract.DetectDocumentText", + Self::AnalyzeDocument => "Textract.AnalyzeDocument", + } + } +} + +#[derive(Debug, Deserialize, Serialize)] +pub struct TextractDocument { + #[serde(rename = "Bytes")] + pub bytes: String, +} + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub enum FeatureType { + Tables, + Forms, + Queries, + Signatures, + Layout, +} + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub(super) enum BlockType { + KeyValueSet, + Page, + Line, + Word, + Table, + Cell, + SelectionElement, + MergedCell, + Title, + Query, + QueryResult, + Signature, + TableTitle, + TableFooter, + LayoutText, + LayoutTitle, + LayoutHeader, + LayoutFooter, + LayoutSectionHeader, + LayoutPageNumber, + LayoutList, + LayoutFigure, + LayoutTable, + LayoutKeyValue, + #[serde(other)] + Unknown, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum LayoutType { + Text, + Title, + Header, + Footer, + SectionHeader, + PageNumber, + List, + Figure, + Table, + KeyValue, +} + +impl BlockType { + pub fn layout(self) -> Option { + match self { + Self::LayoutText => Some(LayoutType::Text), + Self::LayoutTitle => Some(LayoutType::Title), + Self::LayoutHeader => Some(LayoutType::Header), + Self::LayoutFooter => Some(LayoutType::Footer), + Self::LayoutSectionHeader => Some(LayoutType::SectionHeader), + Self::LayoutPageNumber => Some(LayoutType::PageNumber), + Self::LayoutList => Some(LayoutType::List), + Self::LayoutFigure => Some(LayoutType::Figure), + Self::LayoutTable => Some(LayoutType::Table), + Self::LayoutKeyValue => Some(LayoutType::KeyValue), + Self::KeyValueSet + | Self::Page + | Self::Line + | Self::Word + | Self::Table + | Self::Cell + | Self::SelectionElement + | Self::MergedCell + | Self::Title + | Self::Query + | Self::QueryResult + | Self::Signature + | Self::TableTitle + | Self::TableFooter + | Self::Unknown => None, + } + } +} + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub(super) enum RelationshipType { + Value, + Child, + ComplexFeatures, + MergedCell, + Title, + Answer, + Table, + TableTitle, + TableFooter, + #[serde(other)] + Unknown, +} + +#[derive(Deserialize)] +#[serde(rename_all = "PascalCase")] +pub(super) struct Block { + #[serde(default)] + pub id: String, + pub block_type: BlockType, + pub text: Option, + pub page: Option, + pub row_index: Option, + pub column_index: Option, + #[serde(default)] + pub relationships: Vec, +} + +#[derive(Deserialize)] +#[serde(rename_all = "PascalCase")] +pub(super) struct Relationship { + pub r#type: RelationshipType, + #[serde(default)] + pub ids: Vec, +} + +impl Block { + pub fn page(&self) -> i64 { + self.page.unwrap_or(1) + } + + pub fn children(&self) -> impl Iterator { + self.relationships + .iter() + .filter(|relationship| relationship.r#type == RelationshipType::Child) + .flat_map(|relationship| relationship.ids.iter().map(String::as_str)) + } +} + +#[derive(Deserialize)] +#[serde(rename_all = "PascalCase")] +pub(super) struct DocumentMetadata { + pub pages: Option, +} + +#[derive(Deserialize)] +#[serde(rename_all = "PascalCase")] +pub struct TextractResponse { + #[serde(default)] + pub(super) blocks: Vec, + pub(super) document_metadata: Option, +} + +pub struct TextractEnvironment { + headers: Vec<(String, String)>, + region: String, + signer: SigV4Signer, +} + +impl OcrEnvironment for TextractEnvironment { + fn headers(&self) -> &[(String, String)] { + &self.headers + } + + fn signer(&self) -> Option<&dyn RequestSigner> { + Some(&self.signer) + } +} + +pub(super) fn health_check_document() -> OcrDocument { + OcrDocument::ImageUrl { + image_url: HEALTH_CHECK_IMAGE_DATA_URI.into(), + extra_fields: Default::default(), + } +} + +pub(super) async fn environment( + request: &PreparedOcrRequest, + operation: TextractOperation, +) -> Result { + let env_lookup = |name: &str| request.connection.secret(name); + let region = + resolve_aws_region(None, &request.optional_params, &env_lookup).ok_or_else(|| { + Error::InvalidRequest( + "Missing AWS region - pass aws_region_name or set AWS_REGION_NAME or AWS_REGION" + .into(), + ) + })?; + let signer = SigV4Signer::resolve( + region.clone(), + TEXTRACT_SERVICE, + &request.optional_params, + &env_lookup, + ) + .await + .map_err(litellm_auth::Error::from)?; + Ok(TextractEnvironment { + headers: operation_headers(&request.connection.extra_headers, operation), + region, + signer, + }) +} + +/// A caller's copy of an operation header would reach the wire next to ours +/// while the signature covers only one value, which Textract rejects. +fn operation_headers( + extra_headers: &[(String, String)], + operation: TextractOperation, +) -> Vec<(String, String)> { + let operation = [ + (TARGET_HEADER, operation.target()), + (CONTENT_TYPE_HEADER, AWS_JSON_CONTENT_TYPE), + ]; + extra_headers + .iter() + .filter(|(name, _)| { + !operation + .iter() + .any(|(operation_name, _)| name.eq_ignore_ascii_case(operation_name)) + }) + .cloned() + .chain( + operation + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())), + ) + .collect() +} + +pub(super) fn endpoint(request: &PreparedOcrRequest, environment: &TextractEnvironment) -> String { + request + .connection + .api_base + .clone() + .unwrap_or_else(|| format!("https://textract.{}.amazonaws.com/", environment.region)) +} + +pub(super) fn document_bytes(document: &OcrDocument) -> Result { + let inline = InlineDocument::parse(document.source())?.ok_or(Error::InvalidDataUri)?; + Ok(TextractDocument { + bytes: STANDARD.encode(inline.decode(SYNC_DOCUMENT_MAX_BYTES)?), + }) +} + +pub(super) async fn inline_document( + document: OcrDocument, + context: OcrRequestContext<'_>, +) -> Result { + inline_remote_document( + context.client.document_fetcher(), + document, + context.connection, + ) + .await +} + +#[derive(Deserialize)] +struct AwsError { + #[serde(rename = "__type", default)] + kind: String, + #[serde(rename = "Message", alias = "message", default)] + message: String, +} + +/// Textract answers both an unsupported format and a multi-page PDF or TIFF +/// with a bare "unsupported document format", which reads like a corrupt file. +/// Say what the synchronous API accepts. +pub(super) fn error_class(body: String, status: u16, headers: Vec<(String, String)>) -> Error { + let unsupported = serde_json::from_str::(&body) + .ok() + .filter(|error| error.kind.ends_with(UNSUPPORTED_DOCUMENT)); + Error::Provider { + status, + body: match unsupported { + Some(error) => format!( + "{UNSUPPORTED_DOCUMENT}: {}. aws_textract uses Textract's synchronous API, which reads a JPEG, PNG, or a single-page PDF or TIFF; other formats and multi-page documents are not supported", + error.message + ), + None => body, + }, + headers, + } +} + +pub(super) fn lines_by_page(blocks: &[Block]) -> Vec<(i64, String)> { + let pages: std::collections::BTreeSet = blocks.iter().map(Block::page).collect(); + pages + .into_iter() + .map(|page| { + let lines: Vec<&str> = blocks + .iter() + .filter(|block| block.block_type == BlockType::Line && block.page() == page) + .filter_map(|block| block.text.as_deref()) + .collect(); + (page, lines.join("\n")) + }) + .filter(|(_, markdown)| !markdown.is_empty()) + .collect() +} + +pub(super) fn ocr_response( + model: &str, + page_markdown: Vec<(i64, String)>, + document_metadata: Option, +) -> LiteLLMOcrResponse { + let pages: Vec = page_markdown + .into_iter() + .map(|(page, markdown)| OcrPage { + index: page - 1, + markdown, + ..Default::default() + }) + .collect(); + let pages_processed = document_metadata + .and_then(|metadata| metadata.pages) + .or_else(|| i64::try_from(pages.len()).ok()); + LiteLLMOcrResponse { + usage_info: Some(OcrUsageInfo { + pages_processed, + ..Default::default() + }), + ..LiteLLMOcrResponse::new(model, pages) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::{Value, json}; + + use super::*; + + const HINT: &str = "other formats and multi-page documents are not supported"; + + fn blocks(value: Value) -> Vec { + serde_json::from_value(value).unwrap() + } + + #[rstest] + #[case::detect("detect-document-text", TextractOperation::DetectDocumentText)] + #[case::analyze("analyze-document", TextractOperation::AnalyzeDocument)] + #[case::any_case("Analyze-Document", TextractOperation::AnalyzeDocument)] + fn a_model_names_its_operation(#[case] model: &str, #[case] expected: TextractOperation) { + assert_eq!(TextractOperation::from_model(model).unwrap(), expected); + } + + #[rstest] + #[case::misspelled("analyse-document")] + #[case::operation_name_from_the_api("AnalyzeDocument")] + #[case::operation_litellm_does_not_call("analyze-expense")] + #[case::empty("")] + fn a_model_outside_the_operations_is_refused_with_the_supported_names(#[case] model: &str) { + let error = TextractOperation::from_model(model).unwrap_err(); + + assert_eq!( + error.to_string(), + format!( + "invalid model: aws_textract has no model {model:?} - use one of: detect-document-text, analyze-document" + ) + ); + assert_eq!(error.http_status_code(), Some(400)); + } + + #[rstest] + #[case::line("LINE", BlockType::Line)] + #[case::key_value_set("KEY_VALUE_SET", BlockType::KeyValueSet)] + #[case::layout_section_header("LAYOUT_SECTION_HEADER", BlockType::LayoutSectionHeader)] + #[case::layout_key_value("LAYOUT_KEY_VALUE", BlockType::LayoutKeyValue)] + #[case::added_by_textract_later("LAYOUT_SIDEBAR", BlockType::Unknown)] + fn block_type_reads_the_documented_names(#[case] wire: &str, #[case] expected: BlockType) { + let block: Block = serde_json::from_value(json!({"BlockType": wire})).unwrap(); + + assert_eq!(block.block_type, expected); + } + + #[rstest] + #[case::layout_title(BlockType::LayoutTitle, Some(LayoutType::Title))] + #[case::layout_table(BlockType::LayoutTable, Some(LayoutType::Table))] + #[case::table_is_not_layout(BlockType::Table, None)] + #[case::title_is_not_layout(BlockType::Title, None)] + #[case::unknown_is_not_layout(BlockType::Unknown, None)] + fn only_layout_block_types_have_a_layout_type( + #[case] block_type: BlockType, + #[case] expected: Option, + ) { + assert_eq!(block_type.layout(), expected); + } + + #[rstest] + #[case::child_only(json!([{"Type": "CHILD", "Ids": ["a", "b"]}]), vec!["a", "b"])] + #[case::other_relationships_are_skipped( + json!([ + {"Type": "TABLE_TITLE", "Ids": ["t"]}, + {"Type": "CHILD", "Ids": ["a"]}, + {"Type": "MERGED_CELL", "Ids": ["m"]}, + {"Type": "ADDED_LATER", "Ids": ["x"]}, + {"Type": "CHILD", "Ids": ["b"]} + ]), + vec!["a", "b"] + )] + #[case::no_relationships(json!([]), vec![])] + fn children_are_the_ids_of_child_relationships( + #[case] relationships: Value, + #[case] expected: Vec<&str>, + ) { + let block: Block = + serde_json::from_value(json!({"BlockType": "LINE", "Relationships": relationships})) + .unwrap(); + + assert_eq!(block.children().collect::>(), expected); + } + + #[rstest] + #[case::tables("TABLES", Some(FeatureType::Tables))] + #[case::forms("FORMS", Some(FeatureType::Forms))] + #[case::queries("QUERIES", Some(FeatureType::Queries))] + #[case::signatures("SIGNATURES", Some(FeatureType::Signatures))] + #[case::layout("LAYOUT", Some(FeatureType::Layout))] + #[case::lowercase_is_not_a_feature("layout", None)] + #[case::undocumented("HANDWRITING", None)] + fn feature_type_accepts_only_the_documented_values( + #[case] wire: &str, + #[case] expected: Option, + ) { + assert_eq!( + serde_json::from_value::(json!(wire)).ok(), + expected + ); + if let Some(feature) = expected { + assert_eq!(serde_json::to_value(feature).unwrap(), json!(wire)); + } + } + + #[rstest] + #[case::image_url( + OcrDocument::ImageUrl { + image_url: "data:image/png;base64,aGVsbG8=".into(), + extra_fields: Default::default(), + }, + "aGVsbG8=" + )] + #[case::document_url( + OcrDocument::DocumentUrl { + document_url: "data:application/pdf;base64,YWJj".into(), + extra_fields: Default::default(), + }, + "YWJj" + )] + #[case::percent_encoded_data_uri_is_re_encoded_as_base64( + OcrDocument::DocumentUrl { + document_url: "data:,abc".into(), + extra_fields: Default::default(), + }, + "YWJj" + )] + fn document_bytes_are_the_base64_payload_without_the_data_uri_envelope( + #[case] document: OcrDocument, + #[case] expected: &str, + ) { + assert_eq!(document_bytes(&document).unwrap().bytes, expected); + } + + #[rstest] + #[case::remote_url("https://example.com/a.pdf".to_string(), Error::InvalidDataUri)] + #[case::invalid_base64("data:image/png;base64,@@@".to_string(), Error::InvalidDataUri)] + #[case::over_the_sync_limit( + format!("data:,{}", "a".repeat(SYNC_DOCUMENT_MAX_BYTES + 1)), + Error::InlineDocumentTooLarge + )] + fn document_bytes_refuse_what_the_sync_api_cannot_take( + #[case] document_url: String, + #[case] expected: Error, + ) { + let error = document_bytes(&OcrDocument::DocumentUrl { + document_url, + extra_fields: Default::default(), + }) + .unwrap_err(); + + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&expected) + ); + } + + #[rstest] + #[case::bare_type( + r#"{"__type":"UnsupportedDocumentException","Message":"Request has unsupported document format"}"#, + Some("Request has unsupported document format") + )] + #[case::namespaced_type( + r#"{"__type":"com.amazonaws.textract#UnsupportedDocumentException","Message":"bad"}"#, + Some("bad") + )] + #[case::lowercase_message( + r#"{"__type":"UnsupportedDocumentException","message":"bad"}"#, + Some("bad") + )] + #[case::other_exception(r#"{"__type":"AccessDeniedException","Message":"no"}"#, None)] + #[case::json_without_a_type(r#"{"Message":"no"}"#, None)] + #[case::not_json("bad gateway", None)] + fn only_an_unsupported_document_gains_the_sync_api_hint( + #[case] body: &str, + #[case] hinted_message: Option<&str>, + ) { + let response_headers = vec![("x-amzn-requestid".to_string(), "abc".to_string())]; + + let Error::Provider { + status, + body: reported, + headers, + } = error_class(body.into(), 400, response_headers.clone()) + else { + panic!("expected a provider error"); + }; + + assert_eq!(status, 400); + assert_eq!(headers, response_headers); + match hinted_message { + Some(message) => { + assert!(reported.contains(message), "{reported}"); + assert!(reported.contains(HINT), "{reported}"); + } + None => assert_eq!(reported, body), + } + } + + #[rstest] + #[case::no_caller_headers(vec![], vec![])] + #[case::unrelated_headers_are_kept(vec![("x-trace", "1")], vec![("x-trace", "1")])] + #[case::a_caller_content_type_is_replaced( + vec![("content-type", "application/json"), ("x-trace", "1")], + vec![("x-trace", "1")] + )] + #[case::a_caller_target_is_replaced( + vec![("X-AMZ-TARGET", "Textract.AnalyzeDocument")], + vec![] + )] + fn operation_headers_are_sent_once( + #[case] extra_headers: Vec<(&str, &str)>, + #[case] kept: Vec<(&str, &str)>, + ) { + let owned = |headers: Vec<(&str, &str)>| -> Vec<(String, String)> { + headers + .into_iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() + }; + + let headers = + operation_headers(&owned(extra_headers), TextractOperation::DetectDocumentText); + + let mut expected = owned(kept); + expected.extend(owned(vec![ + ("X-Amz-Target", "Textract.DetectDocumentText"), + ("Content-Type", "application/x-amz-json-1.1"), + ])); + assert_eq!(headers, expected); + } + + #[rstest] + #[case::words_are_not_repeated( + json!([ + {"BlockType": "PAGE"}, + {"BlockType": "LINE", "Text": "Invoice 12345"}, + {"BlockType": "WORD", "Text": "Invoice"}, + {"BlockType": "WORD", "Text": "12345"}, + {"BlockType": "LINE", "Text": "total 67.89"} + ]), + vec![(1, "Invoice 12345\ntotal 67.89")] + )] + #[case::pages_are_sorted_and_keep_line_order( + json!([ + {"BlockType": "LINE", "Text": "second", "Page": 2}, + {"BlockType": "LINE", "Text": "first", "Page": 1}, + {"BlockType": "LINE", "Text": "also second", "Page": 2} + ]), + vec![(1, "first"), (2, "second\nalso second")] + )] + #[case::a_page_without_lines_is_dropped( + json!([ + {"BlockType": "PAGE", "Page": 1}, + {"BlockType": "LINE", "Text": "only", "Page": 2} + ]), + vec![(2, "only")] + )] + #[case::no_blocks(json!([]), vec![])] + fn lines_are_grouped_by_page(#[case] input: Value, #[case] expected: Vec<(i64, &str)>) { + let pages = lines_by_page(&blocks(input)); + + let pages: Vec<(i64, &str)> = pages + .iter() + .map(|(page, markdown)| (*page, markdown.as_str())) + .collect(); + assert_eq!(pages, expected); + } + + #[rstest] + #[case::metadata_wins(Some(3), Some(3))] + #[case::metadata_without_pages_falls_back_to_the_page_count(None, Some(2))] + fn pages_are_zero_indexed_and_usage_reports_pages_processed( + #[case] metadata_pages: Option, + #[case] expected: Option, + ) { + let response = ocr_response( + "detect-document-text", + vec![(1, "first".into()), (3, "third".into())], + Some(DocumentMetadata { + pages: metadata_pages, + }), + ); + + let pages: Vec<(i64, &str)> = response + .pages + .iter() + .map(|page| (page.index, page.markdown.as_str())) + .collect(); + assert_eq!(pages, vec![(0, "first"), (2, "third")]); + assert_eq!(response.usage_info.unwrap().pages_processed, expected); + } +} diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/mod.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/mod.rs new file mode 100644 index 00000000000..ef07c1f24e1 --- /dev/null +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/mod.rs @@ -0,0 +1,3 @@ +pub mod analyze_transformation; +pub mod common_utils; +pub mod transformation; diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs new file mode 100644 index 00000000000..ad630a1ca4c --- /dev/null +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs @@ -0,0 +1,240 @@ +use litellm_core_utils::call_arguments::CallArguments; +use serde::{Deserialize, Serialize}; + +use super::common_utils::{ + TextractDocument, TextractEnvironment, TextractOperation, TextractResponse, document_bytes, + endpoint, environment, error_class, health_check_document, inline_document, lines_by_page, + ocr_response, +}; +use crate::base_llm::ocr::{ + error::Error, + handler::OcrClient, + transformation::{ + BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrRequestContext, OcrResponseFormat, + PreparedOcrRequest, decode_and_normalize_response, + }, +}; + +#[derive(Debug, Deserialize, Serialize)] +pub struct DetectDocumentTextRequest { + #[serde(rename = "Document")] + pub document: TextractDocument, +} + +#[derive(Clone, Copy, Debug, Default)] +pub struct TextractDetectTextConfig; + +impl BaseOcrConfig for TextractDetectTextConfig { + type OcrParams = (); + type ProviderRequest = DetectDocumentTextRequest; + type Environment = TextractEnvironment; + + fn get_health_check_document(&self) -> OcrDocument { + health_check_document() + } + + fn map_ocr_params( + &self, + _non_default_params: &CallArguments, + _model: &str, + ) -> Result<(), Error> { + Ok(()) + } + + async fn validate_environment( + &self, + request: &PreparedOcrRequest, + _client: &OcrClient, + ) -> Result { + environment(request, TextractOperation::DetectDocumentText).await + } + + fn get_complete_url( + &self, + request: &PreparedOcrRequest, + _optional_params: &(), + environment: &TextractEnvironment, + ) -> Result { + Ok(endpoint(request, environment)) + } + + fn transform_ocr_request( + &self, + _model: &str, + document: OcrDocument, + _optional_params: &(), + _headers: &[(String, String)], + ) -> Result { + Ok(DetectDocumentTextRequest { + document: document_bytes(&document)?, + }) + } + + async fn async_transform_ocr_request( + &self, + model: &str, + document: OcrDocument, + optional_params: &(), + headers: &[(String, String)], + context: OcrRequestContext<'_>, + ) -> Result { + let document = inline_document(document, context).await?; + self.transform_ocr_request(model, document, optional_params, headers) + } + + fn transform_ocr_response( + &self, + model: &str, + raw_response: &[u8], + request_format: OcrResponseFormat, + ) -> Result { + decode_and_normalize_response(model, raw_response, request_format, normalize_response) + } + + fn get_error_class( + &self, + error_message: String, + status_code: u16, + headers: Vec<(String, String)>, + ) -> Error { + error_class(error_message, status_code, headers) + } +} + +fn normalize_response( + model: &str, + response: TextractResponse, +) -> Result { + Ok(ocr_response( + model, + lines_by_page(&response.blocks), + response.document_metadata, + )) +} + +#[cfg(test)] +mod tests { + use rstest::{fixture, rstest}; + use serde_json::{Value, json}; + + use super::*; + + const MODEL: &str = "detect-document-text"; + + #[fixture] + fn document(#[default("data:image/png;base64,aGVsbG8=")] source: &str) -> OcrDocument { + OcrDocument::DocumentUrl { + document_url: source.into(), + extra_fields: Default::default(), + } + } + + #[rstest] + #[case::one_page_without_page_numbers( + json!({ + "DetectDocumentTextModelVersion": "1.0", + "DocumentMetadata": {"Pages": 1}, + "Blocks": [ + {"BlockType": "PAGE"}, + {"BlockType": "LINE", "Text": "Invoice 12345"}, + {"BlockType": "WORD", "Text": "Invoice"}, + {"BlockType": "WORD", "Text": "12345"}, + {"BlockType": "LINE", "Text": "total 67.89"} + ] + }), + vec![(0, "Invoice 12345\ntotal 67.89")], + Some(1) + )] + #[case::pages_out_of_order( + json!({ + "DocumentMetadata": {"Pages": 2}, + "Blocks": [ + {"BlockType": "LINE", "Text": "second", "Page": 2}, + {"BlockType": "LINE", "Text": "first", "Page": 1}, + {"BlockType": "LINE", "Text": "also second", "Page": 2} + ] + }), + vec![(0, "first"), (1, "second\nalso second")], + Some(2) + )] + #[case::missing_metadata_counts_the_pages_with_text( + json!({"Blocks": [{"BlockType": "LINE", "Text": "only"}]}), + vec![(0, "only")], + Some(1) + )] + #[case::blank_document(json!({"DocumentMetadata": {"Pages": 1}}), vec![], Some(1))] + fn response_lines_become_one_markdown_page_per_document_page( + #[case] raw_response: Value, + #[case] expected_pages: Vec<(i64, &str)>, + #[case] expected_pages_processed: Option, + ) { + let response = TextractDetectTextConfig + .transform_ocr_response( + MODEL, + &serde_json::to_vec(&raw_response).unwrap(), + OcrResponseFormat::Litellm, + ) + .unwrap(); + + let pages: Vec<(i64, &str)> = response + .pages + .iter() + .map(|page| (page.index, page.markdown.as_str())) + .collect(); + assert_eq!(pages, expected_pages); + assert_eq!(response.model, MODEL); + assert_eq!( + response.usage_info.unwrap().pages_processed, + expected_pages_processed + ); + } + + #[rstest] + fn the_request_is_only_the_document_bytes(document: OcrDocument) { + let request = TextractDetectTextConfig + .transform_ocr_request(MODEL, document, &(), &[]) + .unwrap(); + + assert_eq!( + serde_json::to_value(request).unwrap(), + json!({"Document": {"Bytes": "aGVsbG8="}}) + ); + } + + #[rstest] + fn a_remote_url_is_refused_by_the_sync_transform( + #[with("https://example.com/a.pdf")] document: OcrDocument, + ) { + let error = TextractDetectTextConfig + .transform_ocr_request(MODEL, document, &(), &[]) + .unwrap_err(); + + assert!(matches!(error, Error::InvalidDataUri)); + } + + #[rstest] + fn the_health_check_document_is_an_inline_image_the_request_accepts() { + let document = TextractDetectTextConfig.get_health_check_document(); + + assert!( + TextractDetectTextConfig + .transform_ocr_request(MODEL, document, &(), &[]) + .is_ok() + ); + } + + #[rstest] + fn provider_errors_go_through_the_shared_textract_error_class() { + let error = TextractDetectTextConfig.get_error_class( + r#"{"__type":"UnsupportedDocumentException","Message":"Request has unsupported document format"}"#.into(), + 400, + Vec::new(), + ); + + assert!( + error + .to_string() + .contains("multi-page documents are not supported") + ); + } +} diff --git a/litellm-rust/crates/llms/src/base_llm/audio_transcription/transformation.rs b/litellm-rust/crates/llms/src/base_llm/audio_transcription/transformation.rs index dd4588732be..1257bbf0d6a 100644 --- a/litellm-rust/crates/llms/src/base_llm/audio_transcription/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/audio_transcription/transformation.rs @@ -21,14 +21,7 @@ impl AudioTranscriptionResponseData { } } -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum AudioTranscriptionAuth { - Bearer, - AwsSigV4 { - region: String, - service: &'static str, - }, -} +pub use litellm_auth::RequestAuth; pub trait BaseAudioTranscriptionConfig: Sync { fn get_supported_openai_params(&self) -> &'static [&'static str]; @@ -70,5 +63,5 @@ pub trait BaseAudioTranscriptionConfig: Sync { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> Result; + ) -> Result; } diff --git a/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs b/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs index ac0450c25f0..c7d1a27c71e 100644 --- a/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs @@ -41,14 +41,7 @@ pub const STREAM_PARAM: &str = "stream"; /// presence does not make a request untranslatable. const IGNORABLE_MESSAGE_FIELDS: &[&str] = &["name"]; -/// How the upstream call is authenticated. API-key strategies are resolved in -/// `prepare`; SigV4 needs the serialized body, so the handler signs it. -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum ChatCompletionsAuth { - Header { name: &'static str, value: String }, - Bearer { token: String }, - AwsSigV4 { region: String }, -} +pub use litellm_auth::RequestAuth; /// Why a request cannot be served by the Rust path. /// @@ -91,7 +84,7 @@ pub trait BaseConfig: Sync { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> Result; + ) -> Result; fn default_headers(&self) -> &'static [(&'static str, &'static str)] { &[("content-type", "application/json")] diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs index 9fce387beb5..e09842e2856 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs @@ -76,6 +76,12 @@ pub enum Error { Unsupported(&'static str), #[error("invalid provider: {0}")] InvalidProvider(String), + #[error("invalid model: {provider} has no model {model:?} - use one of: {}", supported.join(", "))] + InvalidModel { + provider: &'static str, + model: String, + supported: &'static [&'static str], + }, #[error("invalid request: {0}")] InvalidRequest(String), #[error("invalid response: {0}")] @@ -100,6 +106,8 @@ pub enum Error { Params(#[from] litellm_core_utils::params::Error), #[error(transparent)] Headers(#[from] litellm_http::request::HeaderError), + #[error(transparent)] + Http(#[from] litellm_http::Error), } impl From for Error { @@ -153,8 +161,10 @@ impl Error { | Self::DotModel | Self::InvalidRequest(_) | Self::InvalidProvider(_) + | Self::InvalidModel { .. } | Self::Params(_) | Self::Headers(_) + | Self::Http(_) ) } diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs index 91fb6461770..245261d9f92 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs @@ -5,7 +5,7 @@ use litellm_host::event::WireRequest; use litellm_http::{ ClientVariant, HttpClientConfig, HttpClientPool, media::{MediaFetcher, UrlPolicy}, - request::{HeaderPolicy, execute_http_request, with_headers}, + outbound::{OutboundRequest, RequestSigner}, transport, }; use serde::{Serialize, de::DeserializeOwned}; @@ -117,8 +117,9 @@ pub async fn ocr( ) -> Result { let http = config.prepare_request(request, client, hooks).await?; let url = http.url().to_string(); - let headers = request_headers(&http)?; - let response = execute_http_request(client.provider_http(), http) + let headers = http.headers().to_vec(); + let response = http + .send(client.provider_http()) .await .map_err(transport_error)?; if !response.status().is_success() { @@ -153,21 +154,6 @@ pub async fn ocr( .await } -fn request_headers(request: &reqwest::Request) -> Result, Error> { - request - .headers() - .iter() - .map(|(name, value)| { - value - .to_str() - .map(|value| (name.to_string(), value.to_string())) - .map_err(|_| Error::RequestField { - path: "headers".into(), - }) - }) - .collect() -} - pub async fn read_json_response( response: reqwest::Response, native: bool, @@ -222,13 +208,13 @@ pub fn transport_error(error: reqwest::Error) -> Error { pub async fn transform_request_body( config: &C, - client: &OcrClient, request: &PreparedOcrRequest, url: &str, headers: &[(String, String)], body: B, + signer: Option<&dyn RequestSigner>, hooks: &dyn CallHooks, -) -> Result { +) -> Result { let composed = litellm_core_utils::call_arguments::compose_body( &request.optional_params, &body, @@ -244,7 +230,17 @@ pub async fn transform_request_body( }); } config.validate_request_body(&changed.body)?; - build_http_request(client, request, url, &changed.headers, &changed.body) + let timeout = Some(request.connection.timeout); + Ok(match signer { + Some(signer) => OutboundRequest::signed_json( + url.into(), + changed.headers, + &changed.body, + timeout, + signer, + ), + None => OutboundRequest::json(url.into(), changed.headers, &changed.body, timeout), + }?) } fn wire_request(url: &str, headers: &[(String, String)], body: Value) -> WireRequest { @@ -255,22 +251,18 @@ fn wire_request(url: &str, headers: &[(String, String)], body: Value) -> WireReq } } -pub fn build_http_request( - client: &OcrClient, +pub fn build_http_request( request: &PreparedOcrRequest, - url: &str, - headers: &[(String, String)], - body: &B, -) -> Result { - let builder = client - .provider_http() - .post(url) - .json(body) - .timeout(request.connection.timeout); - with_headers(builder, headers, HeaderPolicy::All) - .build() - .map_err(transport::Error::from) - .map_err(Error::from) + url: String, + headers: Vec<(String, String)>, + body: &impl Serialize, +) -> Result { + Ok(OutboundRequest::json( + url, + headers, + body, + Some(request.connection.timeout), + )?) } pub async fn guardrail_document( diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index 3960282b580..e02a4b7f266 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -6,6 +6,7 @@ use litellm_core_utils::{ serde_compat::{FiniteF64, LaxI64}, settings::ProcessEnvironment, }; +use litellm_http::outbound::{OutboundRequest, RequestSigner}; use serde::{ Deserialize, Serialize, de::{DeserializeOwned, IntoDeserializer}, @@ -394,6 +395,10 @@ const HEALTH_CHECK_PDF_DATA_URI: &str = "data:application/pdf;base64,JVBERi0xLjQ /// (headers at minimum; Vertex also carries the project id). pub trait OcrEnvironment: Send + Sync { fn headers(&self) -> &[(String, String)]; + + fn signer(&self) -> Option<&dyn RequestSigner> { + None + } } impl OcrEnvironment for Vec<(String, String)> { @@ -536,7 +541,7 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static { request: &PreparedOcrRequest, client: &OcrClient, hooks: &dyn CallHooks, - ) -> impl Future> + Send { + ) -> impl Future> + Send { async move { let params = self.map_ocr_params(&request.optional_params, &request.model)?; let environment = self.validate_environment(request, client).await?; @@ -554,7 +559,16 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static { }, ) .await?; - transform_request_body(self, client, request, &url, headers, body, hooks).await + transform_request_body( + self, + request, + &url, + headers, + body, + environment.signer(), + hooks, + ) + .await } } } diff --git a/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs b/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs index 39734d844da..cfabcb12341 100644 --- a/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs +++ b/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs @@ -8,8 +8,8 @@ use serde_json::{Map, Value, json}; use crate::base_llm::{ audio_transcription::transformation::{ - AudioTranscriptionAuth, AudioTranscriptionRequestData, AudioTranscriptionResponseData, - BaseAudioTranscriptionConfig, + AudioTranscriptionRequestData, AudioTranscriptionResponseData, + BaseAudioTranscriptionConfig, RequestAuth, }, chat::transformation::Error, }; @@ -136,9 +136,9 @@ impl BaseAudioTranscriptionConfig for BedrockAudioTranscriptionConfig { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { + ) -> Result { let (_, model_region) = bedrock_model_id_and_region(model); - Ok(AudioTranscriptionAuth::AwsSigV4 { + Ok(RequestAuth::AwsSigV4 { region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup), service: BEDROCK_SERVICE, }) diff --git a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs index 23c6c5c61bd..09c456f1d0a 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs @@ -1,6 +1,6 @@ use litellm_auth_aws::{ bedrock_model_id_and_region, - constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE}, + constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}, resolve_bedrock_region, }; use litellm_core_utils::{ @@ -17,8 +17,8 @@ use litellm_types::{ use serde_json::{Map, Value, json}; use crate::base_llm::chat::transformation::{ - BaseConfig, ChatCompletionsAuth, Error, ProviderChatRequestData, ProviderChatResponseData, - Unsupported, unsupported_message, unsupported_param, + BaseConfig, Error, ProviderChatRequestData, ProviderChatResponseData, RequestAuth, Unsupported, + unsupported_message, unsupported_param, }; /// Converse parameter names, post `map_openai_params`, that the Rust path can @@ -186,7 +186,7 @@ impl BaseConfig for AmazonConverseConfig { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { + ) -> Result { // Python reads `api_key` as the Bedrock bearer token and consults the // env only when the caller passed none, so a caller-supplied empty key // falls through to SigV4 without reaching for the environment. An @@ -199,11 +199,12 @@ impl BaseConfig for AmazonConverseConfig { } .filter(|token| !token.is_empty()); if let Some(token) = bearer { - return Ok(ChatCompletionsAuth::Bearer { token }); + return Ok(RequestAuth::Bearer { token }); } let (_, model_region) = bedrock_model_id_and_region(model); - Ok(ChatCompletionsAuth::AwsSigV4 { + Ok(RequestAuth::AwsSigV4 { region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup), + service: BEDROCK_SERVICE, }) } diff --git a/litellm-rust/crates/llms/src/bedrock/chat/tests.rs b/litellm-rust/crates/llms/src/bedrock/chat/tests.rs index cca5cbda41a..d7ecde47c6b 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/tests.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/tests.rs @@ -281,8 +281,9 @@ fn signs_with_sigv4_in_the_resolved_region() { &|_| None ) .expect("auth resolves"), - ChatCompletionsAuth::AwsSigV4 { - region: "eu-central-1".to_string() + RequestAuth::AwsSigV4 { + region: "eu-central-1".to_string(), + service: "bedrock", } ); } @@ -306,11 +307,12 @@ fn a_bearer_token_outranks_sigv4_the_way_python_resolves_it() { ) .expect("auth resolves") }; - let bearer = |token: &str| ChatCompletionsAuth::Bearer { + let bearer = |token: &str| RequestAuth::Bearer { token: token.to_string(), }; - let sigv4 = ChatCompletionsAuth::AwsSigV4 { + let sigv4 = RequestAuth::AwsSigV4 { region: "eu-central-1".to_string(), + service: "bedrock", }; // A caller-supplied key is the bearer token, and outranks the env. diff --git a/litellm-rust/crates/llms/src/lib.rs b/litellm-rust/crates/llms/src/lib.rs index 8d1bb366ed4..701eaff4374 100644 --- a/litellm-rust/crates/llms/src/lib.rs +++ b/litellm-rust/crates/llms/src/lib.rs @@ -1,4 +1,5 @@ pub mod anthropic; +pub mod aws_textract; pub mod azure_ai; pub mod base_llm; pub mod bedrock; diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index 307ba697316..5272be97c24 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -5,6 +5,7 @@ use litellm_core_utils::{ params::OpaqueParams, url_utils::ApiUrl, }; +use litellm_http::outbound::OutboundRequest; use serde::{Deserialize, Deserializer, Serialize}; use serde_json::{Map, Value, json}; @@ -166,7 +167,7 @@ impl BaseOcrConfig for ReductoParseV3Config { request: &PreparedOcrRequest, client: &OcrClient, hooks: &dyn CallHooks, - ) -> Result { + ) -> Result { prepare_upload_request(self, request, client, hooks).await } } @@ -251,7 +252,7 @@ impl BaseOcrConfig for ReductoParseLegacyConfig { request: &PreparedOcrRequest, client: &OcrClient, hooks: &dyn CallHooks, - ) -> Result { + ) -> Result { prepare_upload_request(self, request, client, hooks).await } } @@ -264,7 +265,7 @@ async fn prepare_upload_request, -) -> Result { +) -> Result { let params = config.map_ocr_params(&request.optional_params, &request.model)?; let headers = config.validate_environment(request, client).await?; let url = config.get_complete_url(request, ¶ms, &headers)?; @@ -286,7 +287,7 @@ async fn prepare_upload_request Result { diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs index f0b035621fa..9a23deefb89 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs @@ -19,6 +19,13 @@ const MODEL_PREFIX: &str = "deepseek-ai/"; const DEFAULT_LOCATION: &str = "us-central1"; const DEEPSEEK_OCR_PARAMS: &[&str] = &["stream", "temperature", "max_tokens", "top_p", "n", "stop"]; +/// DeepSeek-OCR is a transcription model: at the endpoint's default sampling temperature it +/// hallucinates extra text, so requests are greedy unless the caller sets a temperature. +const DEFAULT_TEMPERATURE: f64 = 0.0; +/// Greedy decoding on dense screenshots falls into repetition loops that run to the token limit; +/// a mild penalty breaks them without changing clean-document output. +const DEFAULT_REPETITION_PENALTY: f64 = 1.05; + pub type DeepSeekOcrParams = OpaqueParams; #[derive(Clone, Debug, Serialize, Deserialize)] @@ -171,11 +178,19 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig { image_url: document.source().to_string(), }], }], - params: optional_params - .iter() - .filter(|(name, _)| DEEPSEEK_OCR_PARAMS.contains(&name.as_str())) - .map(|(name, value)| (name.clone(), value.clone())) - .collect(), + params: [ + ("temperature", DEFAULT_TEMPERATURE), + ("repetition_penalty", DEFAULT_REPETITION_PENALTY), + ] + .into_iter() + .map(|(name, value)| (name.to_string(), Value::from(value))) + .chain( + optional_params + .iter() + .filter(|(name, _)| DEEPSEEK_OCR_PARAMS.contains(&name.as_str())) + .map(|(name, value)| (name.clone(), value.clone())), + ) + .collect(), }) } } @@ -484,6 +499,46 @@ mod tests { assert!(result.get("ignored").is_none()); } + #[test] + fn request_uses_greedy_defaults_unless_the_caller_overrides_them() { + let request = |params: DeepSeekOcrParams| { + serde_json::to_value( + VertexAIDeepSeekOCRConfig + .transform_ocr_request( + "deepseek-ai/deepseek-ocr-maas", + document(), + ¶ms, + &[], + ) + .unwrap(), + ) + .unwrap() + }; + let defaults = request(DeepSeekOcrParams::default()); + assert_eq!(defaults["temperature"], 0.0); + assert_eq!(defaults["repetition_penalty"], 1.05); + assert_eq!( + request(serde_json::from_value(json!({"temperature":0.7})).unwrap())["temperature"], + 0.7 + ); + } + + #[test] + fn caller_temperature_argument_overrides_the_greedy_default_in_the_composed_body() { + let arguments = serde_json::from_value(json!({"temperature":0.7})).unwrap(); + let body = VertexAIDeepSeekOCRConfig + .transform_ocr_request( + "deepseek-ai/deepseek-ocr-maas", + document(), + &DeepSeekOcrParams::default(), + &[], + ) + .unwrap(); + let composed = + litellm_core_utils::call_arguments::compose_body(&arguments, &body, &[]).unwrap(); + assert_eq!(composed["temperature"], 0.7); + } + #[rstest] #[case(json!({"type":"image_url","image_url":"data:image/png;base64,AA=="}))] #[case(json!({"type":"document_url","document_url":"data:application/pdf;base64,AA=="}))] diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 19d28f76b6f..6c5a65173e3 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -58,6 +58,7 @@ pub(crate) fn core_error_to_pyerr(error: Error) -> PyErr { audio_transcription::Error::InvalidProvider(_) | audio_transcription::Error::InvalidRequest(_) | audio_transcription::Error::Headers(_) + | audio_transcription::Error::Http(_) | audio_transcription::Error::InvalidType { .. } | audio_transcription::Error::MissingField(_) | audio_transcription::Error::Aws(_) => true, @@ -68,6 +69,7 @@ pub(crate) fn core_error_to_pyerr(error: Error) -> PyErr { chat_completions::Error::InvalidProvider(_) | chat_completions::Error::InvalidRequest(_) | chat_completions::Error::Headers(_) + | chat_completions::Error::Http(_) | chat_completions::Error::InvalidType { .. } | chat_completions::Error::MissingField(_) | chat_completions::Error::Aws(_) => true, @@ -105,6 +107,7 @@ pub(crate) fn chat_completions_error_to_pyerr(error: chat_completions::Error) -> | Error::InvalidType { .. } | Error::MissingField(_) | Error::Headers(_) + | Error::Http(_) | Error::Transport(TransportError::Connect(_)) => { RustBridgeDeclined::new_err(error.to_string()) } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs index ed840dec70c..a928e62d5b7 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs @@ -7,7 +7,8 @@ use pyo3::{ gc::{PyTraverseError, PyVisit}, prelude::*, pybacked::PyBackedBytes, - types::{PyBytes, PyString}, + sync::PyOnceLock, + types::{PyBytes, PyString, PyType}, }; #[derive(Debug)] @@ -84,7 +85,8 @@ impl FromPyObject<'_, '_> for FileDocumentInput { "OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.", )); } - if file.is_instance(&py.import("os")?.getattr("PathLike")?)? { + static PATH_LIKE: PyOnceLock> = PyOnceLock::new(); + if file.is_instance(PATH_LIKE.import(py, "os", "PathLike")?)? { return Ok(Self { input: OcrDocumentInput::Path { path: file.extract::()?, diff --git a/litellm/experimental_mcp_client/Readme.md b/litellm/experimental_mcp_client/Readme.md index 4fbd624369c..0c7b0aa76b9 100644 --- a/litellm/experimental_mcp_client/Readme.md +++ b/litellm/experimental_mcp_client/Readme.md @@ -1,6 +1,17 @@ # LiteLLM MCP Client -LiteLLM MCP Client is a client that allows you to use MCP tools with LiteLLM. +LiteLLM MCP Client allows you to use MCP tools with LiteLLM +## MCP Python SDK compatibility +The `mcp` and `proxy` extras require MCP Python SDK 2.2 or newer within the 2.x release line. Installing core LiteLLM without these extras does not require MCP +Existing MCP SDK1 clients can continue connecting to the gateway over the supported legacy MCP protocols. The client and gateway can use different SDK versions in separate Python environments. Modern protocol advertisement remains disabled during the Phase 0 upgrade. An initialize body requesting `2026-07-28` falls back to the supported legacy version `2025-11-25`; an explicit `MCP-Protocol-Version: 2026-07-28` HTTP header is rejected with HTTP 400 + +Code sharing the gateway's Python environment must support SDK2. Its Python API has breaking changes, including renamed imports and snake_case model attributes such as `input_schema`, `is_error`, and `structured_content`. This also applies to callers consuming SDK objects returned by LiteLLM's experimental MCP client. MCP JSON fields retain their protocol spelling, such as `inputSchema` and `isError` + +Upgrade SDK1-dependent libraries before installing them alongside `litellm[mcp]` or `litellm[proxy]`, or keep those clients in a separate environment and connect over the network. For example, `langchain-mcp-adapters==0.2.1` uses SDK1 Python APIs and is tested as a separate legacy client, not as a shared SDK2 dependency + +The shared unit-test workflow runs the MCP integration suite once, with SDK2 in the gateway environment and an isolated SDK1 peer. Keep the SDK1 list/call compatibility test while SDK1 clients are supported; remove it when that support is explicitly retired and the client migration is documented + +See the official [SDK migration guide](https://py.sdk.modelcontextprotocol.io/migration/) for Python API changes diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 56ee5f30d02..4b456710057 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -9,57 +9,30 @@ import json import os from collections.abc import Awaitable, Callable, Generator from contextlib import AbstractAsyncContextManager -from datetime import timedelta from functools import partial -from importlib import metadata from types import MappingProxyType -from typing import Any, Final, Protocol, TypeAlias, TypeVar +from typing import Any, Final, TypeAlias, TypeVar -import httpx -from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream -from mcp import ClientSession, McpError, ReadResourceResult, Resource, StdioServerParameters +import httpx2 +from httpx2._client import UseClientDefault +from httpx2._types import AuthTypes +from mcp import ClientSession, MCPError, ReadResourceResult, Resource, StdioServerParameters from mcp.client.sse import sse_client from mcp.client.stdio import stdio_client +from mcp.client.streamable_http import streamable_http_client +from mcp.shared._stream_protocols import ReadStream, WriteStream from mcp.shared.message import SessionMessage -from mcp.shared.session import RequestResponder -from typing_extensions import Unpack _TransportStreams: TypeAlias = tuple[ - MemoryObjectReceiveStream[SessionMessage | Exception], - MemoryObjectSendStream[SessionMessage], - Unpack[tuple[object, ...]], + ReadStream[SessionMessage | Exception], + WriteStream[SessionMessage], ] _TransportContext: TypeAlias = AbstractAsyncContextManager[_TransportStreams] -class _StreamableHttpClientFactory(Protocol): - """The ``streamable_http_client`` entry point this module calls on the installed MCP SDK.""" - - def __call__(self, *, url: str, http_client: httpx.AsyncClient | None) -> _TransportContext: ... - - -streamable_http_client: _StreamableHttpClientFactory | None = None -try: - import mcp.client.streamable_http as streamable_http_module - - streamable_http_client = getattr(streamable_http_module, "streamable_http_client", None) -except ImportError: - pass - -MCP_STREAMABLE_HTTP_REQUIREMENT: Final = "mcp>=1.28.1" - - -def missing_streamable_http_client_error() -> ImportError: - return ImportError( - f"MCP streamable HTTP transport requires {MCP_STREAMABLE_HTTP_REQUIREMENT}, but the installed " - f"mcp {metadata.version('mcp')} does not provide streamable_http_client. " - "Fix with: pip install 'litellm[mcp]' (or upgrade mcp directly: pip install -U mcp)" - ) - - from mcp.types import ( METHOD_NOT_FOUND, - ClientResult, + REQUEST_TIMEOUT, GetPromptRequestParams, GetPromptResult, ListPromptsResult, @@ -68,7 +41,6 @@ from mcp.types import ( Prompt, ResourceTemplate, ServerNotification, - ServerRequest, TextContent, ) from mcp.types import CallToolRequestParams as MCPCallToolRequestParams @@ -153,23 +125,21 @@ def _first_non_cancelled_cause(exc: BaseException) -> BaseException | None: return None -_SDK_READ_TIMEOUT_CODE: Final = int(httpx.codes.REQUEST_TIMEOUT) -"""The code the MCP SDK puts on its own elapsed read timeout, an HTTP status in a field that -otherwise carries JSON-RPC error codes.""" +_SDK_READ_TIMEOUT_CODE: Final = REQUEST_TIMEOUT +"""The code the MCP SDK puts on its own elapsed read timeout.""" def as_mcp_read_timeout(exc: BaseException) -> TimeoutError | None: """Normalize an MCP SDK read timeout for client and gateway diagnostics, or return ``None``. - The SDK reports its own elapsed read timeout as ``McpError`` carrying an HTTP status code in a - field that otherwise holds JSON-RPC error codes, and it relays an upstream's JSON-RPC error - through that same class and field. The numeric code alone therefore cannot separate the two, and - an upstream answering with application code 408 would be reported as a gateway timeout it never - caused. The SDK raises its own from inside an ``except TimeoutError``, so the elapsed timeout is + The SDK reports its own elapsed read timeout as ``MCPError`` carrying ``REQUEST_TIMEOUT`` in a + field that also carries relayed upstream JSON-RPC errors. The numeric code alone therefore + cannot separate the two, and an upstream answering with the same application code would be + reported as a gateway timeout it never caused. The SDK raises its own from inside an ``except TimeoutError``, so the elapsed timeout is on the context chain, while a relayed error is built from a received message and has no such chain; that is the discriminator. """ - if not isinstance(exc, McpError) or exc.error.code != _SDK_READ_TIMEOUT_CODE: + if not isinstance(exc, MCPError) or exc.error.code != _SDK_READ_TIMEOUT_CODE: return None if not isinstance(exc.__context__, TimeoutError): return None @@ -179,9 +149,25 @@ def as_mcp_read_timeout(exc: BaseException) -> TimeoutError | None: TSessionResult = TypeVar("TSessionResult") -class MCPSigV4Auth(httpx.Auth): +class _MCPHTTPClient(httpx2.AsyncClient): + async def send( + self, + request: httpx2.Request, + *, + stream: bool = False, + auth: AuthTypes | UseClientDefault | None = httpx2.USE_CLIENT_DEFAULT, + follow_redirects: bool | UseClientDefault = httpx2.USE_CLIENT_DEFAULT, + ) -> httpx2.Response: + response: Final = await super().send(request, stream=stream, auth=auth, follow_redirects=follow_redirects) + if request.method == "POST" and response.is_error and response.status_code != 404: + await response.aclose() + response.raise_for_status() + return response + + +class MCPSigV4Auth(httpx2.Auth): """ - httpx Auth class that signs each request with AWS SigV4. + httpx2 Auth class that signs each request with AWS SigV4. This is used for MCP servers that require AWS SigV4 authentication, such as AWS Bedrock AgentCore MCP servers. httpx calls auth_flow() for every outgoing request, enabling per-request signature computation. @@ -270,7 +256,7 @@ class MCPSigV4Auth(httpx.Auth): token=sts_creds["SessionToken"], ) - def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: + def auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]: from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest @@ -314,8 +300,8 @@ class MCPClient: stdio_config: MCPStdioConfig | None = None, extra_headers: dict[str, str] | None = None, ssl_verify: VerifyTypes | None = None, - aws_auth: httpx.Auth | None = None, - resolved_auth: httpx.Auth | None = None, + aws_auth: httpx2.Auth | None = None, + resolved_auth: httpx2.Auth | None = None, sampling_callback: Callable | None = None, elicitation_callback: Callable | None = None, logging_callback: Callable | None = None, @@ -333,10 +319,10 @@ class MCPClient: self.stdio_config: MCPStdioConfig | None = stdio_config self.extra_headers: dict[str, str] | None = extra_headers self.ssl_verify: VerifyTypes | None = ssl_verify - self._aws_auth: httpx.Auth | None = aws_auth - # A pre-resolved httpx.Auth (e.g. from the v2 credential resolver) attached to the + self._aws_auth: httpx2.Auth | None = aws_auth + # A pre-resolved httpx2.Auth (e.g. from the v2 credential resolver) attached to the # upstream client's auth= slot, taking precedence over the SigV4 aws_auth. - self._resolved_auth: httpx.Auth | None = resolved_auth + self._resolved_auth: httpx2.Auth | None = resolved_auth self._last_initialize_instructions: str | None = None self._sampling_callback: Callable | None = sampling_callback self._elicitation_callback: Callable | None = elicitation_callback @@ -348,9 +334,11 @@ class MCPClient: async def discovery_auth_fingerprint(self) -> str: return self._hash_discovery_auth(await self.prepare_request_auth()) - async def prepare_request_auth(self) -> httpx.Request: + async def prepare_request_auth(self) -> httpx2.Request: """Preview the authenticated request without sending it, closing the auth flow afterwards.""" - request: Final = httpx.Request("POST", self.server_url or "http://localhost/", headers=self._get_auth_headers()) + request: Final = httpx2.Request( + "POST", self.server_url or "http://localhost/", headers=self._get_auth_headers() + ) if self._resolved_auth is None: return request flow: Final = self._resolved_auth.async_auth_flow(request) @@ -361,20 +349,20 @@ class MCPClient: await flow.aclose() @staticmethod - def _hash_discovery_auth(request: httpx.Request) -> str: + def _hash_discovery_auth(request: httpx2.Request) -> str: material: Final = json.dumps((str(request.url), tuple(sorted(request.headers.multi_items())))) return hashlib.sha256(material.encode()).hexdigest() def _create_transport_context( self, - ) -> tuple[_TransportContext, httpx.AsyncClient | None]: + ) -> tuple[_TransportContext, httpx2.AsyncClient | None]: """ Create the appropriate transport context based on transport type. Returns: Tuple of (transport_context, http_client). http_client is only set for HTTP transport and needs cleanup. """ - http_client: httpx.AsyncClient | None = None + http_client: httpx2.AsyncClient | None = None if self.transport_type == MCPTransport.stdio: if not self.stdio_config: raise ValueError("stdio_config is required for stdio transport") @@ -397,14 +385,12 @@ class MCPClient: None, ) # HTTP transport (default) - if streamable_http_client is None: - raise missing_streamable_http_client_error() headers = self._get_auth_headers() httpx_client_factory = self._create_httpx_client_factory() verbose_logger.debug("litellm headers for streamable_http_client: %s", headers) http_client = httpx_client_factory( headers=headers, - timeout=httpx.Timeout(self.timeout), + timeout=httpx2.Timeout(self.timeout), ) transport_ctx: Final = streamable_http_client( url=self.server_url, @@ -473,13 +459,14 @@ class MCPClient: transport: Final = await transport_ctx.__aenter__() in_flight_error: BaseException | None = None try: - read_stream, write_stream = transport[0], transport[1] + read_stream: Final = transport[0] + write_stream: Final = transport[1] stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future() async def receive_message( - message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception, + message: ServerNotification | Exception, ) -> None: - if not isinstance(message, (ValueError, httpx.RequestError, OSError)): + if not isinstance(message, (ValueError, httpx2.HTTPError, OSError)): return if not stream_error.done(): stream_error.set_result(message) @@ -499,7 +486,7 @@ class MCPClient: session_ctx: Final = ClientSession( read_stream, write_stream, - read_timeout_seconds=timedelta(seconds=self.timeout), + read_timeout_seconds=self.timeout, message_handler=receive_message, **session_kwargs, ) @@ -512,7 +499,7 @@ class MCPClient: if isinstance(ins, str) and ins.strip(): self._last_initialize_instructions = ins.strip() return await operation(session) - except McpError: + except MCPError: if stream_error.done(): raise stream_error.result() raise @@ -544,7 +531,7 @@ class MCPClient: quiet_on_error demotes the failure line to debug for callers that own the exception (call_tool / list_tools under raise_on_error), so an expected pass-through re-auth does not emit a warning per call; every other caller keeps the operator-visible warning.""" - http_client: httpx.AsyncClient | None = None + http_client: httpx2.AsyncClient | None = None try: self._last_initialize_instructions = None transport_ctx, http_client = self._create_transport_context() @@ -609,7 +596,7 @@ class MCPClient: elif isinstance(self._mcp_auth_value, dict): headers.update(self._mcp_auth_value) # Note: aws_sigv4 auth is not handled here — SigV4 requires per-request - # signing (including the body hash), so it uses httpx.Auth flow instead + # signing (including the body hash), so it uses httpx2.Auth flow instead # of static headers. See MCPSigV4Auth and _create_httpx_client_factory(). # update the headers with the extra headers if self.extra_headers: @@ -623,9 +610,11 @@ class MCPClient: headers.update(injected or {}) return _strip_header_whitespace(headers) - def _create_httpx_client_factory(self) -> Callable[..., httpx.AsyncClient]: + def _create_httpx_client_factory( + self, *, transport: httpx2.AsyncBaseTransport | None = None + ) -> Callable[..., httpx2.AsyncClient]: """ - Create a custom httpx client factory that uses LiteLLM's SSL configuration. + Create a custom httpx2 client factory that uses LiteLLM's SSL configuration. This factory follows the same CA bundle path logic as http_handler.py: 1. Check ssl_verify parameter (can be SSLContext, bool, or path to CA bundle) 2. Check SSL_VERIFY environment variable @@ -636,10 +625,10 @@ class MCPClient: def factory( *, headers: dict[str, str] | None = None, - timeout: httpx.Timeout | None = None, - auth: httpx.Auth | None = None, - ) -> httpx.AsyncClient: - """Create an httpx.AsyncClient with LiteLLM's SSL configuration.""" + timeout: httpx2.Timeout | None = None, + auth: httpx2.Auth | None = None, + ) -> httpx2.AsyncClient: + """Create an httpx2.AsyncClient with LiteLLM's SSL configuration.""" # Get unified SSL configuration using the same logic as http_handler.py ssl_config: Final = get_ssl_configuration(self.ssl_verify) verbose_logger.debug("MCP client using SSL configuration: %s", type(ssl_config).__name__) @@ -649,7 +638,8 @@ class MCPClient: fallback_auth: Final = self._resolved_auth if self._resolved_auth is not None else self._aws_auth effective_auth: Final = auth if auth is not None else fallback_auth guard: Final = credential_redirect_hook(self.server_url, self._credential_slot) - return httpx.AsyncClient( + return _MCPHTTPClient( + transport=transport, headers=headers, timeout=timeout, auth=effective_auth, @@ -723,7 +713,7 @@ class MCPClient: """The error result ``call_tool`` returns when it swallows a failure (no re-execution).""" return MCPCallToolResult( content=[TextContent(type="text", text=f"{type(exc).__name__}: {exc}")], - isError=True, + is_error=True, ) async def call_tool( @@ -808,12 +798,12 @@ class MCPClient: verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio") async def _list_prompts_operation(session: ClientSession) -> ListPromptsResult: - capabilities: Final = session.get_server_capabilities() + capabilities: Final = session.server_capabilities if capabilities is not None and capabilities.prompts is None: return ListPromptsResult(prompts=[]) try: return await session.list_prompts() - except McpError as error: + except MCPError as error: if error.error.code != METHOD_NOT_FOUND: raise verbose_logger.debug( @@ -898,12 +888,12 @@ class MCPClient: verbose_logger.debug("MCP client listing resources from %s", self.server_url or "stdio") async def _list_resources_operation(session: ClientSession) -> ListResourcesResult: - capabilities: Final = session.get_server_capabilities() + capabilities: Final = session.server_capabilities if capabilities is not None and capabilities.resources is None: return ListResourcesResult(resources=[]) try: return await session.list_resources() - except McpError as error: + except MCPError as error: if error.error.code != METHOD_NOT_FOUND: raise verbose_logger.debug( @@ -947,30 +937,30 @@ class MCPClient: verbose_logger.debug("MCP client listing resource templates from %s", self.server_url or "stdio") async def _list_resource_templates_operation(session: ClientSession) -> ListResourceTemplatesResult: - capabilities: Final = session.get_server_capabilities() + capabilities: Final = session.server_capabilities if capabilities is not None and capabilities.resources is None: - return ListResourceTemplatesResult(resourceTemplates=[]) + return ListResourceTemplatesResult(resource_templates=[]) # mutable-ok: MCP result payload try: return await session.list_resource_templates() - except McpError as error: + except MCPError as error: if error.error.code != METHOD_NOT_FOUND: raise verbose_logger.debug( "MCP client list_resource_templates is unsupported by %s: %s", self.server_url or "stdio", error ) - return ListResourceTemplatesResult(resourceTemplates=[]) + return ListResourceTemplatesResult(resource_templates=[]) # mutable-ok: MCP result payload try: result: Final = await self.run_with_session(_list_resource_templates_operation) - resource_template_count: Final = len(result.resourceTemplates) - resource_template_names: Final = [resourceTemplate.name for resourceTemplate in result.resourceTemplates] + resource_template_count: Final = len(result.resource_templates) + resource_template_names: Final = [resource_template.name for resource_template in result.resource_templates] verbose_logger.info( "MCP client listed %s resource templates from %s: %s", resource_template_count, self.server_url or "stdio", resource_template_names, ) - return result.resourceTemplates + return result.resource_templates except asyncio.CancelledError: verbose_logger.warning("MCP client list_resource_templates was cancelled") raise @@ -1000,7 +990,7 @@ class MCPClient: async def _read_resource_operation(session: ClientSession): verbose_logger.debug("MCP client sending read_resource request to session") - return await session.read_resource(url) + return await session.read_resource(str(url)) try: read_resource_result: Final = await self.run_with_session(_read_resource_operation) diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py index 51d2139ef3b..a9ee851d529 100644 --- a/litellm/experimental_mcp_client/tools.py +++ b/litellm/experimental_mcp_client/tools.py @@ -26,7 +26,7 @@ from litellm.types.utils import ChatCompletionMessageToolCall ######################################################## def transform_mcp_tool_to_openai_tool(mcp_tool: MCPTool) -> ChatCompletionToolParam: """Convert an MCP tool to an OpenAI tool.""" - normalized_parameters: Final = _normalize_mcp_input_schema(mcp_tool.inputSchema) + normalized_parameters: Final = _normalize_mcp_input_schema(mcp_tool.input_schema) return ChatCompletionToolParam( type="function", @@ -73,7 +73,7 @@ def transform_mcp_tool_to_openai_responses_api_tool( mcp_tool: MCPTool, ) -> FunctionToolParam: """Convert an MCP tool to an OpenAI Responses API tool.""" - normalized_parameters: Final = _normalize_mcp_input_schema(mcp_tool.inputSchema) + normalized_parameters: Final = _normalize_mcp_input_schema(mcp_tool.input_schema) return FunctionToolParam( name=mcp_tool.name, @@ -93,7 +93,7 @@ def transform_mcp_tool_to_anthropic_tool(mcp_tool: MCPTool) -> AnthropicMessages return AnthropicMessagesTool( name=mcp_tool.name, description=mcp_tool.description or "", - input_schema=sanitize_input_schema_for_anthropic(mcp_tool.inputSchema), + input_schema=sanitize_input_schema_for_anthropic(mcp_tool.input_schema), type="custom", ) @@ -129,7 +129,7 @@ async def list_tools_with_pagination( ) tools.extend(result.tools) - next_cursor = getattr(result, "nextCursor", None) + next_cursor = getattr(result, "next_cursor", None) if not isinstance(next_cursor, str) or not next_cursor: return tools if next_cursor in seen_cursors: diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 4f9b18713d0..494d9e0935a 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -25,7 +25,10 @@ from litellm.integrations.prompt_management_base import PromptManagementClient from litellm.litellm_core_utils.prompt_templates.common_utils import ( with_prompt_cache_breakpoint, ) -from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request +from litellm.llms.anthropic.common_utils import ( + is_claude_code_one_shot_subagent_request, + supports_anthropic_cache_control, +) from litellm.types.integrations.anthropic_cache_control_hook import ( GATEWAY_INJECTED_CACHE_METADATA_KEY, GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT, @@ -574,8 +577,9 @@ class AnthropicCacheControlHook(CustomPromptManagement): custom_llm_provider: str | None, api_base: object, prompt_cache_options: object, + request_kwargs: object, ) -> Sequence[Mapping[str, object]] | None: - if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools, cache_control): + if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools, cache_control, request_kwargs): return None return AnthropicCacheControlHook._stamped_with_dialect( points, model, custom_llm_provider, api_base, prompt_cache_options @@ -612,6 +616,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): system: str | list | None, tools: list | None, cache_control: object = None, + request_kwargs: object = None, ) -> bool: """Whether configured injection points must yield to client-set cache_control. @@ -624,7 +629,9 @@ class AnthropicCacheControlHook(CustomPromptManagement): """ if all(point.get("_litellm_judged") for point in points): return False - return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control) + return AnthropicCacheControlHook._request_has_cache_control( + messages, system, tools, cache_control, request_kwargs + ) @staticmethod def _request_has_cache_control( @@ -632,31 +639,29 @@ class AnthropicCacheControlHook(CustomPromptManagement): system: str | list | None, tools: list | None = None, cache_control: object = None, + request_kwargs: object = None, ) -> bool: - """Return True if the request already carries any client-supplied cache_control. - - When the client (e.g. Claude Code) already marks its own breakpoints we - stand down entirely rather than add more, per the auto-caching contract. - Tools count: they are a breakpoint the client can mark, they count toward - the provider's four-block limit, and caching only the tool definitions is - a common pattern, so injecting alongside them can exceed the cap. Tools - carry the mark either at the top level (Anthropic shape) or nested under - ``function`` (OpenAI shape); the Anthropic chat transform accepts both. - """ - if cache_control is not None: - return True - if AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) > 0: - return True - if tools is not None: - return any( - isinstance(tool, dict) - and ( - tool.get("cache_control") is not None - or (isinstance(tool.get("function"), dict) and tool["function"].get("cache_control") is not None) - ) - for tool in tools + """Client breakpoints own caching in both the request and its extra_body envelope.""" + bodies: Final = ( + {"messages": messages, "system": system, "tools": tools, "cache_control": cache_control}, + _validated_object_mapping(AnthropicCacheControlHook._request_value(request_kwargs, "extra_body")) or {}, + ) + return any( + body.get("cache_control") is not None + or AnthropicCacheControlHook.count_request_cache_breakpoints( + _validated_object_list(body.get("messages")) or (), body.get("system") ) - return False + > 0 + or any( + AnthropicCacheControlHook._request_value(tool, "cache_control") is not None + or AnthropicCacheControlHook._request_value( + AnthropicCacheControlHook._request_value(tool, "function"), "cache_control" + ) + is not None + for tool in (_validated_object_list(body.get("tools")) or ()) + ) + for body in bodies + ) @staticmethod def get_default_injection_points( @@ -676,36 +681,19 @@ class AnthropicCacheControlHook(CustomPromptManagement): even when the global flag is off. Caches the system prompt and the trailing turn, so the stable prefix (system + tools + history) is reused while the breakpoint advances with the conversation. Returns [] - (stand down) when neither flag is on, the provider does not consume - cache_control breakpoints (only anthropic / bedrock do), the model - lacks prompt-caching support, or the request already carries - client-supplied cache_control. + (stand down) when neither flag is on, the model is not Claude on a + supported explicit-cache transport, the model lacks prompt-caching + support, or the request already carries client-supplied cache_control. """ import litellm if litellm.enable_anthropic_prompt_caching is not True and enable_prompt_caching is not True: return [] - provider = custom_llm_provider - if provider is None: - from litellm.litellm_core_utils.get_llm_provider_logic import ( - get_llm_provider, - ) - - try: - _, provider, _, _ = get_llm_provider(model=model) - except Exception: # noqa: BLE001 # unroutable model must never block the call, just skip auto-caching - return [] - - if provider not in ("anthropic", "bedrock"): + if not supports_anthropic_cache_control(model, custom_llm_provider): return [] - from litellm.utils import supports_prompt_caching - - if not supports_prompt_caching(model=model, custom_llm_provider=provider): - return [] - - if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control): + if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control, request_kwargs): return [] if is_claude_code_one_shot_subagent_request( @@ -737,13 +725,15 @@ class AnthropicCacheControlHook(CustomPromptManagement): prompt and trailing turn) do not depend on which deployment serves the call. Returns the input list itself when auto-injection would not apply """ + import litellm + points: Final = next( ( candidate for candidate in ( AnthropicCacheControlHook.get_default_injection_points( messages=messages, - model=model, + model=litellm.model_alias_map.get(model, model), custom_llm_provider=None, tools=tools, enable_prompt_caching=enable_prompt_caching, @@ -789,6 +779,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): prompt-management gate and the AnthropicCacheControlHook run unchanged. """ + import litellm + if non_default_params.get("cache_control_injection_points"): judged: Final = AnthropicCacheControlHook._judged_configured_points( non_default_params["cache_control_injection_points"], @@ -799,6 +791,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): custom_llm_provider, api_base, non_default_params.get("prompt_cache_options"), + non_default_params, ) if judged is None: non_default_params.pop("cache_control_injection_points") @@ -808,7 +801,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): points: Final = AnthropicCacheControlHook.get_default_injection_points( messages=messages, system=None, - model=model, + model=litellm.model_alias_map.get(model, model), custom_llm_provider=custom_llm_provider, tools=tools, enable_prompt_caching=enable_prompt_caching, @@ -925,7 +918,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None) ) if configured and AnthropicCacheControlHook._should_stand_down( - configured, typed_messages, system, tools, cache_control + configured, typed_messages, system, tools, cache_control, kwargs ): return messages, system injection_points: list[CacheControlInjectionPoint] = configured or [] diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py index 5a5324eae5e..0271cf1e03c 100644 --- a/litellm/integrations/arize/_utils.py +++ b/litellm/integrations/arize/_utils.py @@ -1139,7 +1139,10 @@ def _set_mcp_tool_output(span: "Span", coerced_response_obj: object) -> None: safe_set_attribute(span, SpanAttributes.OUTPUT_MIME_TYPE, OpenInferenceMimeTypeValues.TEXT.value) return - structured: Final[object] = coerced_response_obj.get("structuredContent") + structured: Final[object] = coerced_response_obj.get( + "structured_content", + coerced_response_obj.get("structuredContent"), # pyright: ignore[reportUnknownMemberType] # tolerant dual-spelling lookup on untyped payloads + ) payload: Final[object] = content if content else structured if structured is not None else content if payload is None: return diff --git a/litellm/integrations/otel/mappers/langfuse.py b/litellm/integrations/otel/mappers/langfuse.py index e76cffde881..9aff944cff0 100644 --- a/litellm/integrations/otel/mappers/langfuse.py +++ b/litellm/integrations/otel/mappers/langfuse.py @@ -39,7 +39,7 @@ LANGFUSE_TRACE_TAGS: Final = "langfuse.trace.tags" class LangfuseMapper: _LLM_CALL_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = { - "langfuse.observation.type": lambda d: "generation", + "langfuse.observation.type": lambda _: "generation", "langfuse.observation.model.name": lambda d: d.request_model or None, "langfuse.observation.metadata.provider": lambda d: d.provider or None, "langfuse.observation.id": lambda d: d.identity.call_id or None, @@ -68,7 +68,9 @@ class LangfuseMapper: collect(LangfuseMapper._MODEL_PARAMS, d.request_params) ), LANGFUSE_OBSERVATION_INPUT: lambda d: serialize_messages(d.messages_in), - LANGFUSE_OBSERVATION_OUTPUT: lambda d: serialize_messages(output_messages(d)), + LANGFUSE_OBSERVATION_OUTPUT: lambda d: ( + d.embedding_output.as_json() if d.embedding_output is not None else serialize_messages(output_messages(d)) + ), "langfuse.observation.usage_details": lambda d: json_if(collect(LangfuseMapper._USAGE_FIELDS, d.usage)), "langfuse.observation.cost_details": lambda d: ( json.dumps({"total": d.response_cost}) if d.response_cost is not None else None diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index 33da1549fd5..467c286db9d 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -3,7 +3,7 @@ from __future__ import annotations import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from enum import Enum from types import MappingProxyType @@ -353,6 +353,24 @@ class ToolDefinition: parameters_json: str | None = None # JSON-serialized schema (str so it's an AttrValue) +@dataclass(frozen=True, slots=True) +class EmbeddingOutput: + count: int + dimensions: int | None + + @classmethod + def from_response(cls, response: Mapping[str, object]) -> EmbeddingOutput | None: + vectors: Final = tuple(row.get("embedding") for row in _dicts(response.get("data"))) + if not vectors: + return None + first: Final = vectors[0] + width: Final = len(cast(Sequence[object], first)) if isinstance(first, list) else None + return cls(count=len(vectors), dimensions=width) + + def as_json(self) -> str: + return json.dumps({"count": self.count, "dimensions": self.dimensions}) + + @dataclass(frozen=True) class LLMCallSpanData: operation: GenAIOperation @@ -386,6 +404,7 @@ class LLMCallSpanData: call_type: str | None = None request_route: str | None = None trace: TraceControls = field(default_factory=TraceControls) + embedding_output: EmbeddingOutput | None = None @classmethod def from_standard_logging_payload( @@ -413,8 +432,12 @@ class LLMCallSpanData: # no prompt/response text. finish_reasons: Final = _finish_reasons(choices_out) call_type: Final = as_str(payload.get("call_type")) + operation: Final = resolve_operation(call_type) + embedding_output: Final = ( + EmbeddingOutput.from_response(response) if operation is GenAIOperation.EMBEDDINGS else None + ) return cls( - operation=resolve_operation(call_type), + operation=operation, provider=resolve_provider(as_str(payload.get("custom_llm_provider"))), request_model=context.request_model, response_model=context.response_model, @@ -437,6 +460,7 @@ class LLMCallSpanData: call_type=call_type or None, request_route=request_route or context.identity.request_route, trace=trace or TraceControls(), + embedding_output=embedding_output if capture_content else None, ) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 3328e20e2ea..98e2f6d5bde 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -77,6 +77,21 @@ _CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) _CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/") +def supports_anthropic_cache_control(model: str, custom_llm_provider: str | None) -> bool: + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + from litellm.utils import supports_prompt_caching + + try: + provider: Final = custom_llm_provider if custom_llm_provider is not None else get_llm_provider(model=model)[1] + except Exception: # noqa: BLE001 # Optional caching must not block an unroutable request + return False + return ( + provider in ("anthropic", "bedrock", "vertex_ai", "azure_ai") + and "claude" in model.lower() + and supports_prompt_caching(model=model, custom_llm_provider=provider) + ) + + def is_claude_code_user_agent(user_agent: str) -> bool: """Claude Code sends its API calls through the Anthropic SDK as `claude-cli/` and its own fetches, such as gateway model discovery, as `claude-code/`""" diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4dbf0337894..53c0807e86c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5300,7 +5300,7 @@ "supports_web_search": false }, "azure/gpt-4.1-nano": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -5334,7 +5334,7 @@ "supports_vision": true }, "azure/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -10207,7 +10207,7 @@ "supports_web_search": false }, "azure/us/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -14246,7 +14246,6 @@ "source": "https://developers.openai.com/api/docs/pricing" }, "claude-haiku-4-5-20251001": { - "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -14270,7 +14269,6 @@ "prompt_cache_min_tokens": 4096 }, "claude-haiku-4-5": { - "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -14420,7 +14418,6 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-4-5": { - "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, @@ -14455,7 +14452,6 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-sonnet-4-5-20250929": { - "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, @@ -14491,7 +14487,6 @@ "source": "https://docs.anthropic.com/en/docs/about-claude/pricing" }, "claude-sonnet-5": { - "deprecation_date": "2027-06-30", "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -14530,7 +14525,6 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-sonnet-4-6": { - "deprecation_date": "2027-02-17", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -14684,7 +14678,6 @@ "prompt_cache_min_tokens": 1024 }, "claude-opus-4-5-20251101": { - "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -14714,7 +14707,6 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-5": { - "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -14745,7 +14737,6 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-4-6": { - "deprecation_date": "2027-02-05", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -14783,7 +14774,6 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-4-6-20260205": { - "deprecation_date": "2027-02-05", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -14820,7 +14810,6 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-7": { - "deprecation_date": "2027-04-16", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -14859,7 +14848,6 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-4-7-20260416": { - "deprecation_date": "2027-04-16", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -14897,7 +14885,6 @@ "prompt_cache_min_tokens": 2048 }, "claude-fable-5": { - "deprecation_date": "2027-06-09", "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "cache_read_input_token_cost": 1e-06, @@ -14937,7 +14924,6 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-fable-5-1": { - "deprecation_date": "2027-09-01", "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "cache_read_input_token_cost": 2.5e-07, @@ -14978,7 +14964,6 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-5": { - "deprecation_date": "2027-07-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -15020,7 +15005,6 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-4-8": { - "deprecation_date": "2027-05-28", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -40942,7 +40926,7 @@ "supports_prompt_caching": true, "supports_reasoning": false, "supports_response_schema": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-3.5-sonnet": { "input_cost_per_token": 3e-06, @@ -40998,7 +40982,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.1": { "input_cost_per_image": 0.0048, @@ -41024,7 +41008,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-sonnet-4": { "input_cost_per_image": 0.0048, @@ -41054,7 +41038,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-sonnet-4.6": { "supports_adaptive_thinking": true, @@ -41086,7 +41070,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.5": { "cache_creation_input_token_cost": 6.25e-06, @@ -41112,7 +41096,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.6": { "supports_adaptive_thinking": true, @@ -41140,7 +41124,7 @@ "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_pdf_input": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-sonnet-4.5": { "input_cost_per_image": 0.0048, @@ -41170,7 +41154,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-haiku-4.5": { "cache_creation_input_token_cost": 1.25e-06, @@ -41195,7 +41179,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.7": { "supports_adaptive_thinking": true, @@ -41223,7 +41207,7 @@ "prompt_cache_min_tokens": 2048, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-5": { "prompt_cache_min_tokens": 512, @@ -41250,7 +41234,7 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, "openrouter/bytedance/ui-tars-1.5-7b": { @@ -41420,35 +41404,36 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 1.6e-06, + "input_cost_per_token": 4.22298e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 3.2e-06, + "output_cost_per_token": 8.44596e-07, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.35e-07, + "cache_read_input_token_cost": 3.51915e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, "supports_web_search": false }, "openrouter/deepseek/deepseek-v4.1-flash": { - "input_cost_per_token": 1.5e-07, - "output_cost_per_token": 6e-07, - "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 6e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":1.5e-7,"output_cost_per_token":6e-7,"cache_read_input_token_cost":3e-9}, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -41523,7 +41508,7 @@ "supports_audio_input": true, "supports_pdf_input": true, "supports_reasoning": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-2.5-pro": { @@ -41553,7 +41538,7 @@ "supports_audio_input": true, "supports_pdf_input": true, "supports_reasoning": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3-pro-preview": { @@ -41638,7 +41623,7 @@ "supports_tool_choice": true, "supports_url_context": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "tpm": 800000, "supports_video_input": true }, @@ -41684,7 +41669,7 @@ "supports_url_context": true, "supports_video_input": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "tpm": 800000 }, "openrouter/google/gemini-3.1-flash-lite": { @@ -41729,7 +41714,7 @@ "supports_url_context": true, "supports_video_input": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "tpm": 800000 }, "openrouter/google/gemini-3.1-pro-preview": { @@ -41767,7 +41752,7 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/gryphe/mythomax-l2-13b": { @@ -42053,11 +42038,11 @@ }, "openrouter/nvidia/nemotron-3.5-lightning": { "cache_read_input_token_cost": 4e-08, - "input_cost_per_token": 8e-08, + "input_cost_per_token": 7e-08, "litellm_provider": "openrouter", "max_input_tokens": 262144, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 235929, + "max_tokens": 235929, "mode": "chat", "output_cost_per_token": 2e-07, "source": "https://openrouter.ai/api/v1/models", @@ -42148,7 +42133,7 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-4.1-mini": { "cache_read_input_token_cost": 1e-07, @@ -42170,7 +42155,7 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-4.1-nano": { "cache_read_input_token_cost": 2.5e-08, @@ -42192,7 +42177,7 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-4o": { "input_cost_per_token": 2.5e-06, @@ -42298,7 +42283,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5": { "cache_read_input_token_cost": 1.25e-07, @@ -42325,7 +42310,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5-mini": { "cache_read_input_token_cost": 2.5e-08, @@ -42352,7 +42337,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5-nano": { "cache_read_input_token_cost": 5e-09, @@ -42379,7 +42364,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.1-codex-max": { "cache_read_input_token_cost": 1.25e-07, @@ -42406,7 +42391,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.2": { "input_cost_per_image": 0, @@ -42427,7 +42412,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.2-chat": { "input_cost_per_image": 0, @@ -42448,7 +42433,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.2-pro": { "input_cost_per_image": 0, @@ -42468,7 +42453,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-sol": { "cache_creation_input_token_cost": 2.5e-06, @@ -42509,7 +42494,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-sol-pro": { "input_cost_per_token": 2e-06, @@ -42534,15 +42519,15 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-oss-120b": { "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 1.5e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, - "max_output_tokens": 117964, - "max_tokens": 117964, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 6e-07, "source": "https://openrouter.ai/api/v1/models", @@ -42550,7 +42535,7 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": false, - "supports_prompt_caching": false, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -42598,7 +42583,7 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/o3-mini": { "input_cost_per_token": 1.1e-06, @@ -42619,7 +42604,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/o3-mini-high": { "input_cost_per_token": 1.1e-06, @@ -42640,7 +42625,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/qwen/qwen-2.5-coder-32b-instruct": { "input_cost_per_token": 6.6e-07, @@ -60378,7 +60363,6 @@ } }, "claude-mythos-5-1": { - "deprecation_date": "2027-09-01", "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "cache_read_input_token_cost": 2.5e-07, @@ -65630,7 +65614,7 @@ "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "prompt_cache_min_tokens": 512, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-fable-5.1": { "input_cost_per_token": 1e-05, @@ -65656,7 +65640,7 @@ "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "prompt_cache_min_tokens": 512, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.8": { "input_cost_per_token": 5e-06, @@ -65680,7 +65664,7 @@ "supports_prompt_caching": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-sonnet-5": { "input_cost_per_token": 2e-06, @@ -65704,7 +65688,7 @@ "supports_prompt_caching": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, - "supports_web_search": false + "supports_web_search": true }, "openrouter/google/gemini-2.5-flash-lite": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -65728,7 +65712,7 @@ "deprecation_date": "2026-10-20", "input_cost_per_audio_token": 3e-07, "supports_prompt_caching": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.5-flash": { @@ -65752,7 +65736,7 @@ "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 3e-06, "supports_prompt_caching": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.5-flash-lite": { @@ -65776,7 +65760,7 @@ "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 3e-07, "supports_prompt_caching": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.6-flash": { @@ -65800,7 +65784,7 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.7-flash": { @@ -65824,7 +65808,7 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.8-flash": { @@ -65848,7 +65832,7 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/openai/gpt-4o-mini": { @@ -65889,7 +65873,7 @@ "supports_audio_input": false, "cache_read_input_token_cost": 1.25e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.3-codex": { "input_cost_per_token": 1.75e-06, @@ -65909,7 +65893,7 @@ "supports_audio_input": false, "cache_read_input_token_cost": 1.75e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.4": { "input_cost_per_token": 2.5e-06, @@ -65932,7 +65916,7 @@ "input_cost_per_token_above_272k_tokens": 5e-06, "output_cost_per_token_above_272k_tokens": 2.25e-05, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.4-mini": { "input_cost_per_token": 7.5e-07, @@ -65952,7 +65936,7 @@ "supports_audio_input": false, "cache_read_input_token_cost": 7.5e-08, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.4-nano": { "input_cost_per_token": 2e-07, @@ -65972,7 +65956,7 @@ "supports_audio_input": false, "cache_read_input_token_cost": 2e-08, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.5": { "input_cost_per_token": 5e-06, @@ -65995,7 +65979,7 @@ "input_cost_per_token_above_272k_tokens": 1e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-luna": { "cache_creation_input_token_cost": 2.5e-07, @@ -66020,7 +66004,7 @@ "input_cost_per_token_above_272k_tokens": 4e-07, "output_cost_per_token_above_272k_tokens": 1.8e-06, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-luna-pro": { "input_cost_per_token": 2e-07, @@ -66045,7 +66029,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-terra": { "cache_creation_input_token_cost": 2.5e-06, @@ -66070,7 +66054,7 @@ "input_cost_per_token_above_272k_tokens": 4e-06, "output_cost_per_token_above_272k_tokens": 1.8e-05, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-terra-pro": { "input_cost_per_token": 2e-06, @@ -66095,7 +66079,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/o3": { "input_cost_per_token": 2e-06, @@ -66115,7 +66099,7 @@ "supports_audio_input": false, "cache_read_input_token_cost": 5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/o4-mini": { "input_cost_per_token": 1.1e-06, @@ -66135,7 +66119,7 @@ "supports_audio_input": false, "cache_read_input_token_cost": 2.75e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/x-ai/grok-4.20": { "input_cost_per_token": 1.25e-06, @@ -66158,7 +66142,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/x-ai/grok-4.20-multi-agent": { "input_cost_per_token": 1.25e-06, @@ -66181,7 +66165,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/x-ai/grok-4.3": { "input_cost_per_token": 1.25e-06, @@ -66204,7 +66188,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/x-ai/grok-4.5": { "input_cost_per_token": 2e-06, @@ -66227,7 +66211,7 @@ "input_cost_per_token_above_200k_tokens": 4e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/x-ai/grok-4.6": { "input_cost_per_token": 2e-06, @@ -66250,7 +66234,7 @@ "input_cost_per_token_above_200k_tokens": 4e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/x-ai/grok-build-0.1": { "input_cost_per_token": 1e-06, @@ -66273,7 +66257,7 @@ "input_cost_per_token_above_200k_tokens": 2e-06, "output_cost_per_token_above_200k_tokens": 4e-06, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "baseten/zai-org/GLM-5.3": { "cache_read_input_token_cost": 1.4e-07, @@ -66362,7 +66346,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-6-astra-pro": { "input_cost_per_token": 1e-05, @@ -66387,7 +66371,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/qwen/qwen3.8-flash": { "input_cost_per_token": 1.5e-07, @@ -66436,8 +66420,8 @@ "cache_read_input_token_cost": 6.86e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -66456,8 +66440,8 @@ "cache_read_input_token_cost": 1.69e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 943717, - "max_tokens": 943717, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -66570,9 +66554,9 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-flash-0731": { - "input_cost_per_token": 6e-08, - "output_cost_per_token": 1.2e-07, - "cache_read_input_token_cost": 1.2e-08, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 8e-08, + "cache_read_input_token_cost": 1.6e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 943718, @@ -66655,9 +66639,9 @@ "supports_web_search": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 2.1e-06, - "output_cost_per_token": 1.095e-05, - "cache_read_input_token_cost": 2.3e-07, + "input_cost_per_token": 1.7e-06, + "output_cost_per_token": 8.5e-06, + "cache_read_input_token_cost": 1.7e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -66731,7 +66715,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/google/gemini-3.1-flash-image": { "input_cost_per_token": 5e-07, @@ -66751,7 +66735,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/google/gemini-3-pro-image": { "input_cost_per_token": 2e-06, @@ -66775,7 +66759,7 @@ "supports_response_schema": true, "supports_vision": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/z-ai/glm-5.2": { "input_cost_per_token": 5.544e-07, @@ -66877,13 +66861,13 @@ "supports_web_search": false }, "openrouter/nvidia/nemotron-3-ultra-550b-a55b": { - "input_cost_per_token": 6.25e-07, - "output_cost_per_token": 3.125e-06, - "cache_read_input_token_cost": 1.875e-07, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 2.4e-06, + "cache_read_input_token_cost": 1.2e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, - "max_output_tokens": 32768, - "max_tokens": 32768, + "max_output_tokens": 182520, + "max_tokens": 182520, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -67117,7 +67101,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-chat-latest": { "input_cost_per_token": 5e-06, @@ -67137,12 +67121,12 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 4.984e-08, - "output_cost_per_token": 9.968e-08, - "cache_read_input_token_cost": 9.968e-09, + "input_cost_per_token": 4.06e-08, + "output_cost_per_token": 8.12e-08, + "cache_read_input_token_cost": 8.12e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -67431,7 +67415,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/google/gemini-3.1-flash-image-preview": { "input_cost_per_token": 5e-07, @@ -67451,7 +67435,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/google/gemini-3.1-pro-preview-customtools": { "input_cost_per_token": 2e-06, @@ -67477,7 +67461,7 @@ "supports_pdf_input": true, "supports_audio_input": true, "supports_prompt_caching": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/qwen/qwen3-max-thinking": { @@ -67645,7 +67629,7 @@ "supports_response_schema": true, "supports_vision": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.1-codex": { "input_cost_per_token": 1.25e-06, @@ -67665,7 +67649,7 @@ "supports_response_schema": true, "supports_vision": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.1-codex-mini": { "input_cost_per_token": 2.5e-07, @@ -67685,7 +67669,7 @@ "supports_response_schema": true, "supports_vision": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/moonshotai/kimi-k2-thinking": { "input_cost_per_token": 6e-07, @@ -67828,7 +67812,7 @@ "supports_vision": true, "supports_prompt_caching": true, "supports_reasoning": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/qwen/qwen3-vl-30b-a3b-thinking": { "input_cost_per_token": 2e-07, @@ -67885,7 +67869,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/qwen/qwen3-vl-235b-a22b-thinking": { "input_cost_per_token": 4e-07, @@ -68051,7 +68035,7 @@ "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, - "supports_prompt_caching": false, + "supports_prompt_caching": true, "supports_reasoning": false, "supports_tool_choice": true, "supports_response_schema": true, @@ -68290,7 +68274,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/google/gemini-2.5-pro-preview": { "input_cost_per_token": 1.25e-06, @@ -68316,7 +68300,7 @@ "supports_pdf_input": true, "supports_audio_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/mistralai/mistral-medium-3": { "input_cost_per_token": 4e-07, @@ -68494,7 +68478,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/meta-llama/llama-4-maverick": { "input_cost_per_token": 1.875e-07, @@ -68551,7 +68535,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/google/gemma-3-4b-it": { "input_cost_per_token": 5e-08, @@ -69338,7 +69322,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-4.1-nano": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -69692,7 +69676,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-4.1-nano": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -71106,7 +71090,7 @@ "supports_response_schema": true, "supports_tool_choice": false, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~anthropic/claude-haiku-latest": { "cache_creation_input_token_cost": 1.25e-06, @@ -71128,7 +71112,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~anthropic/claude-opus-latest": { "cache_creation_input_token_cost": 6.25e-06, @@ -71150,7 +71134,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~anthropic/claude-sonnet-latest": { "cache_creation_input_token_cost": 2.5e-06, @@ -71172,17 +71156,17 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~deepseek/deepseek-flash-latest": { - "cache_read_input_token_cost": 4.2e-09, - "input_cost_per_token": 1.4e-07, + "cache_read_input_token_cost": 2.6e-09, + "input_cost_per_token": 1.3e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 4.2e-07, + "output_cost_per_token": 5.2e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71215,14 +71199,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-v4-flash-latest": { - "cache_read_input_token_cost": 1.75e-09, - "input_cost_per_token": 5.5e-08, + "cache_read_input_token_cost": 1.6e-08, + "input_cost_per_token": 4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 1.65e-07, + "output_cost_per_token": 8e-08, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71255,7 +71239,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~google/gemini-pro-latest": { "cache_creation_input_token_cost": 3.75e-07, @@ -71281,17 +71265,17 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~moonshotai/kimi-latest": { - "cache_read_input_token_cost": 2.3e-07, - "input_cost_per_token": 2.1e-06, + "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 1.7e-06, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 1.095e-05, + "output_cost_per_token": 8.5e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71326,7 +71310,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~openai/gpt-luna-latest": { "cache_creation_input_token_cost": 2.5e-07, @@ -71351,7 +71335,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~openai/gpt-mini-latest": { "cache_read_input_token_cost": 7.5e-08, @@ -71371,7 +71355,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~openai/gpt-sol-latest": { "cache_creation_input_token_cost": 2.5e-06, @@ -71396,7 +71380,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~openai/gpt-terra-latest": { "cache_creation_input_token_cost": 2.5e-06, @@ -71421,7 +71405,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~x-ai/grok-latest": { "cache_read_input_token_cost": 5e-07, @@ -71444,7 +71428,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~z-ai/glm-flash-latest": { "cache_read_input_token_cost": 1.5e-08, @@ -71467,14 +71451,14 @@ "supports_web_search": false }, "openrouter/~z-ai/glm-latest": { - "cache_read_input_token_cost": 1.46625e-07, - "input_cost_per_token": 9e-07, + "cache_read_input_token_cost": 1.5678e-07, + "input_cost_per_token": 8.442e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 235929, - "max_tokens": 235929, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.805e-06, + "output_cost_per_token": 2.6532e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71700,7 +71684,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-fable-5.1:batch": { "cache_creation_input_token_cost": 6.25e-06, @@ -71722,7 +71706,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-haiku-4.5:batch": { "cache_creation_input_token_cost": 6.25e-07, @@ -71744,7 +71728,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.1:batch": { "cache_creation_input_token_cost": 9.375e-06, @@ -71766,7 +71750,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.5:batch": { "cache_creation_input_token_cost": 3.125e-06, @@ -71788,7 +71772,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.6:batch": { "cache_creation_input_token_cost": 3.125e-06, @@ -71810,7 +71794,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.7:batch": { "cache_creation_input_token_cost": 3.125e-06, @@ -71832,7 +71816,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.8:batch": { "cache_creation_input_token_cost": 3.125e-06, @@ -71854,7 +71838,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-5:batch": { "cache_creation_input_token_cost": 3.125e-06, @@ -71876,7 +71860,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-sonnet-4.5:batch": { "cache_creation_input_token_cost": 1.875e-06, @@ -71902,7 +71886,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-sonnet-4.6:batch": { "cache_creation_input_token_cost": 1.875e-06, @@ -71924,7 +71908,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-sonnet-5:batch": { "cache_creation_input_token_cost": 1.25e-06, @@ -71946,7 +71930,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/arcee-ai/trinity-large-thinking": { "cache_read_input_token_cost": 6e-08, @@ -72344,7 +72328,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-2.5-flash:batch": { @@ -72368,7 +72352,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-2.5-pro:batch": { @@ -72395,7 +72379,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3-flash-preview:batch": { @@ -72416,7 +72400,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.1-flash-lite:batch": { @@ -72439,7 +72423,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.1-pro-preview:batch": { @@ -72462,7 +72446,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.5-flash-lite:batch": { @@ -72485,7 +72469,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.5-flash:batch": { @@ -72508,7 +72492,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.6-flash:batch": { @@ -72532,7 +72516,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.7-flash:batch": { @@ -72556,7 +72540,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.8-flash:batch": { @@ -72580,7 +72564,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/ibm-granite/granite-4.0-h-micro": { @@ -72956,7 +72940,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/meta/muse-spark-1.2": { "cache_read_input_token_cost": 1.5e-07, @@ -72976,7 +72960,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/meta/muse-spark-1.2-contributor": { "cache_read_input_token_cost": 2e-09, @@ -72996,7 +72980,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/meta/muse-spark-1.3": { "cache_read_input_token_cost": 1.5e-07, @@ -73016,7 +73000,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/meta/muse-spark-1.3-contributor": { "cache_read_input_token_cost": 2e-09, @@ -73036,7 +73020,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/microsoft/phi-4": { "input_cost_per_token": 7e-08, @@ -73404,7 +73388,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-4-turbo:batch": { "input_cost_per_token": 5e-06, @@ -73423,7 +73407,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-4.1-mini:batch": { "cache_read_input_token_cost": 5e-08, @@ -73443,7 +73427,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-4.1-nano:batch": { "cache_read_input_token_cost": 1.25e-08, @@ -73463,7 +73447,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-4.1:batch": { "cache_read_input_token_cost": 2.5e-07, @@ -73483,7 +73467,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-4o-mini:batch": { "cache_read_input_token_cost": 3.75e-08, @@ -73543,7 +73527,7 @@ "supports_response_schema": true, "supports_tool_choice": false, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5-image-mini": { "cache_read_input_token_cost": 2.5e-07, @@ -73563,7 +73547,7 @@ "supports_response_schema": true, "supports_tool_choice": false, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5-mini:batch": { "cache_read_input_token_cost": 1.25e-08, @@ -73583,7 +73567,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5-nano:batch": { "cache_read_input_token_cost": 2.5e-09, @@ -73603,7 +73587,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5-pro:batch": { "input_cost_per_token": 7.5e-06, @@ -73622,7 +73606,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5:batch": { "cache_read_input_token_cost": 6.25e-08, @@ -73642,7 +73626,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.1:batch": { "cache_read_input_token_cost": 6.25e-08, @@ -73662,7 +73646,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.2-pro:batch": { "input_cost_per_token": 1.05e-05, @@ -73681,7 +73665,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.2:batch": { "cache_read_input_token_cost": 8.75e-08, @@ -73701,7 +73685,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.4-image-2": { "cache_read_input_token_cost": 2e-06, @@ -73721,7 +73705,7 @@ "supports_response_schema": true, "supports_tool_choice": false, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.4-mini:batch": { "cache_read_input_token_cost": 3.75e-08, @@ -73741,7 +73725,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.4-nano:batch": { "cache_read_input_token_cost": 1e-08, @@ -73761,7 +73745,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.4-pro:batch": { "input_cost_per_token": 1.5e-05, @@ -73782,7 +73766,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.4:batch": { "cache_read_input_token_cost": 1.25e-07, @@ -73805,7 +73789,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.5-pro:batch": { "input_cost_per_token": 1.5e-05, @@ -73826,7 +73810,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.5:batch": { "cache_read_input_token_cost": 2.5e-07, @@ -73849,7 +73833,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-luna-pro:batch": { "cache_read_input_token_cost": 1e-08, @@ -73872,7 +73856,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-luna:batch": { "cache_read_input_token_cost": 1e-08, @@ -73895,7 +73879,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-sol-pro:batch": { "cache_creation_input_token_cost": 1.25e-06, @@ -73920,7 +73904,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-sol:batch": { "cache_creation_input_token_cost": 1.25e-06, @@ -73945,7 +73929,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-terra-pro:batch": { "cache_read_input_token_cost": 1e-07, @@ -73968,7 +73952,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-terra:batch": { "cache_read_input_token_cost": 1e-07, @@ -73991,7 +73975,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-6-astra-pro:batch": { "cache_creation_input_token_cost": 6.25e-06, @@ -74016,7 +74000,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-6-astra:batch": { "cache_creation_input_token_cost": 6.25e-06, @@ -74041,7 +74025,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-oss-120b:batch": { "input_cost_per_token": 1.5e-07, @@ -74080,7 +74064,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/o3:batch": { "cache_read_input_token_cost": 2.5e-07, @@ -74100,7 +74084,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/o4-mini:batch": { "cache_read_input_token_cost": 1.375e-07, @@ -74120,7 +74104,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/perceptron/perceptron-mk1": { "input_cost_per_token": 1.5e-07, @@ -74629,14 +74613,15 @@ "supports_web_search": false }, "openrouter/tencent/hy3": { - "cache_read_input_token_cost": 2.0625e-08, - "input_cost_per_token": 8.25e-08, + "cache_read_input_token_cost": 3.3e-08, + "input_cost_per_token": 1.32e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3.3e-07, + "off_peak_pricing": {"hours_utc":"16:00-00:00","input_cost_per_token":8.25e-8,"output_cost_per_token":3.3e-7,"cache_read_input_token_cost":2.0625e-8}, + "output_cost_per_token": 5.28e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -74860,7 +74845,7 @@ "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": false, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false @@ -74945,7 +74930,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/z-ai/glm-5.2:batch": { "cache_read_input_token_cost": 7e-08, @@ -75006,5 +74991,44 @@ "supports_tool_choice": true, "supports_vision": false, "supports_web_search": false + }, + "openrouter/prism-ml/ternary-bonsai-2-27b": { + "input_cost_per_token": 7.5e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_pdf_input": false, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": false + }, + "openrouter/z-ai/glm-5.3-flashx": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 3.7e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.25e-06, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_pdf_input": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": false } } diff --git a/litellm/ocr/dispatch.py b/litellm/ocr/dispatch.py index 80c93273d1e..55b19458b7a 100644 --- a/litellm/ocr/dispatch.py +++ b/litellm/ocr/dispatch.py @@ -53,7 +53,9 @@ _PYTHON_AOCR: Final = cast( # cast-ok: forward the original call shape through def _context(request: LiteLLMOcrRequest) -> Context: - return Context(Route.OCR, provider=request.custom_llm_provider, model=request.model) + prefix, separator, _ = request.model.partition("/") + provider: Final = request.custom_llm_provider or (prefix if separator else None) + return Context(Route.OCR, provider=provider, model=request.model) _DISPATCH: Final = PublicDispatch( diff --git a/litellm/proxy/_experimental/mcp_server/elicitation_handler.py b/litellm/proxy/_experimental/mcp_server/elicitation_handler.py index bbd1c9aaf1e..6155f1f215c 100644 --- a/litellm/proxy/_experimental/mcp_server/elicitation_handler.py +++ b/litellm/proxy/_experimental/mcp_server/elicitation_handler.py @@ -42,9 +42,9 @@ class _DownstreamElicitSession(Protocol): async def elicit_url(self, message: str, url: str, elicitation_id: str) -> "ElicitResult": ... - async def elicit_form(self, message: str, requestedSchema: dict[str, object]) -> "ElicitResult": ... + async def elicit_form(self, message: str, requested_schema: dict[str, object]) -> "ElicitResult": ... - async def elicit(self, message: str, requestedSchema: dict[str, object]) -> "ElicitResult": ... + async def elicit(self, message: str, requested_schema: dict[str, object]) -> "ElicitResult": ... async def handle_elicitation_request( @@ -145,22 +145,22 @@ async def _relay_elicitation_to_downstream( result = await downstream_session.elicit_url( message=params.message, url=params.url, - elicitation_id=params.elicitationId, + elicitation_id=params.elicitation_id, ) elif isinstance(params, ElicitRequestFormParams): # Form mode: relay structured form to client verbose_logger.info("MCP elicitation: relaying form mode to downstream") result = await downstream_session.elicit_form( message=params.message, - requestedSchema=params.requestedSchema, + requested_schema=params.requested_schema, ) else: # Fallback for generic ElicitRequestParams — pass an empty schema - # since elicit() requires requestedSchema as a positional arg. + # since elicit() requires requested_schema as a positional arg. verbose_logger.info("MCP elicitation: relaying generic elicitation to downstream") result = await downstream_session.elicit( message=getattr(params, "message", ""), - requestedSchema=getattr(params, "requestedSchema", {}), + requested_schema=getattr(params, "requested_schema", {}), # mutable-ok: elicitation default schema ) verbose_logger.info( "MCP elicitation: downstream responded with action=%s", diff --git a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py index 42b2d29cd52..b96a7a74e4a 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py +++ b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py @@ -14,6 +14,7 @@ from collections.abc import Iterator from typing import Final, Literal, NamedTuple, NoReturn, TypeAlias import httpx +import httpx2 from mcp.types import Tool as MCPTool from pydantic import BaseModel, ConfigDict from typing_extensions import assert_never @@ -63,8 +64,8 @@ class AggregateToolListing(NamedTuple): outcomes: dict[str, ServerOutcome] -def _iter_upstream_responses(exc: BaseException) -> Iterator[httpx.Response]: - """Yield every ``httpx.Response`` in the exception tree, in the shared traversal's deliberate +def _iter_upstream_responses(exc: BaseException) -> Iterator[httpx.Response | httpx2.Response]: + """Yield every upstream ``httpx``/``httpx2`` ``Response`` in the exception tree, in the shared traversal's deliberate order (explicit causes first, ExceptionGroup members in raise order, the incidental ``__context__`` chain last), so a response raised while handling the real failure can never shadow one on the explicit causal chain. Consumers apply their own predicate over the stream: @@ -72,11 +73,11 @@ def _iter_upstream_responses(exc: BaseException) -> Iterator[httpx.Response]: behind an unrelated earlier one.""" for current in iter_exception_tree(exc): response = getattr(current, "response", None) - if isinstance(response, httpx.Response): + if isinstance(response, (httpx.Response, httpx2.Response)): yield response -def _find_upstream_response(exc: BaseException) -> httpx.Response | None: +def _find_upstream_response(exc: BaseException) -> httpx.Response | httpx2.Response | None: return next(_iter_upstream_responses(exc), None) @@ -136,9 +137,9 @@ def classify_list_exception(exc: BaseException) -> ServerListFault: response: Final = _find_upstream_response(exc) if response is not None: return ServerListFault(tag="upstream_error", status_code=response.status_code) - if isinstance(exc, (httpx.TimeoutException,)): + if isinstance(exc, (httpx.TimeoutException, httpx2.TimeoutException)): return ServerListFault(tag="timeout") - if isinstance(exc, httpx.TransportError): + if isinstance(exc, (httpx.TransportError, httpx2.TransportError)): return ServerListFault(tag="unreachable") return ServerListFault(tag="internal") diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py index c0235077ecd..08a5d2b4135 100644 --- a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py +++ b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py @@ -135,7 +135,7 @@ class MCPGuardrailTranslationHandler(BaseTranslation): mcp_tool: Final = MCPTool( name=mcp_tool_name, description=mcp_tool_description or "", - inputSchema={}, # Call payload has no schema; guardrail gets args from request_data + input_schema={}, # mutable-ok: call payload has no schema; guardrail gets args from request_data ) openai_tool: Final = transform_mcp_tool_to_openai_tool(mcp_tool) fn: Final = openai_tool["function"] diff --git a/litellm/proxy/_experimental/mcp_server/mcp_context.py b/litellm/proxy/_experimental/mcp_server/mcp_context.py index 74cc0c900d9..11325a9f127 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_context.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_context.py @@ -6,7 +6,23 @@ mcp_server_manager.py and server.py. """ from contextvars import ContextVar -from typing import Final +from typing import TYPE_CHECKING, Final + +if TYPE_CHECKING: + from mcp.server.context import ServerRequestContext + +# The SDK 1.x ``mcp.server.lowlevel.server.request_ctx`` ContextVar was removed in +# SDK 2, which hands each request handler a ``ServerRequestContext`` argument +# instead. The handlers set this var so downstream helpers (session auth caching, +# debug diagnostics, progress forwarding) can reach the same request-scoped state. +active_mcp_request_ctx_var: Final[ContextVar["ServerRequestContext | None"]] = ContextVar( + "active_mcp_request_ctx", default=None +) + + +def get_active_mcp_request_ctx() -> "ServerRequestContext | None": + return active_mcp_request_ctx_var.get() + # Set server-side in proxy_server.py route handlers when a request arrives via # /toolset/{name}/mcp or the toolset fallback in dynamic_mcp_route. diff --git a/litellm/proxy/_experimental/mcp_server/mcp_debug.py b/litellm/proxy/_experimental/mcp_server/mcp_debug.py index 1f157aefdc3..ff482b80b50 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_debug.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_debug.py @@ -100,6 +100,8 @@ Usage with curl:: http://localhost:4000/mcp/atlassian_mcp """ +from __future__ import annotations + import asyncio import base64 import io @@ -109,17 +111,20 @@ from collections.abc import AsyncIterator, Callable, Mapping from http.cookies import CookieError, SimpleCookie from itertools import islice from types import MappingProxyType -from typing import Final +from typing import TYPE_CHECKING, Final from urllib.parse import parse_qsl, quote, quote_plus, unquote_plus, urlencode import httpx +import httpx2 from pydantic import JsonValue, TypeAdapter from starlette.requests import HTTPConnection from starlette.types import Message, Send from litellm.litellm_core_utils.secret_redaction import REDACTED, redact_string from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker -from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution + +if TYPE_CHECKING: + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution # Header the client sends to opt into debug mode MCP_DEBUG_REQUEST_HEADER: Final = "x-litellm-mcp-debug" @@ -132,9 +137,9 @@ MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: Final = "litellm.mcp.auth_diagnostics" def record_auth_resolution(server_id: str, source: AuthResolution) -> None: - from mcp.server.lowlevel.server import request_ctx + from litellm.proxy._experimental.mcp_server.mcp_context import get_active_mcp_request_ctx - context: Final[object] = request_ctx.get(None) + context: Final[object] = get_active_mcp_request_ctx() request: Final[object] = getattr(context, "request", None) if isinstance(request, HTTPConnection): diagnostics: Final[object] = request.scope.get(MCP_AUTH_DIAGNOSTICS_SCOPE_KEY) @@ -150,6 +155,8 @@ class MCPAuthDiagnostics: self._outcomes = tuple(item for item in self._outcomes if item[0] != server_id) + ((server_id, resolution),) def resolution(self) -> str: + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution + match self._outcomes: case (): return AuthResolution.unresolved.value @@ -159,6 +166,8 @@ class MCPAuthDiagnostics: return AuthResolution.multiple.value def headers(self) -> Mapping[str, str]: + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution + if len(self._outcomes) <= 1: return MappingProxyType({"x-mcp-debug-auth-resolution": self.resolution()}) return MappingProxyType( @@ -372,6 +381,8 @@ class MCPDebug: server_url: str | None = None server_auth_type: str | None = None + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution + auth_resolution: Final = AuthResolution.unresolved.value for server_name in mcp_servers or []: @@ -409,7 +420,7 @@ def _safe_text(value: str, limit: int = _BODY_PREVIEW_CHARS) -> str: return escaped if len(escaped) <= limit else f"{escaped[:limit]}...(truncated)" -def safe_upstream_url(url: httpx.URL) -> str: +def safe_upstream_url(url: httpx.URL | httpx2.URL) -> str: return _safe_text(str(url.copy_with(username="", password="", path="/", query=None, fragment=None))) @@ -449,10 +460,10 @@ def _header_secret_values(name: str, value: str) -> tuple[str, ...]: return (value, credential, decoded, password, unquote_plus(password)) -def _body_secret_values(request: httpx.Request) -> tuple[str, ...] | None: +def _body_secret_values(request: httpx.Request | httpx2.Request) -> tuple[str, ...] | None: try: raw: Final = request.content - except httpx.RequestNotRead: + except (httpx.RequestNotRead, httpx2.RequestNotRead): return None if not raw: return () @@ -478,7 +489,7 @@ def _body_secret_values(request: httpx.Request) -> tuple[str, ...] | None: ) -def _request_secret_values(request: httpx.Request) -> tuple[str, ...] | None: +def _request_secret_values(request: httpx.Request | httpx2.Request) -> tuple[str, ...] | None: body_values: Final = _body_secret_values(request) if body_values is None: return None @@ -537,18 +548,18 @@ def _preview(raw: bytes, content_type: str = "", secrets: tuple[str, ...] = ()) return _safe_text(redact_string(_mask_known_values(json.dumps(parsed, separators=(",", ":")), secrets))) -def _masked_headers(headers: httpx.Headers) -> str: +def _masked_headers(headers: httpx.Headers | httpx2.Headers) -> str: return _safe_text(", ".join(f"{name}={value}" for name, value in headers.items() if name in _SAFE_HEADER_NAMES)) -def _request_body_preview(request: httpx.Request, secrets: tuple[str, ...] | None) -> str: +def _request_body_preview(request: httpx.Request | httpx2.Request, secrets: tuple[str, ...] | None) -> str: try: return _preview(request.content, request.headers.get("content-type", ""), secrets or ()) - except httpx.RequestNotRead: + except (httpx.RequestNotRead, httpx2.RequestNotRead): return "(streamed, not captured)" -def _response_body_preview(response: httpx.Response, secrets: tuple[str, ...] | None) -> str: +def _response_body_preview(response: httpx.Response | httpx2.Response, secrets: tuple[str, ...] | None) -> str: if secrets is None: return "(omitted: request credentials unavailable)" captured: Final = response.extensions.get(_CAPTURE_EXTENSION) @@ -556,7 +567,7 @@ def _response_body_preview(response: httpx.Response, secrets: tuple[str, ...] | return captured try: return _preview(response.content, response.headers.get("content-type", ""), secrets) - except httpx.ResponseNotRead: + except (httpx.ResponseNotRead, httpx2.ResponseNotRead): return "(not read)" @@ -569,7 +580,7 @@ async def _read_error_prefix(chunks: AsyncIterator[bytes], limit: int) -> bytes: return buffer.getvalue() -async def capture_upstream_error_response(response: httpx.Response) -> None: +async def capture_upstream_error_response(response: httpx.Response | httpx2.Response) -> None: if not response.is_error: return try: @@ -584,7 +595,7 @@ async def capture_upstream_error_response(response: httpx.Response) -> None: if secrets is not None else "(omitted: request credentials unavailable)" ) - except (asyncio.TimeoutError, httpx.HTTPError, httpx.StreamError): + except (asyncio.TimeoutError, httpx.HTTPError, httpx.StreamError, httpx2.HTTPError, httpx2.StreamError): response._content = b"" # pyright: ignore[reportPrivateUsage] # rebind-ok: httpx auth retries must survive diagnostic read failures response.extensions[_CAPTURE_EXTENSION] = ( "(unavailable: error body read failed)" # rebind-ok: httpx response hooks communicate through extensions @@ -593,7 +604,7 @@ async def capture_upstream_error_response(response: httpx.Response) -> None: response.extensions[_CAPTURE_EXTENSION] = preview # rebind-ok: httpx response hooks communicate through extensions -def describe_upstream_response(response: httpx.Response) -> str: +def describe_upstream_response(response: httpx.Response | httpx2.Response) -> str: try: request: Final = response.request except RuntimeError: @@ -616,6 +627,6 @@ def describe_upstream_http_failure(exc: BaseException) -> str | None: describe_upstream_response(response) for current in islice(iter_exception_tree(exc), 16) for response in (getattr(current, "response", None),) - if isinstance(response, httpx.Response) + if isinstance(response, (httpx.Response, httpx2.Response)) ) return " | ".join(lines) or None diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 469ea86ad4b..36ecb05208b 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -34,6 +34,7 @@ from urllib.parse import ParseResult, urlparse import anyio import httpx +import httpx2 from fastapi import HTTPException from httpx import HTTPStatusError from mcp import ReadResourceResult, Resource @@ -194,8 +195,7 @@ from litellm.types.mcp_server.mcp_server_manager import ( from litellm.types.utils import CallTypes if TYPE_CHECKING: - from mcp.client.session import ClientSession - from mcp.shared.context import RequestContext + from mcp.client.session import ClientRequestContext from mcp.types import CreateMessageRequestParams from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -1297,8 +1297,8 @@ def _passthrough_token_from_mcp_auth_header( return None -async def _materialize_auth_headers(auth: httpx.Auth | None) -> dict[str, str] | None: - """Extract the header a resolved ``httpx.Auth`` would set, as a plain dict, or None. +async def _materialize_auth_headers(auth: httpx2.Auth | None) -> dict[str, str] | None: + """Extract the header a resolved ``httpx2.Auth`` would set, as a plain dict, or None. OpenAPI tool closures egress through ``AsyncHTTPHandler`` methods that accept headers but no ``auth``, so a resolved credential must be materialized into a header value. Driving one step @@ -1313,7 +1313,7 @@ async def _materialize_auth_headers(auth: httpx.Auth | None) -> dict[str, str] | header_name: Final = getattr(auth, "header_name", None) if not isinstance(header_name, str) or not header_name: return None - probe: Final = httpx.Request("GET", "http://localhost/") + probe: Final = httpx2.Request("GET", "http://localhost/") flow: Final = auth.async_auth_flow(probe) try: first_request: Final = await flow.__anext__() @@ -1587,7 +1587,7 @@ def _create_sampling_callback(user_api_key_auth: UserAPIKeyAuth | None = None): return None async def _sampling_callback( - context: "RequestContext[ClientSession, object]", + context: "ClientRequestContext", params: "CreateMessageRequestParams", ): import litellm @@ -4012,7 +4012,7 @@ class MCPServerManager: subject_token: str | None, user_api_key_auth: UserAPIKeyAuth | None, extra_headers: dict[str, str] | None, - ) -> tuple[httpx.Auth | None, dict[str, str] | None]: + ) -> tuple[httpx2.Auth | None, dict[str, str] | None]: """Resolve a v2-owned server's upstream credential into ``(resolved_auth, extra_headers)``. On a missing/rejected per-user credential this raises the mode's discovery challenge @@ -5552,7 +5552,7 @@ class MCPServerManager: verbose_logger.error(error_msg) return CallToolResult( content=[TextContent(type="text", text=error_msg)], - isError=True, + is_error=True, ) try: @@ -5563,7 +5563,7 @@ class MCPServerManager: # Convert the handler result (string response) to CallToolResult format result: Final = CallToolResult( content=[TextContent(type="text", text=str(handler_result))], - isError=False, + is_error=False, ) return result @@ -5579,7 +5579,7 @@ class MCPServerManager: verbose_logger.error(error_msg) return CallToolResult( content=[TextContent(type="text", text=error_msg)], - isError=True, + is_error=True, ) async def pre_call_tool_check( diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py index 43d97abe4db..3a8e2b3840a 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py @@ -34,6 +34,7 @@ from dataclasses import dataclass from typing import Annotated, Final, Literal import httpx +import httpx2 from pydantic import BaseModel, ConfigDict, Field, SecretStr, TypeAdapter, ValidationError from typing_extensions import assert_never @@ -337,7 +338,7 @@ def _identity_key(config: ClientCredentialsConfig) -> str: return hashlib.sha256(material.encode("utf-8")).hexdigest() -class ClientCredentialsBearerAuth(httpx.Auth): +class ClientCredentialsBearerAuth(httpx2.Auth): """Bearer auth that retries an upstream 401 exactly once with a freshly minted token. The initial token was already resolved (so config/IdP failures surfaced as typed errors @@ -356,7 +357,7 @@ class ClientCredentialsBearerAuth(httpx.Auth): self._access_token = SecretStr(access_token) self._refetch = refetch - async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]: + async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]: token: Final = self._access_token.get_secret_value() name, value = self._carrier.header(token) request.headers[name] = value @@ -371,5 +372,5 @@ class ClientCredentialsBearerAuth(httpx.Auth): request.headers[fresh_name] = fresh_value yield request - def sync_auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: - raise RuntimeError("ClientCredentialsBearerAuth only supports async httpx clients") + def sync_auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]: + raise RuntimeError("ClientCredentialsBearerAuth only supports async httpx2 clients") diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/httpx_auth.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/httpx_auth.py index e4d8fd25748..aa04469a502 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/httpx_auth.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/httpx_auth.py @@ -1,29 +1,29 @@ -"""Concrete `httpx.Auth` objects the resolver returns for the self-contained modes. +"""Concrete `httpx2.Auth` objects the resolver returns for the self-contained modes. -These are the egress credential as the SDK consumes it: an `httpx.Auth` attached to the +These are the egress credential as the SDK consumes it: an `httpx2.Auth` attached to the upstream `AsyncClient`. The OAuth-flow modes (`authorization_code`, `client_credentials`, `token_exchange`) return SDK-provided auth objects instead and land later. -`auth_flow` mutating the outbound request is the `httpx.Auth` contract, not a house-style -violation: the request is httpx's object, and these carry no state of their own. +`auth_flow` mutating the outbound request is the `httpx2.Auth` contract, not a house-style +violation: the request is httpx2's object, and these carry no state of their own. """ from __future__ import annotations from collections.abc import Generator -import httpx +import httpx2 from pydantic import SecretStr -class NoOpAuth(httpx.Auth): +class NoOpAuth(httpx2.Auth): """Attaches nothing — the `none` mode (and the seam-level default).""" - def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: + def auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]: yield request -class StaticHeaderAuth(httpx.Auth): +class StaticHeaderAuth(httpx2.Auth): """Sets one fixed header on every request — the `api_key` family and `passthrough`. The header value is a live credential (a bearer token, an API key, a forwarded user @@ -36,6 +36,6 @@ class StaticHeaderAuth(httpx.Auth): self.header_name = header_name self._header_value = SecretStr(header_value) - def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: + def auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]: request.headers[self.header_name] = self._header_value.get_secret_value() yield request diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 85c7f68719d..e71353e479c 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -1,7 +1,7 @@ """The one credential resolver: dispatch on the declared mode, fail closed. `resolve_credentials` selects exactly one arm off the server's typed `config` and either -produces an `httpx.Auth` or returns a typed `CredError`. The `match` is over the `AuthConfig` +produces an `httpx2.Auth` or returns a typed `CredError`. The `match` is over the `AuthConfig` variant, so each arm receives its own fully-typed config with no field-presence inference and no precedence cascade. It is wildcard-free with an `assert_never` tail, so adding a mode without an arm fails the type gate (basedpyright `reportMatchNotExhaustive`); a bypassed gate fails loudly @@ -25,6 +25,7 @@ from functools import partial from typing import Final import httpx +import httpx2 from typing_extensions import assert_never from litellm._logging import verbose_proxy_logger @@ -135,7 +136,7 @@ class UpstreamCredentialProvider: self._client_credentials_source = client_credentials_source or ClientCredentialsTokenSource() self._sso_assertion_store: SSOAssertionStore = sso_assertion_store or default_sso_assertion_store() - async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]: + async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx2.Auth, CredError]: match server.config: case NoneConfig(): return self._none(server) @@ -155,7 +156,7 @@ class UpstreamCredentialProvider: return _not_implemented(AuthSpecKind.aws_sigv4) assert_never(server.config) - def _none(self, server: ServerSpec) -> Result[httpx.Auth, CredError]: + def _none(self, server: ServerSpec) -> Result[httpx2.Auth, CredError]: try: resource: Final = httpx.URL(server.resource) except httpx.InvalidURL: @@ -169,12 +170,12 @@ class UpstreamCredentialProvider: Reads from the same per-user store as the ``authorization_code`` arm, so the discovery challenge and the egress agree on whether the user is authorized. Returns a typed ``bool`` - (no ``httpx.Auth``), unlike ``resolve_credentials``. A non-per-user mode has no token in the + (no ``httpx2.Auth``), unlike ``resolve_credentials``. A non-per-user mode has no token in the store, so it reads as False without a per-mode branch here. """ return await self._authz_token(subject, server) is not None - def _passthrough(self, subject: Subject) -> Result[httpx.Auth, CredError]: + def _passthrough(self, subject: Subject) -> Result[httpx2.Auth, CredError]: """Forward the caller's own upstream credential verbatim; the gateway mints nothing. The inbound token is the caller's already-disambiguated ``Authorization`` (never the LiteLLM @@ -186,7 +187,7 @@ class UpstreamCredentialProvider: return Ok(NoOpAuth()) return Ok(StaticHeaderAuth(subject.inbound_token.get_secret_value(), header_name="Authorization")) - def _api_key(self, config: ApiKeyConfig) -> Result[httpx.Auth, CredError]: + def _api_key(self, config: ApiKeyConfig) -> Result[httpx2.Auth, CredError]: match config.key_source: case SharedKey() as source: header_name, header_value = config.header(source.value.get_secret_value()) @@ -196,7 +197,9 @@ class UpstreamCredentialProvider: return Error(CredError.of_not_implemented("api_key BYOK source not implemented yet")) assert_never(config.key_source) - async def _id_jag(self, subject: Subject, server: ServerSpec, config: IdJagConfig) -> Result[httpx.Auth, CredError]: + async def _id_jag( + self, subject: Subject, server: ServerSpec, config: IdJagConfig + ) -> Result[httpx2.Auth, CredError]: match await self._id_jag_subject_token(subject): case Error(err): return Error(err) @@ -261,7 +264,7 @@ class UpstreamCredentialProvider: async def _id_jag_exchange( self, subject: Subject, token: str, server: ServerSpec, config: IdJagConfig - ) -> Result[httpx.Auth, CredError]: + ) -> Result[httpx2.Auth, CredError]: slot: Final = _id_jag_slot_key(subject, server) fingerprint: Final = _id_jag_fingerprint(token, server.server_id, config) @@ -313,7 +316,7 @@ class UpstreamCredentialProvider: async def _client_credentials( self, server_id: str, config: ClientCredentialsConfig - ) -> Result[httpx.Auth, CredError]: + ) -> Result[httpx2.Auth, CredError]: """The M2M arm: resolve a cached (or freshly minted) gateway token; no user context. The token is resolved here, before any upstream request, so a misconfigured grant or an @@ -448,7 +451,7 @@ def _client_auth_fingerprint(client_auth: ClientAuth) -> str: assert_never(client_auth) -def _not_implemented(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]: +def _not_implemented(kind: AuthSpecKind) -> Result[httpx2.Auth, CredError]: return Error(CredError.of_not_implemented(f"{kind.value}: resolver arm not implemented yet")) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index d186724fd9f..33c3a854058 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -30,7 +30,7 @@ from dataclasses import dataclass, field from enum import Enum from typing import Annotated, Final, Literal -import httpx +import httpx2 from expression import case, tag, tagged_union from pydantic import BaseModel, ConfigDict, Field, SecretStr, field_validator from typing_extensions import assert_never @@ -66,7 +66,7 @@ class AuthResolution(str, Enum): @dataclass(frozen=True, slots=True) class ResolvedCredential: - auth: httpx.Auth = field(repr=False) + auth: httpx2.Auth = field(repr=False) source: AuthResolution @@ -110,7 +110,7 @@ class Unauthorized: @tagged_union(frozen=True) class CredError: - """Why a credential could not be produced. Fail-closed: an arm yields this or an `httpx.Auth`. + """Why a credential could not be produced. Fail-closed: an arm yields this or an `httpx2.Auth`. Discriminated on the `Literal` `tag`; consumers `match self.tag` (see `summary`) so the type checker can prove exhaustiveness. Construct via the `of_*` factories. diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 9d895d755dd..15f97a15b73 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -10,6 +10,7 @@ from uuid import uuid4 import anyio import httpx +import httpx2 from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from pydantic import ValidationError from starlette.datastructures import Headers @@ -120,20 +121,29 @@ def _known_connection_error_message(exc: BaseException, url: str | None, timeout f"within {timeout_seconds:.0f}s. Check that the LiteLLM proxy can reach this URL " "from its network (DNS, egress rules, firewalls) and that the server answers MCP requests." ) - if isinstance(exc, httpx.LocalProtocolError): + if isinstance(exc, (httpx.LocalProtocolError, httpx2.LocalProtocolError)): return ( "Failed to connect to MCP server: a request header is malformed. " "Check static headers for leading/trailing spaces or illegal characters." ) - if isinstance(exc, (httpx.ConnectError, httpx.ConnectTimeout)): + if isinstance(exc, (httpx.ConnectError, httpx.ConnectTimeout, httpx2.ConnectError, httpx2.ConnectTimeout)): return ( "Failed to connect to MCP server: the server is unreachable. Check the URL and that the server is running." ) - if isinstance(exc, httpx.TimeoutException): + if isinstance(exc, (httpx.TimeoutException, httpx2.TimeoutException)): return "Failed to connect to MCP server: the connection timed out." - if isinstance(exc, httpx.HTTPStatusError): + if isinstance(exc, (httpx.HTTPStatusError, httpx2.HTTPStatusError)): return f"Failed to connect to MCP server: it returned HTTP {exc.response.status_code}." - if isinstance(exc, (httpx.NetworkError, httpx.RemoteProtocolError, ConnectionError)): + if isinstance( + exc, + ( + httpx.NetworkError, + httpx.RemoteProtocolError, + httpx2.NetworkError, + httpx2.RemoteProtocolError, + ConnectionError, + ), + ): return ( "Failed to connect to MCP server: the connection was interrupted. " "Check the server and network connection, then retry." @@ -148,7 +158,18 @@ def _known_connection_error_message(exc: BaseException, url: str | None, timeout "Failed to connect to MCP server: the endpoint returned invalid JSON or an invalid MCP response. " "Check the MCP endpoint URL and the server's protocol implementation." ) - if MCP_AVAILABLE and isinstance(exc, McpError): + if MCP_AVAILABLE and isinstance(exc, MCPError): + if exc.error.message.startswith("Unexpected content type:"): + return ( + "Failed to connect to MCP server: the endpoint returned an unsupported content type. " + "Check that the URL is an MCP endpoint, not a web page, and matches the selected transport." + ) + if exc.error.code == -32700 or exc.error.message.startswith("Failed to parse"): + return ( + f"Failed to connect to MCP server: the endpoint returned invalid JSON or an invalid MCP response " + f"(JSON-RPC code {exc.error.code}). " + "Check the MCP endpoint URL and the server's protocol implementation." + ) if exc.error.code == -32000 and exc.error.message == "Connection closed": return ( "Failed to connect to MCP server: the connection was closed before the request completed. " @@ -168,7 +189,7 @@ def _known_connection_error_message(exc: BaseException, url: str | None, timeout if MCP_AVAILABLE: - from mcp.shared.exceptions import McpError + from mcp.shared.exceptions import MCPError from mcp.types import Tool as MCPTool from litellm.experimental_mcp_client.client import MCPClient, as_mcp_read_timeout @@ -518,7 +539,7 @@ if MCP_AVAILABLE: ListMCPToolsRestAPIResponseObject( name=tool.name, description=tool.description, - inputSchema=tool.inputSchema, + inputSchema=tool.input_schema, mcp_info=enriched_mcp_info, ) for tool in tools @@ -1484,7 +1505,7 @@ if MCP_AVAILABLE: effective_timeout: Final = ( min(request.timeout if request.timeout is not None else MCP_CLIENT_TIMEOUT, timeout_seconds) if any( - isinstance(cause, McpError) and as_mcp_read_timeout(cause) is not None + isinstance(cause, MCPError) and as_mcp_read_timeout(cause) is not None for cause in iter_exception_tree(e) ) else timeout_seconds @@ -1635,7 +1656,7 @@ if MCP_AVAILABLE: "message": f"Timed out listing tools after {listing_deadline} seconds. " "The MCP server may be responding slowly or paginating excessively.", } - model_dumped_tools: Final[list[dict]] = [tool.model_dump() for tool in list_tools_result] + model_dumped_tools: Final[list[dict]] = [tool.model_dump(by_alias=True) for tool in list_tools_result] return { "tools": model_dumped_tools, "error": None, diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py index fec2a1f9ee6..361d8d5ae31 100644 --- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -18,8 +18,7 @@ if typing.TYPE_CHECKING: from collections.abc import Awaitable, Callable from fastapi import Request - from mcp.client.session import ClientSession - from mcp.shared.context import RequestContext + from mcp.client.session import ClientRequestContext from mcp.types import ( ContentBlock, CreateMessageResult, @@ -333,14 +332,14 @@ def _convert_single_content( return {"type": "text", "text": content.text} elif content_type == "image": image_data: Final[str] = getattr(content, "data", "") - image_mime_type: Final[str] = getattr(content, "mimeType", "image/png") + image_mime_type: Final[str] = getattr(content, "mime_type", "image/png") return { "type": "image_url", "image_url": {"url": f"data:{image_mime_type};base64,{image_data}"}, } elif content_type == "audio": audio_data: Final[str] = getattr(content, "data", "") - audio_mime_type: Final[str] = getattr(content, "mimeType", "audio/wav") + audio_mime_type: Final[str] = getattr(content, "mime_type", "audio/wav") # Map MIME type to OpenAI audio format format_map: Final = { "audio/wav": "wav", @@ -375,7 +374,7 @@ def _convert_single_content( # ToolResultContent → proper OpenAI tool-role message. # Marked so the message-level converter can emit it as a # separate ``{"role": "tool", ...}`` message. - tool_result_use_id: Final = getattr(content, "toolUseId", "") + tool_result_use_id: Final = getattr(content, "tool_use_id", "") nested_content: Final[Sequence[ContentBlock]] = getattr(content, "content", []) if isinstance(nested_content, list): text_parts = [getattr(c, "text", str(c)) for c in nested_content if getattr(c, "type", None) == "text"] @@ -538,7 +537,7 @@ def _extract_tool_results( results: Final = [] for item in items: if getattr(item, "type", None) == "tool_result": - tool_use_id = getattr(item, "toolUseId", "") + tool_use_id = getattr(item, "tool_use_id", "") # Extract text from nested content nested_content: Sequence[ContentBlock] = getattr(item, "content", []) if isinstance(nested_content, list): @@ -573,7 +572,7 @@ def _convert_mcp_tools_to_openai( "function": { "name": tool.name, "description": tool.description or "", - "parameters": tool.inputSchema + "parameters": tool.input_schema or { "type": "object", "properties": {}, @@ -718,7 +717,7 @@ def _convert_openai_response_to_mcp_result( role="assistant", content=content_parts, model=actual_model, - stopReason=stop_reason, + stop_reason=stop_reason, ) # Simple text response text: Final = message.content or "" @@ -726,7 +725,7 @@ def _convert_openai_response_to_mcp_result( role="assistant", content=TextContent(type="text", text=text), model=actual_model, - stopReason=stop_reason, + stop_reason=stop_reason, ) @@ -1066,21 +1065,21 @@ async def _build_completion_kwargs( ) -> dict[str, Any]: openai_messages: Final = _convert_mcp_messages_to_openai( messages=params.messages, - system_prompt=params.systemPrompt, + system_prompt=params.system_prompt, ) completion_kwargs: Final[dict[str, object]] = { "model": model, "messages": openai_messages, - "max_tokens": params.maxTokens, + "max_tokens": params.max_tokens, } if params.temperature is not None: completion_kwargs["temperature"] = params.temperature - if params.stopSequences: - completion_kwargs["stop"] = params.stopSequences + if params.stop_sequences: + completion_kwargs["stop"] = params.stop_sequences openai_tools: Final = _convert_mcp_tools_to_openai(params.tools) if openai_tools: completion_kwargs["tools"] = openai_tools - openai_tool_choice: Final = _convert_mcp_tool_choice_to_openai(params.toolChoice) + openai_tool_choice: Final = _convert_mcp_tool_choice_to_openai(params.tool_choice) if openai_tool_choice is not None: completion_kwargs["tool_choice"] = openai_tool_choice completion_kwargs["metadata"] = {"mcp_metadata": params.metadata} if params.metadata else {} @@ -1137,7 +1136,7 @@ async def _run_guardrails_and_call_llm( async def handle_sampling_create_message( - context: "RequestContext[ClientSession, object]", + context: "ClientRequestContext", params: "CreateMessageRequestParams", default_model: str | None = None, user_api_key_auth: "UserAPIKeyAuth | None" = None, @@ -1180,13 +1179,13 @@ async def handle_sampling_create_message( try: model: Final = _resolve_model_from_preferences( - model_preferences=params.modelPreferences, + model_preferences=params.model_preferences, default_model=default_model, ) verbose_logger.info( "MCP sampling: resolved model=%s from preferences=%s", model, - params.modelPreferences, + params.model_preferences, ) access_denial: Final = await _check_model_access(model, user_api_key_auth) @@ -1228,7 +1227,7 @@ async def handle_sampling_create_message( verbose_logger.info( "MCP sampling: completed successfully, model=%s, stopReason=%s", getattr(result, "model", "unknown"), - getattr(result, "stopReason", "unknown"), + getattr(result, "stop_reason", "unknown"), ) return result except Exception as e: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 73366b701d2..4ea22ca1f01 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -15,13 +15,13 @@ import traceback import types import uuid from collections import Counter -from collections.abc import AsyncIterator, Callable, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Iterable, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol import httpx from fastapi import FastAPI, HTTPException -from pydantic import AnyUrl, ConfigDict, TypeAdapter, ValidationError +from pydantic import AnyUrl, ConfigDict, Field, TypeAdapter, ValidationError from starlette.requests import Request as StarletteRequest from starlette.responses import JSONResponse from starlette.types import Message, Receive, Scope, Send @@ -64,6 +64,8 @@ from litellm.proxy._experimental.mcp_server.mcp_context import ( _mcp_gateway_initialize_instructions, _mcp_gateway_server_name, _mcp_proxy_mode, # pyright: ignore[reportPrivateUsage] # server-owned request mode + active_mcp_request_ctx_var, + get_active_mcp_request_ctx, ) from litellm.proxy._experimental.mcp_server.mcp_debug import ( MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, @@ -137,6 +139,24 @@ _MCP_ROUTING_PEEK_MAX_BYTES: Final = 4096 # ASGI scope keys carrying OTel request state into a stateful MCP message handler. _MCP_TRANSPORT_SPAN_SCOPE_KEY: Final = "litellm_otel_transport_span" _MCP_DESTINATIONS_SCOPE_KEY: Final = "litellm_otel_request_destinations" +_MCP_PROTOCOL_VERSION_HEADER: Final = b"mcp-protocol-version" + + +def unsupported_protocol_version(scope: Scope) -> str | None: + """Return the unsupported ``MCP-Protocol-Version`` header value, if any. + + SDK 2's ``StreamableHTTPSessionManager`` routes any version outside + ``HANDSHAKE_PROTOCOL_VERSIONS`` to the modern single-exchange path, which + bypasses litellm's session/auth model, so the ASGI entry rejects it. + """ + headers: Final[Iterable[tuple[bytes, bytes]]] = scope.get("headers") or () + values: Final = tuple( + raw.decode("latin-1").strip() for key, raw in headers if key.lower() == _MCP_PROTOCOL_VERSION_HEADER + ) + for value in values: + if value and value not in HANDSHAKE_PROTOCOL_VERSIONS: + return value + return None async def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None: @@ -156,14 +176,12 @@ try: from mcp import ReadResourceResult, Resource from mcp.server import Server - from mcp.server.lowlevel.helper_types import ReadResourceContents from mcp.server.session import ServerSession as _McpServerSession from mcp.types import ( BlobResourceContents, GetPromptResult, ResourceTemplate, TextResourceContents, - Tool, ) # Robust auth lookup keyed by session_object. @@ -176,7 +194,6 @@ except ImportError as e: # so they will never be accessed at runtime BlobResourceContents = None GetPromptResult = None - ReadResourceContents = None ReadResourceResult = None Resource = None ResourceTemplate = None @@ -277,8 +294,8 @@ def _mcp_meta_trace_carrier(req_ctx: object) -> dict[str, str] | None: span's identity attribution. """ meta: Final = getattr(req_ctx, "meta", None) - extra: Final = getattr(meta, "model_extra", None) - if not isinstance(extra, dict): + extra: Final = meta if isinstance(meta, Mapping) else getattr(meta, "model_extra", None) + if not isinstance(extra, Mapping): return None carrier: Final = {key: extra[key] for key in ("traceparent", "tracestate") if isinstance(extra.get(key), str)} return carrier or None @@ -456,6 +473,7 @@ if MCP_AVAILABLE: AuthContextMiddleware, auth_context_var, ) + from mcp.server.context import ServerRequestContext from mcp.server.lowlevel.server import NotificationOptions from mcp.server.models import InitializationOptions @@ -464,14 +482,23 @@ if MCP_AVAILABLE: except ImportError: StreamableHTTPSessionManager = None from mcp.types import ( + INVALID_REQUEST, + CallToolRequestParams, CallToolResult, + GetPromptRequestParams, Implementation, InitializeRequest, + ListPromptsResult, + ListResourcesResult, + ListResourceTemplatesResult, ListToolsResult, + PaginatedRequestParams, Prompt, + ReadResourceRequestParams, TextContent, ) from mcp.types import Tool as MCPTool + from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import ( MCPAuthenticatedUser, @@ -520,46 +547,20 @@ if MCP_AVAILABLE: Object returned by the /tools/list REST API route. """ - mcp_info: MCPInfo | None = None + mcp_info: MCPInfo | None = Field(default=None, alias="mcp_info") model_config = ConfigDict(arbitrary_types_allowed=True) - def _normalize_resource_contents(contents: list) -> list[ReadResourceContents]: - """Normalize ResourceContents to ReadResourceContents, preserving meta (MCP 1.26.0+).""" - normalized: Final[list[ReadResourceContents]] = [] - for content in contents: - meta = getattr(content, "meta", None) - if meta is None and hasattr(content, "model_dump"): - d = content.model_dump() - meta = d.get("meta") - if meta is None: - meta = d.get("_meta") - if isinstance(content, TextResourceContents): - normalized.append( - ReadResourceContents( - content=content.text, - mime_type=content.mimeType, - meta=meta, - ) - ) - elif isinstance(content, BlobResourceContents): - normalized.append( - ReadResourceContents( - content=content.blob, - mime_type=content.mimeType, - meta=meta, - ) - ) - return normalized - def _gateway_create_initialization_options( self, notification_options: NotificationOptions | None = None, experimental_capabilities: dict[str, dict[str, object]] | None = None, + extensions: dict[str, dict[str, object]] | None = None, ) -> InitializationOptions: base_options: Final = Server.create_initialization_options( self, notification_options=notification_options, experimental_capabilities=experimental_capabilities or {}, + extensions=extensions, ) opts: Final = ( base_options.model_copy( @@ -817,8 +818,7 @@ if MCP_AVAILABLE: ############### MCP Server Routes ####################### ######################################################## - @server.list_tools() - async def handle_list_tools() -> "ListToolsResult | list[Tool]": + async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListToolsResult: """ List all available tools, with each server's listing outcome attached to the result's ``_meta`` (SERVER_OUTCOMES_META_KEY) so a broken upstream is distinguishable from a healthy @@ -826,12 +826,9 @@ if MCP_AVAILABLE: pass the result through unwrapped, which is what lets the ``_meta`` survive to the client. Also captures the active session for propagation to callbacks. """ - from mcp.server.lowlevel.server import request_ctx - - req_ctx: Final = request_ctx.get(None) - _session_reset_token = None - if req_ctx: - _session_reset_token = active_mcp_session_var.set(req_ctx.session) + req_ctx: Final = ctx + _ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx) + _session_reset_token: Final = active_mcp_session_var.set(ctx.session) _trace_token = None _transport_token = None _destinations_token = None @@ -864,13 +861,13 @@ if MCP_AVAILABLE: ) if _mcp_proxy_mode.get(): - return [Tool.model_validate(d) for d in get_mcp_proxy_tool_definitions()] # mutable-ok: MCP SDK list + return ListToolsResult(tools=[Tool.model_validate(d) for d in get_mcp_proxy_tool_definitions()]) if getattr( getattr(user_api_key_auth, "object_permission", None), "mcp_tool_search_enabled", False, ): - return [Tool.model_validate(d) for d in get_virtual_tool_definitions()] + return ListToolsResult(tools=[Tool.model_validate(d) for d in get_virtual_tool_definitions()]) # Get mcp_servers from context variable verbose_logger.debug("MCP list_tools - Calling _list_mcp_tools") @@ -886,7 +883,7 @@ if MCP_AVAILABLE: ) verbose_logger.info("MCP list_tools - Successfully returned %s tools", len(listing.tools)) if not listing.outcomes: - return listing.tools + return ListToolsResult(tools=listing.tools) outcome_meta: Final = { SERVER_OUTCOMES_META_KEY: { key: outcome_wire_value(outcome) for key, outcome in listing.outcomes.items() @@ -894,36 +891,32 @@ if MCP_AVAILABLE: } return ListToolsResult.model_validate({"tools": listing.tools, "_meta": outcome_meta}) except HTTPException as e: - from mcp.shared.exceptions import McpError - from mcp.types import INVALID_REQUEST, ErrorData + from mcp.shared.exceptions import MCPError + from mcp.types import INVALID_REQUEST - raise McpError(ErrorData(code=INVALID_REQUEST, message=_http_detail_message(e.detail))) from e + raise MCPError(code=INVALID_REQUEST, message=_http_detail_message(e.detail)) from e except Exception as e: verbose_logger.exception("Error in list_tools endpoint: %s", e) # Return empty list instead of failing completely # This prevents the HTTP stream from failing and allows the client to get a response - return [] + return ListToolsResult(tools=[]) # mutable-ok: MCP result payload finally: _otel_reset_mcp_request_destinations(_destinations_token) _otel_reset_mcp_transport_span(_transport_token) _otel_reset_mcp_trace_carrier(_trace_token) - if _session_reset_token is not None: - active_mcp_session_var.reset(_session_reset_token) + active_mcp_session_var.reset(_session_reset_token) + active_mcp_request_ctx_var.reset(_ctx_reset_token) - def _capture_host_progress_callback(host_server) -> Callable | None: + def _capture_host_progress_callback(ctx: ServerRequestContext) -> Callable | None: """Return a progress-forwarding callback bound to the host MCP session. Returns ``None`` when the host did not supply a progress token. """ - try: - host_ctx: Final = host_server.request_context - except Exception as e: - verbose_logger.warning("Could not capture host progress context: %s", e) - return None + host_ctx: Final = ctx if not (host_ctx and hasattr(host_ctx, "meta") and host_ctx.meta): return None - host_token: Final = getattr(host_ctx.meta, "progressToken", None) + host_token: Final = host_ctx.meta.get("progress_token") if host_token is None or not (hasattr(host_ctx, "session") and host_ctx.session): return None host_session: Final = host_ctx.session @@ -944,10 +937,10 @@ if MCP_AVAILABLE: return forward_progress def _reject_mcp_proxy_operation() -> NoReturn: - from mcp.shared.exceptions import McpError - from mcp.types import METHOD_NOT_FOUND, ErrorData + from mcp.shared.exceptions import MCPError + from mcp.types import METHOD_NOT_FOUND - raise McpError(ErrorData(code=METHOD_NOT_FOUND, message="Operation unavailable on /mcp/proxy")) + raise MCPError(code=METHOD_NOT_FOUND, message="Operation unavailable on /mcp/proxy") async def _build_virtual_call_logging_obj( name: str, @@ -1022,7 +1015,7 @@ if MCP_AVAILABLE: content=[ # mutable-ok: MCP result content TextContent(type="text", text=f"Tool {name} is unavailable on /mcp/proxy") ], - isError=True, + is_error=True, ) if _mcp_proxy_mode.get() and name in MCP_PROXY_TOOL_NAMES: @@ -1104,7 +1097,7 @@ if MCP_AVAILABLE: text=f"Tool {name} requires mcp_tool_search_enabled on the key", ) ], - isError=True, + is_error=True, ) args: Final = arguments or {} @@ -1154,29 +1147,24 @@ if MCP_AVAILABLE: litellm_logging_obj=virtual_logging_obj, ) - @server.call_tool() - async def mcp_server_tool_call(name: str, arguments: dict[str, object] | None) -> CallToolResult: + async def mcp_server_tool_call(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult: """ Call a specific tool with the provided arguments Args: - name (str): Name of the tool to call - arguments (Dict[str, Any] | None): Arguments to pass to the tool + ctx: SDK request context carrying the client session and HTTP request + params (CallToolRequestParams): Tool name and arguments Returns: - List[Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]]: Tool execution results - Raises: - HTTPException: If tool not found or arguments missing + CallToolResult: Tool execution results """ - from mcp.server.lowlevel.server import request_ctx from mcp.types import CallToolResult from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.proxy_server import proxy_config - req_ctx: Final = request_ctx.get(None) - _session_reset_token = None - if req_ctx: - _session_reset_token = active_mcp_session_var.set(req_ctx.session) + req_ctx: Final = ctx + _ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx) + _session_reset_token: Final = active_mcp_session_var.set(ctx.session) _trace_token = None _transport_token = None _destinations_token = None @@ -1207,8 +1195,8 @@ if MCP_AVAILABLE: # Inside this try so virtual-tool errors convert to isError # CallToolResult instead of raising out of the protocol handler. virtual_tool_result: Final = await _dispatch_virtual_mcp_tool( - name=name, - arguments=arguments, + name=params.name, + arguments=params.arguments, user_api_key_auth=user_api_key_auth, client_ip=_client_ip, mcp_servers=mcp_servers, @@ -1220,9 +1208,9 @@ if MCP_AVAILABLE: if virtual_tool_result is not None: return virtual_tool_result - host_progress_callback: Final = _capture_host_progress_callback(server) + host_progress_callback: Final = _capture_host_progress_callback(ctx) # Create a body date for logging - body_data: Final = {"name": name, "arguments": arguments} + body_data: Final = {"name": params.name, "arguments": params.arguments} # mutable-ok: logging payload # Set trace/session id from raw_headers so spend logs and logging_obj stay consistent (same as A2A) chain_id: Final = get_chain_id_from_headers(raw_headers) if chain_id: @@ -1247,7 +1235,7 @@ if MCP_AVAILABLE: # Authorization is unaffected: it ran before this, and the union is resolved # from the untouched auth object passed to call_mcp_tool below. user_api_key_dict=await MCPRequestHandler.billing_auth_for_tool_call( - user_api_key_auth, tool_name=name + user_api_key_auth, tool_name=params.name ), proxy_config=proxy_config, ) @@ -1273,7 +1261,7 @@ if MCP_AVAILABLE: ) return CallToolResult( content=[TextContent(text=str(e), type="text")], - isError=True, + is_error=True, ) except BlockedPiiEntityError as e: verbose_logger.error("BlockedPiiEntityError in MCP tool call: %s", e) @@ -1284,19 +1272,19 @@ if MCP_AVAILABLE: type="text", ) ], - isError=True, + is_error=True, ) except GuardrailRaisedException as e: verbose_logger.error("GuardrailRaisedException in MCP tool call: %s", e) return CallToolResult( content=[TextContent(text=f"Error: Guardrail violation - {e}", type="text")], - isError=True, + is_error=True, ) except HTTPException as e: verbose_logger.error("HTTPException in MCP tool call: %s", e) return CallToolResult( content=[TextContent(text=f"Error: {_http_detail_message(e.detail)}", type="text")], - isError=True, + is_error=True, ) except MCPUpstreamAuthError as e: # The MCP session manager serializes handler exceptions as JSON-RPC errors, so a @@ -1312,13 +1300,13 @@ if MCP_AVAILABLE: type="text", ) ], - isError=True, + is_error=True, ) except Exception as e: verbose_logger.exception("MCP mcp_server_tool_call - error: %s", e) return CallToolResult( content=[TextContent(text=f"Error: {e}", type="text")], - isError=True, + is_error=True, ) return response @@ -1326,22 +1314,17 @@ if MCP_AVAILABLE: _otel_reset_mcp_request_destinations(_destinations_token) _otel_reset_mcp_transport_span(_transport_token) _otel_reset_mcp_trace_carrier(_trace_token) - if _session_reset_token is not None: - active_mcp_session_var.reset(_session_reset_token) + active_mcp_session_var.reset(_session_reset_token) + active_mcp_request_ctx_var.reset(_ctx_reset_token) - @server.list_prompts() - async def list_prompts() -> list[Prompt]: + async def list_prompts(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListPromptsResult: """ List all available prompts """ if _mcp_proxy_mode.get(): _reject_mcp_proxy_operation() - from mcp.server.lowlevel.server import request_ctx - - req_ctx: Final = request_ctx.get(None) - _session_reset_token = None - if req_ctx: - _session_reset_token = active_mcp_session_var.set(req_ctx.session) + _ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx) + _session_reset_token: Final = active_mcp_session_var.set(ctx.session) try: # Get user authentication from context variable @@ -1371,36 +1354,24 @@ if MCP_AVAILABLE: raw_headers=raw_headers, ) verbose_logger.info("MCP list_prompts - Successfully returned %s prompts", len(prompts)) - return prompts + return ListPromptsResult(prompts=prompts) except Exception as e: verbose_logger.exception("Error in list_prompts endpoint: %s", e) # Return empty list instead of failing completely # This prevents the HTTP stream from failing and allows the client to get a response - return [] + return ListPromptsResult(prompts=[]) # mutable-ok: MCP result payload finally: - if _session_reset_token is not None: - active_mcp_session_var.reset(_session_reset_token) + active_mcp_session_var.reset(_session_reset_token) + active_mcp_request_ctx_var.reset(_ctx_reset_token) - @server.get_prompt() - async def get_prompt(name: str, arguments: dict[str, str] | None) -> GetPromptResult: + async def get_prompt(ctx: ServerRequestContext, params: GetPromptRequestParams) -> GetPromptResult: """ Get a specific prompt with the provided arguments - - Args: - name (str): Name of the prompt to get - arguments (Dict[str, Any] | None): Arguments to pass to the prompt - - Returns: - GetPromptResult: Getting prompt execution results """ if _mcp_proxy_mode.get(): _reject_mcp_proxy_operation() - from mcp.server.lowlevel.server import request_ctx - - req_ctx: Final = request_ctx.get(None) - _session_reset_token = None - if req_ctx: - _session_reset_token = active_mcp_session_var.set(req_ctx.session) + _ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx) + _session_reset_token: Final = active_mcp_session_var.set(ctx.session) try: ( @@ -1415,8 +1386,8 @@ if MCP_AVAILABLE: verbose_logger.debug("MCP mcp_server_tool_call - User API Key Auth from context: %s", user_api_key_auth) return await mcp_get_prompt( - name=name, - arguments=arguments, + name=params.name, + arguments=params.arguments, user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_servers=mcp_servers, @@ -1425,20 +1396,15 @@ if MCP_AVAILABLE: raw_headers=raw_headers, ) finally: - if _session_reset_token is not None: - active_mcp_session_var.reset(_session_reset_token) + active_mcp_session_var.reset(_session_reset_token) + active_mcp_request_ctx_var.reset(_ctx_reset_token) - @server.list_resources() - async def list_resources() -> list[Resource]: + async def list_resources(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListResourcesResult: """List all available resources.""" if _mcp_proxy_mode.get(): _reject_mcp_proxy_operation() - from mcp.server.lowlevel.server import request_ctx - - req_ctx: Final = request_ctx.get(None) - _session_reset_token = None - if req_ctx: - _session_reset_token = active_mcp_session_var.set(req_ctx.session) + _ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx) + _session_reset_token: Final = active_mcp_session_var.set(ctx.session) try: ( @@ -1466,25 +1432,22 @@ if MCP_AVAILABLE: raw_headers=raw_headers, ) verbose_logger.info("MCP list_resources - Successfully returned %s resources", len(resources)) - return resources + return ListResourcesResult(resources=resources) except Exception as e: verbose_logger.exception("Error in list_resources endpoint: %s", e) - return [] + return ListResourcesResult(resources=[]) # mutable-ok: MCP result payload finally: - if _session_reset_token is not None: - active_mcp_session_var.reset(_session_reset_token) + active_mcp_session_var.reset(_session_reset_token) + active_mcp_request_ctx_var.reset(_ctx_reset_token) - @server.list_resource_templates() - async def list_resource_templates() -> list[ResourceTemplate]: + async def list_resource_templates( + ctx: ServerRequestContext, params: PaginatedRequestParams + ) -> ListResourceTemplatesResult: """List all available resource templates.""" if _mcp_proxy_mode.get(): _reject_mcp_proxy_operation() - from mcp.server.lowlevel.server import request_ctx - - req_ctx: Final = request_ctx.get(None) - _session_reset_token = None - if req_ctx: - _session_reset_token = active_mcp_session_var.set(req_ctx.session) + _ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx) + _session_reset_token: Final = active_mcp_session_var.set(ctx.session) try: ( @@ -1514,24 +1477,19 @@ if MCP_AVAILABLE: verbose_logger.info( "MCP list_resource_templates - Successfully returned %s resource templates", len(resource_templates) ) - return resource_templates + return ListResourceTemplatesResult(resource_templates=resource_templates) except Exception as e: verbose_logger.exception("Error in list_resource_templates endpoint: %s", e) - return [] + return ListResourceTemplatesResult(resource_templates=[]) # mutable-ok: MCP result payload finally: - if _session_reset_token is not None: - active_mcp_session_var.reset(_session_reset_token) + active_mcp_session_var.reset(_session_reset_token) + active_mcp_request_ctx_var.reset(_ctx_reset_token) - @server.read_resource() - async def read_resource(url: AnyUrl) -> list[ReadResourceContents]: + async def read_resource(ctx: ServerRequestContext, params: ReadResourceRequestParams) -> ReadResourceResult: if _mcp_proxy_mode.get(): _reject_mcp_proxy_operation() - from mcp.server.lowlevel.server import request_ctx - - req_ctx: Final = request_ctx.get(None) - _session_reset_token = None - if req_ctx: - _session_reset_token = active_mcp_session_var.set(req_ctx.session) + _ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx) + _session_reset_token: Final = active_mcp_session_var.set(ctx.session) try: ( @@ -1545,7 +1503,7 @@ if MCP_AVAILABLE: ) = await get_or_extract_auth_context() read_resource_result: Final = await mcp_read_resource( - url=url, + url=params.uri, user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_servers=mcp_servers, @@ -1554,10 +1512,18 @@ if MCP_AVAILABLE: raw_headers=raw_headers, ) - return _normalize_resource_contents(read_resource_result.contents) + return read_resource_result finally: - if _session_reset_token is not None: - active_mcp_session_var.reset(_session_reset_token) + active_mcp_session_var.reset(_session_reset_token) + active_mcp_request_ctx_var.reset(_ctx_reset_token) + + server.add_request_handler("tools/list", PaginatedRequestParams, handle_list_tools) + server.add_request_handler("tools/call", CallToolRequestParams, mcp_server_tool_call) + server.add_request_handler("prompts/list", PaginatedRequestParams, list_prompts) + server.add_request_handler("prompts/get", GetPromptRequestParams, get_prompt) + server.add_request_handler("resources/list", PaginatedRequestParams, list_resources) + server.add_request_handler("resources/templates/list", PaginatedRequestParams, list_resource_templates) + server.add_request_handler("resources/read", ReadResourceRequestParams, read_resource) ######################################################## ############ End of MCP Server Routes ################## @@ -3296,11 +3262,11 @@ if MCP_AVAILABLE: Guardrails run before the success/failure logging so the masked text, not the raw one, is what gets logged. - A result with ``isError=True`` is logged as a failure (``status="failure"`` + A result with ``is_error=True`` is logged as a failure (``status="failure"`` payload, so OTel marks the span ERROR) while the HTTP wire behavior stays 200 + ``isError: true`` per the MCP spec. The error check runs after ``async_post_mcp_tool_call_hook`` because guardrails may flip the result - to ``isError=True`` in that hook. Raised exceptions never reach here (the + to ``is_error=True`` in that hook. Raised exceptions never reach here (the ``@client`` wrapper and ``call_mcp_tool``'s except path log those), so this cannot double-log a failure. @@ -3635,10 +3601,10 @@ if MCP_AVAILABLE: """Execute a local-registry tool and report whether it succeeded. Returns the result rather than bare content because the verdict is part of it: the content - alone cannot say whether the handler failed, so callers used to stamp isError=False on every + alone cannot say whether the handler failed, so callers used to stamp is_error=False on every outcome and an upstream rejection was served as tool output. - A failure is reported as ``isError=True`` here rather than raised, because the REST surface + A failure is reported as ``is_error=True`` here rather than raised, because the REST surface turns an unrecognized exception into a 500 and an upstream 403 or 429 is not a gateway crash. ``MCPUpstreamAuthError`` is the exception: it propagates so the caller is told to re-authenticate, which both renderers already know how to say. @@ -3660,8 +3626,14 @@ if MCP_AVAILABLE: raise except Exception as e: verbose_logger.exception("Error executing local tool %s: %s", name, e) - return CallToolResult(content=[TextContent(text=f"Error: {e}", type="text")], isError=True) - return CallToolResult(content=[TextContent(text=str(result), type="text")], isError=False) + return CallToolResult( + content=[TextContent(text=f"Error: {e}", type="text")], # mutable-ok: MCP result content + is_error=True, + ) + return CallToolResult( + content=[TextContent(text=str(result), type="text")], # mutable-ok: MCP result content + is_error=False, + ) def _get_mcp_servers_in_path(path: str) -> list[str] | None: """ @@ -3843,7 +3815,7 @@ if MCP_AVAILABLE: def _extract_initialize_client_info(body: bytes) -> Implementation | None: try: - return InitializeRequest.model_validate_json(body).params.clientInfo + return InitializeRequest.model_validate_json(body, by_name=False).params.client_info except ValidationError: return None @@ -4553,6 +4525,21 @@ if MCP_AVAILABLE: async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through StreamableHTTP.""" try: + bad_version: Final = unsupported_protocol_version(scope) + if bad_version is not None: + supported: Final = ", ".join(sorted(HANDSHAKE_PROTOCOL_VERSIONS)) + await JSONResponse( + status_code=400, + content={ # mutable-ok: JSON-RPC error payload + "jsonrpc": "2.0", + "id": None, + "error": { + "code": INVALID_REQUEST, + "message": f"Unsupported MCP-Protocol-Version {bad_version}; supported: {supported}", + }, + }, + )(scope, receive, send) + return path: Final[str] = scope.get("path", "") ( user_api_key_auth, @@ -5179,12 +5166,8 @@ if MCP_AVAILABLE: return None, None, None, None, None, None, None def _get_current_session(): - try: - from mcp.server.lowlevel.server import request_ctx - - return request_ctx.get().session - except (LookupError, ImportError): - return None + ctx: Final = get_active_mcp_request_ctx() + return ctx.session if ctx is not None else None def _cache_auth_context_lazily(): session: Final = _get_current_session() diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index e921ab0331e..a482d02c31d 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -99,11 +99,20 @@ def mcp_tool_search_settings() -> MCPToolSearchSettings | ValidationError: def _tool_result(tool: Tool) -> ToolSearchResult: - return {"name": tool.name, "description": tool.description or "", "inputSchema": tool.inputSchema} + return { + "name": tool.name, + "description": tool.description or "", + "inputSchema": tool.input_schema, + } # mutable-ok: wire schema payload def _scored_result(tool: Tool, score: float) -> ToolSearchResult: - return {"name": tool.name, "description": tool.description or "", "inputSchema": tool.inputSchema, "score": score} + return { + "name": tool.name, + "description": tool.description or "", + "inputSchema": tool.input_schema, + "score": score, + } # mutable-ok: wire schema payload _MCP_PROXY_IDENTITY_META_KEY: Final[str] = "litellm.ai/proxy_tool_identity" @@ -148,11 +157,11 @@ def _proxy_schema_result(tool: Tool) -> MCPProxySchemaResult: "tool_id": mcp_proxy_tool_id(tool), "name": tool.name, "description": tool.description or "", - "inputSchema": tool.inputSchema, + "inputSchema": tool.input_schema, } - if tool.outputSchema is None: + if tool.output_schema is None: return base - return {**base, "outputSchema": tool.outputSchema} # mutable-ok: wire schema payload + return {**base, "outputSchema": tool.output_schema} # mutable-ok: wire schema payload def _tool_text(tool: Tool) -> str: @@ -372,7 +381,7 @@ def _text_tool_result(text: str, is_error: bool) -> CallToolResult: return CallToolResult( content=[TextContent(type="text", text=text)], # mutable-ok: CallToolResult accepts only list content - isError=is_error, + is_error=is_error, ) @@ -565,7 +574,7 @@ async def handle_mcp_proxy_tool( if not isinstance(tool_arguments, dict): return _text_tool_result("arguments must be an object", is_error=True) try: - validate(instance=tool_arguments, schema=tool.inputSchema) + validate(instance=tool_arguments, schema=tool.input_schema) except JsonSchemaValidationError as exc: return _text_tool_result(f"Invalid arguments: {exc.message}", is_error=True) diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index fb3eb06fd15..6bd080f5216 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -536,7 +536,11 @@ def extract_mcp_tool_result_error_message(result: object) -> str | None: Accepts both ``mcp.types.CallToolResult`` objects and their dict equivalents, duck-typed so the ``mcp`` package is not required. """ - is_error: Final[object] = result.get("isError") if isinstance(result, Mapping) else getattr(result, "isError", None) + is_error: Final[object] = ( + (result.get("isError") if result.get("isError") is not None else result.get("is_error")) + if isinstance(result, Mapping) + else getattr(result, "is_error", None) + ) if is_error is not True: return None content: Final[object] = result.get("content") if isinstance(result, Mapping) else getattr(result, "content", None) @@ -870,8 +874,9 @@ def json_unrewritable_labels(value: object, path_depth: int = 0) -> tuple[str, . def mcp_tool_result_structured_content(result: object) -> object: """The ``structuredContent`` of an MCP tool result, or ``None`` when it has none.""" if isinstance(result, Mapping): - return result.get("structuredContent") - return getattr(result, "structuredContent", None) + structured: Final = result.get("structuredContent") + return structured if structured is not None else result.get("structured_content") + return getattr(result, "structured_content", None) def set_mcp_tool_result_structured_content(result: object, value: object) -> bool: @@ -882,12 +887,12 @@ def set_mcp_tool_result_structured_content(result: object, value: object) -> boo unmasked value in the spend log and the OTel span. """ if isinstance(result, MutableMapping): - result["structuredContent"] = value + result["structured_content" if "structured_content" in result else "structuredContent"] = value return True - if not hasattr(result, "structuredContent"): + if not hasattr(result, "structured_content"): return False try: - setattr(result, "structuredContent", value) # attribute name is fixed by the MCP result shape + setattr(result, "structured_content", value) # attribute name is fixed by the MCP result shape return True except (AttributeError, TypeError, ValueError): return False diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 4cfb2bf8c38..06e157498aa 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3247,6 +3247,17 @@ ], "title": "Key Alias" }, + "key_exists": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Key Exists" + }, "team_id": { "anyOf": [ { diff --git a/litellm/proxy/config_resolvers/__init__.py b/litellm/proxy/config_resolvers/__init__.py index ebd339b34c3..eee760df458 100644 --- a/litellm/proxy/config_resolvers/__init__.py +++ b/litellm/proxy/config_resolvers/__init__.py @@ -5,6 +5,6 @@ from litellm.proxy.config_resolvers._descriptors import ( FieldSource, resolve_fields, ) -from litellm.proxy.config_resolvers.settings_store import SettingsStore +from litellm.proxy.config_resolvers.settings_store import SettingsStore, config_ownership_message -__all__ = ("FieldDescriptor", "FieldSource", "SettingsStore", "resolve_fields") +__all__ = ("FieldDescriptor", "FieldSource", "SettingsStore", "config_ownership_message", "resolve_fields") diff --git a/litellm/proxy/config_resolvers/settings_store.py b/litellm/proxy/config_resolvers/settings_store.py index 345f00c35a5..90f1da76bf6 100644 --- a/litellm/proxy/config_resolvers/settings_store.py +++ b/litellm/proxy/config_resolvers/settings_store.py @@ -19,10 +19,21 @@ from litellm.proxy.config_resolvers.settings_rules import ( class ConfigOwnedKeyError(RuntimeError): - def __init__(self, section: Section, key: str) -> None: - super().__init__(f"{section}.{key} is set in the config file and cannot be changed at runtime") + def __init__(self, section: Section, key: str, *, shadows_db_value: bool = False) -> None: + super().__init__(config_ownership_message(section=section, key=key, shadows_db_value=shadows_db_value)) self.section: Final = section self.key: Final = key + self.shadows_db_value: Final = shadows_db_value + + +def config_ownership_message(*, section: Section, key: str, shadows_db_value: bool) -> str: + stored: Final = ( + " The value stored in the database for it is ignored and will never be applied." if shadows_db_value else "" + ) + return ( + f"{section}.{key} is set in the config file, so the config file owns it and it cannot be changed " + f"here.{stored} Edit the config file to change it, or remove it from the file to let the database own it." + ) _EMPTY_VALUES: Final[Mapping[str, JsonValue]] = MappingProxyType({}) @@ -54,6 +65,13 @@ class SettingsStore(MutableMapping[str, JsonValue]): ) ) + def shadowed_db_keys(self) -> tuple[str, ...]: + """Keys the config file owns whose stored value differs, so the stored one never reaches a reader.""" + return tuple(sorted(key for key in self._yaml_values if self._db_value_is_shadowed(key))) + + def shadows_db_value(self, key: str) -> bool: + return self.owned_by_config(key) and self._db_value_is_shadowed(key) + def apply_db_row(self, row: DbRow, db_row: Mapping[str, JsonValue]) -> None: previous_row: Final = self._database_rows.get(row, _EMPTY_VALUES) self._database_rows = MappingProxyType({**self._database_rows, row: MappingProxyType(dict(db_row))}) @@ -81,7 +99,7 @@ class SettingsStore(MutableMapping[str, JsonValue]): def __setitem__(self, key: str, value: JsonValue) -> None: if self.owned_by_config(key) and value != self.get(key): - raise ConfigOwnedKeyError(self._section, key) + raise ConfigOwnedKeyError(self._section, key, shadows_db_value=self._db_value_is_shadowed(key)) self._runtime_values = MappingProxyType({**self._runtime_values, key: value}) self._deleted_runtime_keys = self._deleted_runtime_keys - frozenset((key,)) @@ -89,7 +107,7 @@ class SettingsStore(MutableMapping[str, JsonValue]): if key not in self: raise KeyError(key) if self.owned_by_config(key): - raise ConfigOwnedKeyError(self._section, key) + raise ConfigOwnedKeyError(self._section, key, shadows_db_value=self._db_value_is_shadowed(key)) self._runtime_values = MappingProxyType( {key_: value for key_, value in self._runtime_values.items() if key_ != key} ) @@ -136,8 +154,14 @@ class SettingsStore(MutableMapping[str, JsonValue]): ) ) - def _resolution_for(self, key: str) -> Resolved: + def _db_value(self, key: str) -> SettingValue: rule: Final = rule_for(self._section, key) + return self._database_rows.get(rule.db_row, _EMPTY_VALUES).get(key, ABSENT) + + def _db_value_is_shadowed(self, key: str) -> bool: + db_value: Final = self._db_value(key) + return not isinstance(db_value, Absent) and db_value is not None and db_value != self.get(key) + + def _resolution_for(self, key: str) -> Resolved: yaml_value: Final[SettingValue] = self._yaml_values.get(key, ABSENT) - db_value: Final[SettingValue] = self._database_rows.get(rule.db_row, _EMPTY_VALUES).get(key, ABSENT) - return resolve(yaml_value, db_value) + return resolve(yaml_value, self._db_value(key)) diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py index 5a6be1089b6..67ef05fc324 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py @@ -34,15 +34,26 @@ def _serialize_mcp_content_item(item: object) -> dict[str, object]: model_dump: Final = getattr(item, "model_dump", None) if callable(model_dump): try: - return dict(model_dump(exclude_none=True)) + dumped: Final[dict[str, object]] = model_dump(exclude_none=True, by_alias=True) + return dict(dumped) except TypeError: - return dict(model_dump()) + dumped_fallback: Final[dict[str, object]] = model_dump() + return dict(dumped_fallback) text: Final = getattr(item, "text", None) if isinstance(text, str): return {"type": getattr(item, "type", "text"), "text": text} return {"type": "text", "text": str(item)} +def _source_field(source: object, key: str, snake_key: str) -> object: + if isinstance(source, dict): + for candidate in (key, snake_key): + if candidate in source: + return source[candidate] # pyright: ignore[reportUnknownVariableType] # dict-shaped sources arrive untyped + return None + return getattr(source, snake_key, None) + + class _CiscoAIDefenseMcpMixin: """MCP-specific instance methods for ``CiscoAIDefenseGuardrail``. @@ -219,14 +230,14 @@ class _CiscoAIDefenseMcpMixin: if isinstance(content, list): content[:] = replacement structured_replacement: Final = _CiscoAIDefenseMcpMixin._replacement_structured_content(replacement) - if hasattr(response_obj, "structuredContent"): + if hasattr(response_obj, "structured_content"): try: - setattr(response_obj, "structuredContent", structured_replacement) + setattr(response_obj, "structured_content", structured_replacement) except (AttributeError, TypeError, ValueError): pass - if hasattr(response_obj, "isError"): + if hasattr(response_obj, "is_error"): try: - setattr(response_obj, "isError", True) + setattr(response_obj, "is_error", True) except (AttributeError, TypeError, ValueError): pass return True @@ -487,7 +498,7 @@ class _CiscoAIDefenseMcpMixin: model_dump: Final = getattr(response, "model_dump", None) if callable(model_dump): try: - dumped = model_dump(exclude_none=True) + dumped = model_dump(exclude_none=True, by_alias=True) except TypeError: dumped = model_dump() if isinstance(dumped, dict): @@ -507,8 +518,8 @@ class _CiscoAIDefenseMcpMixin: source: object = None, ) -> dict[str, object]: result: Final[dict[str, object]] = {"content": [_serialize_mcp_content_item(item) for item in content]} - for key in ("structuredContent", "isError"): - value = source.get(key) if isinstance(source, dict) else getattr(source, key, None) + for key, snake_key in (("structuredContent", "structured_content"), ("isError", "is_error")): + value = _source_field(source, key, snake_key) if value is not None and (key != "isError" or isinstance(value, bool)): result[key] = value return result @@ -549,20 +560,21 @@ class _CiscoAIDefenseMcpMixin: and all(isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], str) for item in response_obj) ): for index, item in enumerate(response_obj): - if item[0] == "structuredContent": + if item[0] in ("structuredContent", "structured_content"): response_obj[index] = (item[0], replacement) replaced = True - elif hasattr(response_obj, "structuredContent"): + elif hasattr(response_obj, "structured_content"): try: - setattr(response_obj, "structuredContent", replacement) + setattr(response_obj, "structured_content", replacement) replaced = True except (AttributeError, TypeError, ValueError): pass elif isinstance(response_obj, dict): result: Final = response_obj.get("result") target: Final[dict[object, object]] = result if isinstance(result, dict) else response_obj - if "structuredContent" in target: - target["structuredContent"] = replacement + structured_key: Final = "structured_content" if "structured_content" in target else "structuredContent" + if structured_key in target: + target[structured_key] = replacement replaced = True return replaced diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index a1c92d37871..5a19d743105 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -127,6 +127,7 @@ class _KeyMetadataDict(TypedDict, total=False): team_id: ReadOnly[str | None] user_id: ReadOnly[str | None] user_email: ReadOnly[str | None] + key_exists: ReadOnly[bool] def _key_metadata(api_key_metadata: Mapping[str, _KeyMetadataDict], api_key: str) -> KeyMetadata: @@ -136,6 +137,7 @@ def _key_metadata(api_key_metadata: Mapping[str, _KeyMetadataDict], api_key: str team_id=meta.get("team_id"), user_id=meta.get("user_id"), user_email=meta.get("user_email"), + key_exists=meta.get("key_exists", False), ) @@ -512,6 +514,7 @@ async def get_api_key_metadata( "key_alias": k.key_alias, "team_id": k.team_id, "user_id": getattr(k, "user_id", None), + "key_exists": True, } for k in key_records } diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 4554a85b225..40bd496fdce 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1968,7 +1968,7 @@ async def generate_key_fn( - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules. - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely. - - enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Anthropic and Bedrock Claude models only. + - enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Supported Claude models on Anthropic, Bedrock, Vertex AI, and Azure AI only. - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. @@ -3323,7 +3323,7 @@ async def update_key_fn( - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules. - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely. - - enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Anthropic and Bedrock Claude models only. + - enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Supported Claude models on Anthropic, Bedrock, Vertex AI, and Azure AI only. - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. - blocked: Optional[bool] - Whether the key is blocked - aliases: Optional[dict] - Model aliases for the key - [Docs](https://litellm.vercel.app/docs/proxy/virtual_keys#model-aliases) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6741b8b56c2..3c7d06268ad 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -447,7 +447,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( project_spend_counter_key, tag_cache_key, ) -from litellm.proxy.config_resolvers import SettingsStore, resolve_fields +from litellm.proxy.config_resolvers import SettingsStore, config_ownership_message, resolve_fields from litellm.proxy.config_resolvers.alerting import ( EMAIL_DESCRIPTORS, MS_TEAMS_DESCRIPTORS, @@ -4907,6 +4907,7 @@ class ProxyConfig: self.router_settings: Final[SettingsStore] = SettingsStore("router_settings") self.litellm_settings: Final[SettingsStore] = SettingsStore("litellm_settings") self.environment_variables: Final[SettingsStore] = SettingsStore("environment_variables") + self._warned_shadowed_keys: frozenset[tuple[Section, str]] = frozenset() self._settings_stores: Final[Mapping[Section, SettingsStore]] = MappingProxyType( { "general_settings": self.settings, @@ -5128,12 +5129,20 @@ class ProxyConfig: f"key '{rejected[0]}' is" if len(rejected) == 1 else f"keys {', '.join(repr(key) for key in rejected)} are" ) pronoun: Final = "it" if len(rejected) == 1 else "them" + shadowed: Final = tuple(key for key in rejected if store.shadows_db_value(key)) + stored: Final = ( + f" The {'value' if len(shadowed) == 1 else 'values'} already stored in the database for " + f"{', '.join(shadowed)} {'is' if len(shadowed) == 1 else 'are'} ignored and will never be applied." + if shadowed + else "" + ) raise HTTPException( status_code=400, detail={ - "error": f"{section_name} {subject} set in the config file and cannot be changed here", + "error": f"{section_name} {subject} set in the config file and cannot be changed here.{stored}", "keys": list(rejected), "section": section_name, + "stored_database_values_ignored": list(shadowed), "resolution": ( f"edit {user_config_file_path} to change {pronoun}, " f"or remove {pronoun} from the file to let the database own {pronoun}" @@ -7430,8 +7439,19 @@ class ProxyConfig: self._prepared_db_settings_values(section, param_value), ) + self._warn_about_shadowed_db_settings() return self._config_with_resolved_settings(config) + def _warn_about_shadowed_db_settings(self) -> None: + shadowed: Final[frozenset[tuple[Section, str]]] = frozenset( + (section, key) for section, store in self._settings_stores.items() for key in store.shadowed_db_keys() + ) + for section, key in sorted(shadowed - self._warned_shadowed_keys): + verbose_proxy_logger.warning( + "%s", config_ownership_message(section=section, key=key, shadows_db_value=True) + ) + self._warned_shadowed_keys = shadowed + def _prepared_db_settings_values(self, section: Section, value: object) -> Mapping[str, SettingsJsonValue]: if section == "environment_variables": decrypted: Final = self._decrypt_and_set_db_env_variables( @@ -17688,8 +17708,8 @@ _GENERAL_SETTINGS_UI_LITELLM_FIELDS: Final[dict[str, GeneralSettingsUILiteLLMFie "type": "Boolean", "tab": "prompt_caching", "description": ( - "Auto-adds cache_control to the system prompt and trailing turn for supported Anthropic " - "and Bedrock Claude models. The cache is shared across callers on the same upstream credentials." + "Auto-adds cache_control to the system prompt and trailing turn for supported Claude models on " + "Anthropic, Bedrock, Vertex AI, and Azure AI. The cache is shared across callers on the same upstream credentials." ), }, "anthropic_prompt_caching_ttl": { diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index e6673ec99aa..d7e100dd630 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -209,6 +209,94 @@ ], "default_model_placeholder": "claude-3-opus" }, + { + "provider": "AWS_Textract", + "provider_display_name": "Amazon Textract", + "litellm_provider": "aws_textract", + "credential_fields": [ + { + "key": "aws_access_key_id", + "label": "AWS Access Key ID", + "placeholder": null, + "tooltip": "You can provide the raw key or the environment variable (e.g. `os.environ/MY_SECRET_KEY`).", + "required": false, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "aws_secret_access_key", + "label": "AWS Secret Access Key", + "placeholder": null, + "tooltip": "You can provide the raw key or the environment variable (e.g. `os.environ/MY_SECRET_KEY`).", + "required": false, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "aws_session_token", + "label": "AWS Session Token", + "placeholder": null, + "tooltip": "Temporary credentials session token. You can provide the raw token or the environment variable (e.g. `os.environ/MY_SESSION_TOKEN`).", + "required": false, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "aws_region_name", + "label": "AWS Region Name", + "placeholder": "us-east-1", + "tooltip": "You can provide the raw key or the environment variable (e.g. `os.environ/MY_SECRET_KEY`).", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "aws_session_name", + "label": "AWS Session Name", + "placeholder": "my-session", + "tooltip": "Name for the AWS session. You can provide the raw value or the environment variable (e.g. `os.environ/MY_SESSION_NAME`).", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "aws_profile_name", + "label": "AWS Profile Name", + "placeholder": "default", + "tooltip": "AWS profile name to use for authentication. You can provide the raw value or the environment variable (e.g. `os.environ/MY_PROFILE_NAME`).", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "aws_role_name", + "label": "AWS Role Name", + "placeholder": "MyRole", + "tooltip": "AWS IAM role name to assume. You can provide the raw value or the environment variable (e.g. `os.environ/MY_ROLE_NAME`).", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "aws_web_identity_token", + "label": "AWS Web Identity Token", + "placeholder": null, + "tooltip": "Web identity token for OIDC authentication. You can provide the raw token or the environment variable (e.g. `os.environ/MY_WEB_IDENTITY_TOKEN`).", + "required": false, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "detect-document-text" + }, { "provider": "BedrockMantle", "provider_display_name": "Amazon Bedrock Mantle", diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 3b6afc34063..73ab7e5213f 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -34,6 +34,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_set_request_parsed_body, ) +from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.types.llms.openai import ( REASONING_EFFORT, ResponsesAPIOptionalRequestParams, @@ -296,6 +297,7 @@ async def responses_api( route_type="aresponses", llm_router=llm_router, ) + raise_if_required_body_param_missing(route_type="aresponses", data=data) except Exception as e: raise await processor._handle_llm_api_exception( e=e, diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 20b4708c193..536c58df65a 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -159,6 +159,7 @@ class ProxyModelNotFoundError(HTTPException): REQUIRED_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, ...]]] = { "acompletion": ("messages",), "aembedding": ("input",), + "aresponses": ("input",), "acreate_batch": ("input_file_id", "endpoint", "completion_window"), } diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 29688b61b3d..ee2e1cfeaf7 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -69,6 +69,7 @@ class KeyMetadataDict(TypedDict, total=False): team_id: ReadOnly[str | None] user_id: ReadOnly[str | None] user_email: ReadOnly[str | None] + key_exists: ReadOnly[bool] class _TokenDigestRow(BaseModel): diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 7c2abce60e2..b2baef126e9 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -497,12 +497,10 @@ def _store_allowed_ips(general_settings: MutableMapping[str, object], allowed_ip raise HTTPException( status_code=400, detail={ # mutable-ok: HTTPException serializes its detail as json - "error": f"{owned.section} key '{owned.key}' is set in the config file and cannot be changed here", + "error": str(owned), "keys": (owned.key,), "section": owned.section, - "resolution": ( - "edit the config file to change it, or remove it from the file to let the database own it" - ), + "stored_database_value_ignored": owned.shadows_db_value, }, ) from owned diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 1b19bf77a7d..16e8ac93d59 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -105,8 +105,8 @@ async def create_mcp_list_tools_events( "description": getattr(tool, "description", ""), "annotations": {"read_only": False}, **dict.fromkeys( - ("input_schema",) if hasattr(tool, "inputSchema") or hasattr(tool, "input_schema") else (), - getattr(tool, "inputSchema", getattr(tool, "input_schema", None)), + ("input_schema",) if hasattr(tool, "input_schema") else (), + getattr(tool, "input_schema", None), ), } for tool in filtered_mcp_tools diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index af3d7ddfac7..79ea6dc36ec 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -18,6 +18,7 @@ import httpx import litellm from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( ITPM_RESERVED_KEY, @@ -136,6 +137,26 @@ class ModelRateLimitingCheck(CustomLogger): return tpm_key, rpm_key + def _get_current_tpm(self, tpm_key: str, tpm_limit: int) -> int | None: + local_tpm: Final = self.dual_cache.get_cache(key=tpm_key, local_only=True) + redis_cache: Final = self.dual_cache.redis_cache + if redis_cache is None or (local_tpm is not None and local_tpm >= tpm_limit): + return local_tpm + try: + return redis_cache.get_cache(key=tpm_key) + except RedisCircuitBreakerOpenError: + return local_tpm + + async def _async_get_current_tpm(self, tpm_key: str, tpm_limit: int, parent_otel_span: Span | None) -> int | None: + local_tpm: Final = await self.dual_cache.async_get_cache(key=tpm_key, local_only=True) + redis_cache: Final = self.dual_cache.redis_cache + if redis_cache is None or (local_tpm is not None and local_tpm >= tpm_limit): + return local_tpm + try: + return await redis_cache.async_get_cache(key=tpm_key, parent_otel_span=parent_otel_span) + except RedisCircuitBreakerOpenError: + return local_tpm + def pre_call_check(self, deployment: dict) -> dict | None: """ Synchronous pre-call check for model rate limits. @@ -168,8 +189,7 @@ class ModelRateLimitingCheck(CustomLogger): # Check TPM limit if tpm_limit is not None: - # First check local cache - current_tpm: Final = self.dual_cache.get_cache(key=tpm_key, local_only=True) + current_tpm: Final = self._get_current_tpm(tpm_key, tpm_limit) if current_tpm is not None and current_tpm >= tpm_limit: raise litellm.RateLimitError( message=f"Model rate limit exceeded. TPM limit={tpm_limit}, current usage={current_tpm}", @@ -249,8 +269,7 @@ class ModelRateLimitingCheck(CustomLogger): # Check TPM limit if tpm_limit is not None: - # First check local cache - current_tpm: Final = await self.dual_cache.async_get_cache(key=tpm_key, local_only=True) + current_tpm: Final = await self._async_get_current_tpm(tpm_key, tpm_limit, parent_otel_span) if current_tpm is not None and current_tpm >= tpm_limit: raise litellm.RateLimitError( message=f"Model rate limit exceeded. TPM limit={tpm_limit}, current usage={current_tpm}", diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index d843a874fe3..8794ff2db95 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -58,6 +58,7 @@ class Rule: Rules: TypeAlias = tuple[Rule, ...] RULES: Final[Rules] = ( + Rule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"aws_textract"})), Rule(Route.OCR, Rollout.RUST_OPT_OUT), Rule(Route.MESSAGES, Rollout.RUST_OPT_IN), Rule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})), diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index e0fd3e9a69d..83e719810d5 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import enum import re from collections.abc import Awaitable, Callable, Mapping @@ -12,6 +14,7 @@ from typing_extensions import TypedDict from litellm.types.llms.base import HiddenParams if TYPE_CHECKING: + import httpx2 from mcp.types import EmbeddedResource as MCPEmbeddedResource from mcp.types import ImageContent as MCPImageContent from mcp.types import TextContent as MCPTextContent @@ -348,7 +351,7 @@ def custom_credential_slot(headers: Mapping[str, str] | None) -> str | None: def credential_redirect_hook( configured_url: str, slot: str | None -) -> Callable[[httpx.Request], Awaitable[None]] | None: +) -> Callable[[httpx.Request | httpx2.Request], Awaitable[None]] | None: """An httpx request hook dropping ``slot`` once a redirect leaves ``configured_url``'s origin. None when no guard is needed, so callers do not each repeat the exemption: HTTP clients already @@ -358,7 +361,7 @@ def credential_redirect_hook( if not configured_url or not slot or same_header(slot, DEFAULT_CREDENTIAL_HEADER): return None - async def guard(request: httpx.Request) -> None: + async def guard(request: httpx.Request | httpx2.Request) -> None: if slot in request.headers and crosses_origin(configured_url, str(request.url)): del request.headers[slot] diff --git a/litellm/types/proxy/management_endpoints/common_daily_activity.py b/litellm/types/proxy/management_endpoints/common_daily_activity.py index 5d42b1230a0..2a4f6b2944a 100644 --- a/litellm/types/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/types/proxy/management_endpoints/common_daily_activity.py @@ -47,6 +47,7 @@ class KeyMetadata(BaseModel): team_id: str | None = None user_id: str | None = None user_email: str | None = None + key_exists: bool | None = None class KeyMetricWithMetadata(MetricBase): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 378f8fec9d5..d416e2af33a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4020,6 +4020,7 @@ class LlmProviders(str, Enum): BYTEZ = "bytez" REPLICATE = "replicate" REDUCTO = "reducto" + AWS_TEXTRACT = "aws_textract" RUNWAYML = "runwayml" AWS_POLLY = "aws_polly" TRANSCRIBE = "transcribe" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4dbf0337894..53c0807e86c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5300,7 +5300,7 @@ "supports_web_search": false }, "azure/gpt-4.1-nano": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -5334,7 +5334,7 @@ "supports_vision": true }, "azure/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -10207,7 +10207,7 @@ "supports_web_search": false }, "azure/us/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -14246,7 +14246,6 @@ "source": "https://developers.openai.com/api/docs/pricing" }, "claude-haiku-4-5-20251001": { - "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -14270,7 +14269,6 @@ "prompt_cache_min_tokens": 4096 }, "claude-haiku-4-5": { - "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -14420,7 +14418,6 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-4-5": { - "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, @@ -14455,7 +14452,6 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-sonnet-4-5-20250929": { - "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, @@ -14491,7 +14487,6 @@ "source": "https://docs.anthropic.com/en/docs/about-claude/pricing" }, "claude-sonnet-5": { - "deprecation_date": "2027-06-30", "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -14530,7 +14525,6 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-sonnet-4-6": { - "deprecation_date": "2027-02-17", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -14684,7 +14678,6 @@ "prompt_cache_min_tokens": 1024 }, "claude-opus-4-5-20251101": { - "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -14714,7 +14707,6 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-5": { - "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -14745,7 +14737,6 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-4-6": { - "deprecation_date": "2027-02-05", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -14783,7 +14774,6 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-4-6-20260205": { - "deprecation_date": "2027-02-05", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -14820,7 +14810,6 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-7": { - "deprecation_date": "2027-04-16", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -14859,7 +14848,6 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-4-7-20260416": { - "deprecation_date": "2027-04-16", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -14897,7 +14885,6 @@ "prompt_cache_min_tokens": 2048 }, "claude-fable-5": { - "deprecation_date": "2027-06-09", "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "cache_read_input_token_cost": 1e-06, @@ -14937,7 +14924,6 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-fable-5-1": { - "deprecation_date": "2027-09-01", "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "cache_read_input_token_cost": 2.5e-07, @@ -14978,7 +14964,6 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-5": { - "deprecation_date": "2027-07-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -15020,7 +15005,6 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-opus-4-8": { - "deprecation_date": "2027-05-28", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -40942,7 +40926,7 @@ "supports_prompt_caching": true, "supports_reasoning": false, "supports_response_schema": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-3.5-sonnet": { "input_cost_per_token": 3e-06, @@ -40998,7 +40982,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.1": { "input_cost_per_image": 0.0048, @@ -41024,7 +41008,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-sonnet-4": { "input_cost_per_image": 0.0048, @@ -41054,7 +41038,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-sonnet-4.6": { "supports_adaptive_thinking": true, @@ -41086,7 +41070,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.5": { "cache_creation_input_token_cost": 6.25e-06, @@ -41112,7 +41096,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.6": { "supports_adaptive_thinking": true, @@ -41140,7 +41124,7 @@ "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_pdf_input": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-sonnet-4.5": { "input_cost_per_image": 0.0048, @@ -41170,7 +41154,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-haiku-4.5": { "cache_creation_input_token_cost": 1.25e-06, @@ -41195,7 +41179,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.7": { "supports_adaptive_thinking": true, @@ -41223,7 +41207,7 @@ "prompt_cache_min_tokens": 2048, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-5": { "prompt_cache_min_tokens": 512, @@ -41250,7 +41234,7 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, "openrouter/bytedance/ui-tars-1.5-7b": { @@ -41420,35 +41404,36 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 1.6e-06, + "input_cost_per_token": 4.22298e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 384000, + "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 3.2e-06, + "output_cost_per_token": 8.44596e-07, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 1.35e-07, + "cache_read_input_token_cost": 3.51915e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, "supports_web_search": false }, "openrouter/deepseek/deepseek-v4.1-flash": { - "input_cost_per_token": 1.5e-07, - "output_cost_per_token": 6e-07, - "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 6e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":1.5e-7,"output_cost_per_token":6e-7,"cache_read_input_token_cost":3e-9}, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -41523,7 +41508,7 @@ "supports_audio_input": true, "supports_pdf_input": true, "supports_reasoning": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-2.5-pro": { @@ -41553,7 +41538,7 @@ "supports_audio_input": true, "supports_pdf_input": true, "supports_reasoning": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3-pro-preview": { @@ -41638,7 +41623,7 @@ "supports_tool_choice": true, "supports_url_context": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "tpm": 800000, "supports_video_input": true }, @@ -41684,7 +41669,7 @@ "supports_url_context": true, "supports_video_input": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "tpm": 800000 }, "openrouter/google/gemini-3.1-flash-lite": { @@ -41729,7 +41714,7 @@ "supports_url_context": true, "supports_video_input": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "tpm": 800000 }, "openrouter/google/gemini-3.1-pro-preview": { @@ -41767,7 +41752,7 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/gryphe/mythomax-l2-13b": { @@ -42053,11 +42038,11 @@ }, "openrouter/nvidia/nemotron-3.5-lightning": { "cache_read_input_token_cost": 4e-08, - "input_cost_per_token": 8e-08, + "input_cost_per_token": 7e-08, "litellm_provider": "openrouter", "max_input_tokens": 262144, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 235929, + "max_tokens": 235929, "mode": "chat", "output_cost_per_token": 2e-07, "source": "https://openrouter.ai/api/v1/models", @@ -42148,7 +42133,7 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-4.1-mini": { "cache_read_input_token_cost": 1e-07, @@ -42170,7 +42155,7 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-4.1-nano": { "cache_read_input_token_cost": 2.5e-08, @@ -42192,7 +42177,7 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-4o": { "input_cost_per_token": 2.5e-06, @@ -42298,7 +42283,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5": { "cache_read_input_token_cost": 1.25e-07, @@ -42325,7 +42310,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5-mini": { "cache_read_input_token_cost": 2.5e-08, @@ -42352,7 +42337,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5-nano": { "cache_read_input_token_cost": 5e-09, @@ -42379,7 +42364,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.1-codex-max": { "cache_read_input_token_cost": 1.25e-07, @@ -42406,7 +42391,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.2": { "input_cost_per_image": 0, @@ -42427,7 +42412,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.2-chat": { "input_cost_per_image": 0, @@ -42448,7 +42433,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.2-pro": { "input_cost_per_image": 0, @@ -42468,7 +42453,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-sol": { "cache_creation_input_token_cost": 2.5e-06, @@ -42509,7 +42494,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-sol-pro": { "input_cost_per_token": 2e-06, @@ -42534,15 +42519,15 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-oss-120b": { "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 1.5e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, - "max_output_tokens": 117964, - "max_tokens": 117964, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 6e-07, "source": "https://openrouter.ai/api/v1/models", @@ -42550,7 +42535,7 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": false, - "supports_prompt_caching": false, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -42598,7 +42583,7 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/o3-mini": { "input_cost_per_token": 1.1e-06, @@ -42619,7 +42604,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/o3-mini-high": { "input_cost_per_token": 1.1e-06, @@ -42640,7 +42625,7 @@ "supports_audio_input": false, "supports_pdf_input": true, "supports_response_schema": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/qwen/qwen-2.5-coder-32b-instruct": { "input_cost_per_token": 6.6e-07, @@ -60378,7 +60363,6 @@ } }, "claude-mythos-5-1": { - "deprecation_date": "2027-09-01", "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "cache_read_input_token_cost": 2.5e-07, @@ -65630,7 +65614,7 @@ "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "prompt_cache_min_tokens": 512, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-fable-5.1": { "input_cost_per_token": 1e-05, @@ -65656,7 +65640,7 @@ "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "prompt_cache_min_tokens": 512, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.8": { "input_cost_per_token": 5e-06, @@ -65680,7 +65664,7 @@ "supports_prompt_caching": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-sonnet-5": { "input_cost_per_token": 2e-06, @@ -65704,7 +65688,7 @@ "supports_prompt_caching": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, - "supports_web_search": false + "supports_web_search": true }, "openrouter/google/gemini-2.5-flash-lite": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -65728,7 +65712,7 @@ "deprecation_date": "2026-10-20", "input_cost_per_audio_token": 3e-07, "supports_prompt_caching": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.5-flash": { @@ -65752,7 +65736,7 @@ "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 3e-06, "supports_prompt_caching": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.5-flash-lite": { @@ -65776,7 +65760,7 @@ "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 3e-07, "supports_prompt_caching": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.6-flash": { @@ -65800,7 +65784,7 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.7-flash": { @@ -65824,7 +65808,7 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.8-flash": { @@ -65848,7 +65832,7 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/openai/gpt-4o-mini": { @@ -65889,7 +65873,7 @@ "supports_audio_input": false, "cache_read_input_token_cost": 1.25e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.3-codex": { "input_cost_per_token": 1.75e-06, @@ -65909,7 +65893,7 @@ "supports_audio_input": false, "cache_read_input_token_cost": 1.75e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.4": { "input_cost_per_token": 2.5e-06, @@ -65932,7 +65916,7 @@ "input_cost_per_token_above_272k_tokens": 5e-06, "output_cost_per_token_above_272k_tokens": 2.25e-05, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.4-mini": { "input_cost_per_token": 7.5e-07, @@ -65952,7 +65936,7 @@ "supports_audio_input": false, "cache_read_input_token_cost": 7.5e-08, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.4-nano": { "input_cost_per_token": 2e-07, @@ -65972,7 +65956,7 @@ "supports_audio_input": false, "cache_read_input_token_cost": 2e-08, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.5": { "input_cost_per_token": 5e-06, @@ -65995,7 +65979,7 @@ "input_cost_per_token_above_272k_tokens": 1e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-luna": { "cache_creation_input_token_cost": 2.5e-07, @@ -66020,7 +66004,7 @@ "input_cost_per_token_above_272k_tokens": 4e-07, "output_cost_per_token_above_272k_tokens": 1.8e-06, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-luna-pro": { "input_cost_per_token": 2e-07, @@ -66045,7 +66029,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-terra": { "cache_creation_input_token_cost": 2.5e-06, @@ -66070,7 +66054,7 @@ "input_cost_per_token_above_272k_tokens": 4e-06, "output_cost_per_token_above_272k_tokens": 1.8e-05, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-terra-pro": { "input_cost_per_token": 2e-06, @@ -66095,7 +66079,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/o3": { "input_cost_per_token": 2e-06, @@ -66115,7 +66099,7 @@ "supports_audio_input": false, "cache_read_input_token_cost": 5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/o4-mini": { "input_cost_per_token": 1.1e-06, @@ -66135,7 +66119,7 @@ "supports_audio_input": false, "cache_read_input_token_cost": 2.75e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/x-ai/grok-4.20": { "input_cost_per_token": 1.25e-06, @@ -66158,7 +66142,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/x-ai/grok-4.20-multi-agent": { "input_cost_per_token": 1.25e-06, @@ -66181,7 +66165,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/x-ai/grok-4.3": { "input_cost_per_token": 1.25e-06, @@ -66204,7 +66188,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/x-ai/grok-4.5": { "input_cost_per_token": 2e-06, @@ -66227,7 +66211,7 @@ "input_cost_per_token_above_200k_tokens": 4e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/x-ai/grok-4.6": { "input_cost_per_token": 2e-06, @@ -66250,7 +66234,7 @@ "input_cost_per_token_above_200k_tokens": 4e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/x-ai/grok-build-0.1": { "input_cost_per_token": 1e-06, @@ -66273,7 +66257,7 @@ "input_cost_per_token_above_200k_tokens": 2e-06, "output_cost_per_token_above_200k_tokens": 4e-06, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "baseten/zai-org/GLM-5.3": { "cache_read_input_token_cost": 1.4e-07, @@ -66362,7 +66346,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-6-astra-pro": { "input_cost_per_token": 1e-05, @@ -66387,7 +66371,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/qwen/qwen3.8-flash": { "input_cost_per_token": 1.5e-07, @@ -66436,8 +66420,8 @@ "cache_read_input_token_cost": 6.86e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -66456,8 +66440,8 @@ "cache_read_input_token_cost": 1.69e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 943717, - "max_tokens": 943717, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -66570,9 +66554,9 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-flash-0731": { - "input_cost_per_token": 6e-08, - "output_cost_per_token": 1.2e-07, - "cache_read_input_token_cost": 1.2e-08, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 8e-08, + "cache_read_input_token_cost": 1.6e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 943718, @@ -66655,9 +66639,9 @@ "supports_web_search": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 2.1e-06, - "output_cost_per_token": 1.095e-05, - "cache_read_input_token_cost": 2.3e-07, + "input_cost_per_token": 1.7e-06, + "output_cost_per_token": 8.5e-06, + "cache_read_input_token_cost": 1.7e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -66731,7 +66715,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/google/gemini-3.1-flash-image": { "input_cost_per_token": 5e-07, @@ -66751,7 +66735,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/google/gemini-3-pro-image": { "input_cost_per_token": 2e-06, @@ -66775,7 +66759,7 @@ "supports_response_schema": true, "supports_vision": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/z-ai/glm-5.2": { "input_cost_per_token": 5.544e-07, @@ -66877,13 +66861,13 @@ "supports_web_search": false }, "openrouter/nvidia/nemotron-3-ultra-550b-a55b": { - "input_cost_per_token": 6.25e-07, - "output_cost_per_token": 3.125e-06, - "cache_read_input_token_cost": 1.875e-07, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 2.4e-06, + "cache_read_input_token_cost": 1.2e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, - "max_output_tokens": 32768, - "max_tokens": 32768, + "max_output_tokens": 182520, + "max_tokens": 182520, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -67117,7 +67101,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-chat-latest": { "input_cost_per_token": 5e-06, @@ -67137,12 +67121,12 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 4.984e-08, - "output_cost_per_token": 9.968e-08, - "cache_read_input_token_cost": 9.968e-09, + "input_cost_per_token": 4.06e-08, + "output_cost_per_token": 8.12e-08, + "cache_read_input_token_cost": 8.12e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -67431,7 +67415,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/google/gemini-3.1-flash-image-preview": { "input_cost_per_token": 5e-07, @@ -67451,7 +67435,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/google/gemini-3.1-pro-preview-customtools": { "input_cost_per_token": 2e-06, @@ -67477,7 +67461,7 @@ "supports_pdf_input": true, "supports_audio_input": true, "supports_prompt_caching": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/qwen/qwen3-max-thinking": { @@ -67645,7 +67629,7 @@ "supports_response_schema": true, "supports_vision": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.1-codex": { "input_cost_per_token": 1.25e-06, @@ -67665,7 +67649,7 @@ "supports_response_schema": true, "supports_vision": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.1-codex-mini": { "input_cost_per_token": 2.5e-07, @@ -67685,7 +67669,7 @@ "supports_response_schema": true, "supports_vision": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/moonshotai/kimi-k2-thinking": { "input_cost_per_token": 6e-07, @@ -67828,7 +67812,7 @@ "supports_vision": true, "supports_prompt_caching": true, "supports_reasoning": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/qwen/qwen3-vl-30b-a3b-thinking": { "input_cost_per_token": 2e-07, @@ -67885,7 +67869,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/qwen/qwen3-vl-235b-a22b-thinking": { "input_cost_per_token": 4e-07, @@ -68051,7 +68035,7 @@ "supports_audio_input": false, "supports_function_calling": true, "supports_pdf_input": false, - "supports_prompt_caching": false, + "supports_prompt_caching": true, "supports_reasoning": false, "supports_tool_choice": true, "supports_response_schema": true, @@ -68290,7 +68274,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/google/gemini-2.5-pro-preview": { "input_cost_per_token": 1.25e-06, @@ -68316,7 +68300,7 @@ "supports_pdf_input": true, "supports_audio_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/mistralai/mistral-medium-3": { "input_cost_per_token": 4e-07, @@ -68494,7 +68478,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/meta-llama/llama-4-maverick": { "input_cost_per_token": 1.875e-07, @@ -68551,7 +68535,7 @@ "supports_vision": true, "supports_pdf_input": true, "supports_prompt_caching": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/google/gemma-3-4b-it": { "input_cost_per_token": 5e-08, @@ -69338,7 +69322,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-4.1-nano": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -69692,7 +69676,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-4.1-nano": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -71106,7 +71090,7 @@ "supports_response_schema": true, "supports_tool_choice": false, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~anthropic/claude-haiku-latest": { "cache_creation_input_token_cost": 1.25e-06, @@ -71128,7 +71112,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~anthropic/claude-opus-latest": { "cache_creation_input_token_cost": 6.25e-06, @@ -71150,7 +71134,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~anthropic/claude-sonnet-latest": { "cache_creation_input_token_cost": 2.5e-06, @@ -71172,17 +71156,17 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~deepseek/deepseek-flash-latest": { - "cache_read_input_token_cost": 4.2e-09, - "input_cost_per_token": 1.4e-07, + "cache_read_input_token_cost": 2.6e-09, + "input_cost_per_token": 1.3e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 4.2e-07, + "output_cost_per_token": 5.2e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71215,14 +71199,14 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-v4-flash-latest": { - "cache_read_input_token_cost": 1.75e-09, - "input_cost_per_token": 5.5e-08, + "cache_read_input_token_cost": 1.6e-08, + "input_cost_per_token": 4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 393216, - "max_tokens": 393216, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 1.65e-07, + "output_cost_per_token": 8e-08, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71255,7 +71239,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~google/gemini-pro-latest": { "cache_creation_input_token_cost": 3.75e-07, @@ -71281,17 +71265,17 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~moonshotai/kimi-latest": { - "cache_read_input_token_cost": 2.3e-07, - "input_cost_per_token": 2.1e-06, + "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 1.7e-06, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 1.095e-05, + "output_cost_per_token": 8.5e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71326,7 +71310,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~openai/gpt-luna-latest": { "cache_creation_input_token_cost": 2.5e-07, @@ -71351,7 +71335,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~openai/gpt-mini-latest": { "cache_read_input_token_cost": 7.5e-08, @@ -71371,7 +71355,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~openai/gpt-sol-latest": { "cache_creation_input_token_cost": 2.5e-06, @@ -71396,7 +71380,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~openai/gpt-terra-latest": { "cache_creation_input_token_cost": 2.5e-06, @@ -71421,7 +71405,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~x-ai/grok-latest": { "cache_read_input_token_cost": 5e-07, @@ -71444,7 +71428,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/~z-ai/glm-flash-latest": { "cache_read_input_token_cost": 1.5e-08, @@ -71467,14 +71451,14 @@ "supports_web_search": false }, "openrouter/~z-ai/glm-latest": { - "cache_read_input_token_cost": 1.46625e-07, - "input_cost_per_token": 9e-07, + "cache_read_input_token_cost": 1.5678e-07, + "input_cost_per_token": 8.442e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 235929, - "max_tokens": 235929, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.805e-06, + "output_cost_per_token": 2.6532e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -71700,7 +71684,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-fable-5.1:batch": { "cache_creation_input_token_cost": 6.25e-06, @@ -71722,7 +71706,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-haiku-4.5:batch": { "cache_creation_input_token_cost": 6.25e-07, @@ -71744,7 +71728,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.1:batch": { "cache_creation_input_token_cost": 9.375e-06, @@ -71766,7 +71750,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.5:batch": { "cache_creation_input_token_cost": 3.125e-06, @@ -71788,7 +71772,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.6:batch": { "cache_creation_input_token_cost": 3.125e-06, @@ -71810,7 +71794,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.7:batch": { "cache_creation_input_token_cost": 3.125e-06, @@ -71832,7 +71816,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-4.8:batch": { "cache_creation_input_token_cost": 3.125e-06, @@ -71854,7 +71838,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-opus-5:batch": { "cache_creation_input_token_cost": 3.125e-06, @@ -71876,7 +71860,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-sonnet-4.5:batch": { "cache_creation_input_token_cost": 1.875e-06, @@ -71902,7 +71886,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-sonnet-4.6:batch": { "cache_creation_input_token_cost": 1.875e-06, @@ -71924,7 +71908,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/anthropic/claude-sonnet-5:batch": { "cache_creation_input_token_cost": 1.25e-06, @@ -71946,7 +71930,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/arcee-ai/trinity-large-thinking": { "cache_read_input_token_cost": 6e-08, @@ -72344,7 +72328,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-2.5-flash:batch": { @@ -72368,7 +72352,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-2.5-pro:batch": { @@ -72395,7 +72379,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3-flash-preview:batch": { @@ -72416,7 +72400,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.1-flash-lite:batch": { @@ -72439,7 +72423,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.1-pro-preview:batch": { @@ -72462,7 +72446,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.5-flash-lite:batch": { @@ -72485,7 +72469,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.5-flash:batch": { @@ -72508,7 +72492,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.6-flash:batch": { @@ -72532,7 +72516,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.7-flash:batch": { @@ -72556,7 +72540,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/google/gemini-3.8-flash:batch": { @@ -72580,7 +72564,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false, + "supports_web_search": true, "supports_video_input": true }, "openrouter/ibm-granite/granite-4.0-h-micro": { @@ -72956,7 +72940,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/meta/muse-spark-1.2": { "cache_read_input_token_cost": 1.5e-07, @@ -72976,7 +72960,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/meta/muse-spark-1.2-contributor": { "cache_read_input_token_cost": 2e-09, @@ -72996,7 +72980,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/meta/muse-spark-1.3": { "cache_read_input_token_cost": 1.5e-07, @@ -73016,7 +73000,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/meta/muse-spark-1.3-contributor": { "cache_read_input_token_cost": 2e-09, @@ -73036,7 +73020,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/microsoft/phi-4": { "input_cost_per_token": 7e-08, @@ -73404,7 +73388,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-4-turbo:batch": { "input_cost_per_token": 5e-06, @@ -73423,7 +73407,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-4.1-mini:batch": { "cache_read_input_token_cost": 5e-08, @@ -73443,7 +73427,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-4.1-nano:batch": { "cache_read_input_token_cost": 1.25e-08, @@ -73463,7 +73447,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-4.1:batch": { "cache_read_input_token_cost": 2.5e-07, @@ -73483,7 +73467,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-4o-mini:batch": { "cache_read_input_token_cost": 3.75e-08, @@ -73543,7 +73527,7 @@ "supports_response_schema": true, "supports_tool_choice": false, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5-image-mini": { "cache_read_input_token_cost": 2.5e-07, @@ -73563,7 +73547,7 @@ "supports_response_schema": true, "supports_tool_choice": false, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5-mini:batch": { "cache_read_input_token_cost": 1.25e-08, @@ -73583,7 +73567,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5-nano:batch": { "cache_read_input_token_cost": 2.5e-09, @@ -73603,7 +73587,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5-pro:batch": { "input_cost_per_token": 7.5e-06, @@ -73622,7 +73606,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5:batch": { "cache_read_input_token_cost": 6.25e-08, @@ -73642,7 +73626,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.1:batch": { "cache_read_input_token_cost": 6.25e-08, @@ -73662,7 +73646,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.2-pro:batch": { "input_cost_per_token": 1.05e-05, @@ -73681,7 +73665,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.2:batch": { "cache_read_input_token_cost": 8.75e-08, @@ -73701,7 +73685,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.4-image-2": { "cache_read_input_token_cost": 2e-06, @@ -73721,7 +73705,7 @@ "supports_response_schema": true, "supports_tool_choice": false, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.4-mini:batch": { "cache_read_input_token_cost": 3.75e-08, @@ -73741,7 +73725,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.4-nano:batch": { "cache_read_input_token_cost": 1e-08, @@ -73761,7 +73745,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.4-pro:batch": { "input_cost_per_token": 1.5e-05, @@ -73782,7 +73766,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.4:batch": { "cache_read_input_token_cost": 1.25e-07, @@ -73805,7 +73789,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.5-pro:batch": { "input_cost_per_token": 1.5e-05, @@ -73826,7 +73810,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.5:batch": { "cache_read_input_token_cost": 2.5e-07, @@ -73849,7 +73833,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-luna-pro:batch": { "cache_read_input_token_cost": 1e-08, @@ -73872,7 +73856,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-luna:batch": { "cache_read_input_token_cost": 1e-08, @@ -73895,7 +73879,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-sol-pro:batch": { "cache_creation_input_token_cost": 1.25e-06, @@ -73920,7 +73904,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-sol:batch": { "cache_creation_input_token_cost": 1.25e-06, @@ -73945,7 +73929,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-terra-pro:batch": { "cache_read_input_token_cost": 1e-07, @@ -73968,7 +73952,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-5.6-terra:batch": { "cache_read_input_token_cost": 1e-07, @@ -73991,7 +73975,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-6-astra-pro:batch": { "cache_creation_input_token_cost": 6.25e-06, @@ -74016,7 +74000,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-6-astra:batch": { "cache_creation_input_token_cost": 6.25e-06, @@ -74041,7 +74025,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/gpt-oss-120b:batch": { "input_cost_per_token": 1.5e-07, @@ -74080,7 +74064,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": false, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/o3:batch": { "cache_read_input_token_cost": 2.5e-07, @@ -74100,7 +74084,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/openai/o4-mini:batch": { "cache_read_input_token_cost": 1.375e-07, @@ -74120,7 +74104,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/perceptron/perceptron-mk1": { "input_cost_per_token": 1.5e-07, @@ -74629,14 +74613,15 @@ "supports_web_search": false }, "openrouter/tencent/hy3": { - "cache_read_input_token_cost": 2.0625e-08, - "input_cost_per_token": 8.25e-08, + "cache_read_input_token_cost": 3.3e-08, + "input_cost_per_token": 1.32e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3.3e-07, + "off_peak_pricing": {"hours_utc":"16:00-00:00","input_cost_per_token":8.25e-8,"output_cost_per_token":3.3e-7,"cache_read_input_token_cost":2.0625e-8}, + "output_cost_per_token": 5.28e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -74860,7 +74845,7 @@ "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": false, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false @@ -74945,7 +74930,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": true }, "openrouter/z-ai/glm-5.2:batch": { "cache_read_input_token_cost": 7e-08, @@ -75006,5 +74991,44 @@ "supports_tool_choice": true, "supports_vision": false, "supports_web_search": false + }, + "openrouter/prism-ml/ternary-bonsai-2-27b": { + "input_cost_per_token": 7.5e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_pdf_input": false, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": false + }, + "openrouter/z-ai/glm-5.3-flashx": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 3.7e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.25e-06, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_pdf_input": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": false } } diff --git a/pyproject.toml b/pyproject.toml index a017e39e084..f2ee1d92d7f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,13 +18,15 @@ dependencies = [ "httpx[http2]>=0.28.0,<1.0", "openai>=2.20.0,<3.0.0", "python-dotenv>=1.0.0,<2.0", - "tiktoken>=0.8.0,<1.0", + "tiktoken>=0.8.0,<1.0; python_version < '3.14'", + "tiktoken>=0.12.0,<1.0; python_version >= '3.14'", "importlib-metadata>=8.0.0,<9.0", "tokenizers>=0.21.0,<1.0", "click>=8.0.0,<9.0", "jinja2>=3.1.6,<4.0", "aiohttp>=3.14.2,<4.0", - "pydantic>=2.10.0,<3.0.0", + "pydantic>=2.11.0,<3.0.0; python_version < '3.14'", + "pydantic>=2.12.0,<3.0.0; python_version >= '3.14'", "pydantic-settings>=2.14.1,<3.0", "jsonschema>=4.0.0,<5.0", "boto3>=1.43.1,<2.0", @@ -66,7 +68,9 @@ proxy = [ "boto3>=1.43.1,<2.0", "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", - "mcp>=1.28.1,<2.0", + "mcp>=2.2.0,<3", + "httpx2>=2.5.0,<3", + "pydantic>=2.12.0,<3", "litellm-proxy-extras==0.4.99", "litellm-enterprise==0.1.68", "RestrictedPython>=8.5,<9.0", @@ -113,7 +117,7 @@ utils = [ "numpydoc>=1.8.0,<2.0", ] caching = ["diskcache>=5.6.3,<6.0"] -mcp = ["mcp>=1.28.1,<2.0"] +mcp = ["mcp>=2.2.0,<3", "httpx2>=2.5.0,<3", "pydantic>=2.12.0,<3"] # Driver for the MongoDB Atlas vector store; Atlas Vector Search has no HTTP query API. # The floor is 4.9 because that is the release AsyncMongoClient landed in. # SAML SSO for the admin UI. python3-saml pulls in xmlsec/lxml, whose wheels @@ -232,7 +236,7 @@ e2e-dev = [ "websockets>=15.0.1,<16.0", "locust==2.45.0", "psutil==7.2.2", - "mcp>=1.28.1,<2.0", + "mcp>=2.2.0,<3", ] proxy-dev = [ "prisma==0.11.0", @@ -272,7 +276,6 @@ ci = [ "blockbuster==1.5.26", "beautifulsoup4==4.14.3", "pylint==4.0.5", - "langchain-mcp-adapters==0.2.1", "langchain-openai==1.1.14", "langgraph>=1.2.4,<1.3.0", "langgraph-prebuilt>=1.1.0,<1.3.0", diff --git a/scripts/check_mcp_sdk_install.py b/scripts/check_mcp_sdk_install.py new file mode 100644 index 00000000000..f5ab51b2f55 --- /dev/null +++ b/scripts/check_mcp_sdk_install.py @@ -0,0 +1,77 @@ +import argparse +import importlib +import importlib.metadata +import sys +from typing import Final + +MINIMUM_MCP_VERSION: Final[tuple[int, int, int]] = (2, 2, 0) + +IMPORTED_MODULES: Final[tuple[str, ...]] = ( + "litellm", + "litellm.experimental_mcp_client", + "litellm.experimental_mcp_client.client", + "litellm.proxy._experimental.mcp_server.server", + "litellm.proxy._experimental.mcp_server.mcp_server_manager", + "litellm.proxy._experimental.mcp_server.rest_endpoints", +) + + +def _version_tuple(distribution: str) -> tuple[int, ...]: + return tuple(int(part) for part in importlib.metadata.version(distribution).split(".") if part.isdigit()) + + +def main() -> int: + parser: Final = argparse.ArgumentParser() + parser.add_argument("--extra", choices=("mcp", "proxy"), default="proxy") + extra: Final = parser.parse_args().extra + for module_name in IMPORTED_MODULES if extra == "proxy" else IMPORTED_MODULES[:3]: + try: + importlib.import_module(module_name) + except Exception as exc: + sys.stderr.write(f"failed to import {module_name}: {exc}\n") + return 1 + + mcp_version: Final = _version_tuple("mcp") + if mcp_version < MINIMUM_MCP_VERSION: + sys.stderr.write(f"mcp {importlib.metadata.version('mcp')} below floor 2.2.0\n") + return 1 + + from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS + + for required in ("2024-11-05", "2025-06-18"): + if required not in HANDSHAKE_PROTOCOL_VERSIONS: + sys.stderr.write(f"HANDSHAKE_PROTOCOL_VERSIONS missing {required}\n") + return 1 + + if extra == "proxy": + scope: Final = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [(b"mcp-protocol-version", b"2026-07-28")], + } + mcp_server: Final = sys.modules["litellm.proxy._experimental.mcp_server.server"] + if mcp_server.unsupported_protocol_version(scope) != "2026-07-28": + sys.stderr.write("unsupported_protocol_version accepted a modern-only version\n") + return 1 + if ( + mcp_server.unsupported_protocol_version(dict(scope, headers=[(b"mcp-protocol-version", b"2025-06-18")])) + is not None + ): + sys.stderr.write("unsupported_protocol_version rejected a handshake version\n") + return 1 + + sys.stdout.write( + "python {} mcp {} httpx2 {} pydantic {} litellm {}\n".format( + sys.version.split()[0], + importlib.metadata.version("mcp"), + importlib.metadata.version("httpx2"), + importlib.metadata.version("pydantic"), + importlib.metadata.version("litellm"), + ) + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/base_sdk_tests/check_base_sdk_install.py b/tests/base_sdk_tests/check_base_sdk_install.py index 6b38de75e2e..190a900faf9 100644 --- a/tests/base_sdk_tests/check_base_sdk_install.py +++ b/tests/base_sdk_tests/check_base_sdk_install.py @@ -11,7 +11,7 @@ import sys import traceback from collections.abc import Callable -EXTRAS_ONLY_MODULES = ("fastapi", "uvicorn", "keyring") +EXTRAS_ONLY_MODULES = ("fastapi", "uvicorn", "keyring", "mcp", "mcp_types", "httpx2", "httpcore2") def _require(condition: bool, message: str) -> None: diff --git a/tests/code_coverage_tests/liccheck.ini b/tests/code_coverage_tests/liccheck.ini index 9103d913c36..8a3e880043b 100644 --- a/tests/code_coverage_tests/liccheck.ini +++ b/tests/code_coverage_tests/liccheck.ini @@ -169,7 +169,9 @@ pygithub: >=2.8.1 # LGPL license argon2-cffi: >=25.1.0 # MIT License blockbuster: >=1.5.26 # Apache 2.0 license pylint: >=3.3.9 # GPLv2 license -langchain-mcp-adapters: >=0.2.1 # MIT License +httpx2: >=2.5.0 # BSD 3-Clause License +httpcore2: >=2.5.0 # BSD 3-Clause License +mcp-types: >=2.2.0 # MIT License langgraph: >=1.0.10 # MIT License langgraph-prebuilt: >=1.0.8 # MIT License - https://github.com/langchain-ai/langgraph/blob/main/LICENSE hypothesis: >=6.165.10 # MPL 2.0 license diff --git a/tests/e2e/mcp/oauth_chat_client.py b/tests/e2e/mcp/oauth_chat_client.py index f6e72fb37dc..d2fca790132 100644 --- a/tests/e2e/mcp/oauth_chat_client.py +++ b/tests/e2e/mcp/oauth_chat_client.py @@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, Final from urllib.parse import parse_qsl import httpx +import httpx2 import pytest from e2e_config import PROXY_BASE_URL, REQUEST_TIMEOUT from e2e_http import AuthHeaders, NoBody, unwrap @@ -29,7 +30,7 @@ from idp import Identity from mcp import ClientSession from mcp.client.auth import OAuthClientProvider from mcp.client.streamable_http import streamable_http_client -from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthToken +from mcp.shared.auth import AuthorizationCodeResult, OAuthClientInformationFull, OAuthClientMetadata, OAuthToken from mcp.types import TextContent from models import ( ChatBody, @@ -192,10 +193,10 @@ def _oauth_provider( redirect_handler: Final = _reject_redirect if storage_state_path is None else _follow_redirect - async def callback_handler() -> tuple[str, str | None]: + async def callback_handler() -> AuthorizationCodeResult: code = code_holder.get("code") assert code is not None, "callback_handler ran before the authorize redirect completed" - return code, code_holder.get("state") + return AuthorizationCodeResult(code=code, state=code_holder.get("state")) return OAuthClientProvider( server_url=url, @@ -214,24 +215,24 @@ def _oauth_provider( ) -class _HeaderInjectingTransport(httpx.AsyncBaseTransport): +class _HeaderInjectingTransport(httpx2.AsyncBaseTransport): """Adds the caller's LiteLLM key header to every outgoing SDK request (discovery, DCR, token exchange), so the gateway resolves which user to store the upstream token for from the key on the token exchange, exactly like a production MCP host configured with a LiteLLM key header.""" - def __init__(self, inner: httpx.AsyncBaseTransport, headers: dict[str, str], gateway_url: str) -> None: + def __init__(self, inner: httpx2.AsyncBaseTransport, headers: dict[str, str], gateway_url: str) -> None: self._inner = inner self._headers = headers - self._gateway_url = httpx.URL(gateway_url) + self._gateway_url = httpx2.URL(gateway_url) @staticmethod - def _port(url: httpx.URL) -> int | None: + def _port(url: httpx2.URL) -> int | None: if url.port is not None: return url.port return {"http": 80, "https": 443}.get(url.scheme) - async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response: same_origin: Final = ( request.url.scheme == self._gateway_url.scheme and request.url.host == self._gateway_url.host @@ -253,12 +254,12 @@ class _HeaderInjectingTransport(httpx.AsyncBaseTransport): def _oauth_http_client( headers: dict[str, str], auth: OAuthClientProvider, gateway_url: str = PROXY_BASE_URL -) -> httpx.AsyncClient: - return httpx.AsyncClient( +) -> httpx2.AsyncClient: + return httpx2.AsyncClient( auth=auth, - timeout=httpx.Timeout(REQUEST_TIMEOUT), + timeout=httpx2.Timeout(REQUEST_TIMEOUT), follow_redirects=True, - transport=_HeaderInjectingTransport(httpx.AsyncHTTPTransport(), headers, gateway_url), + transport=_HeaderInjectingTransport(httpx2.AsyncHTTPTransport(), headers, gateway_url), ) @@ -266,7 +267,7 @@ async def _seed_via_dance( url: str, headers: dict[str, str], storage: InMemoryTokenStorage, storage_state_path: str ) -> tuple[str, ...]: async with _oauth_http_client(headers, _oauth_provider(url, storage, storage_state_path)) as http_client: - async with streamable_http_client(url, http_client=http_client) as (read, write, _): + async with streamable_http_client(url, http_client=http_client) as (read, write): async with ClientSession(read, write) as session: await session.initialize() listed = await session.list_tools() @@ -297,7 +298,7 @@ async def _list_and_call( _oauth_provider(url, storage, storage_state_path, identity, server_alias, allow_upstream_consent), gateway_url, ) as http_client: - async with streamable_http_client(url, http_client=http_client) as (read, write, _): + async with streamable_http_client(url, http_client=http_client) as (read, write): async with ClientSession(read, write) as session: await session.initialize() listed: Final = await session.list_tools() @@ -305,7 +306,7 @@ async def _list_and_call( text: Final = "".join(content.text for content in result.content if isinstance(content, TextContent)) return OauthToolRun( tools=tuple(sorted(tool_item.name for tool_item in listed.tools)), - is_error=result.isError, + is_error=result.is_error, text=text, ) diff --git a/tests/mcp_tests/conftest.py b/tests/mcp_tests/conftest.py index eff32f27aec..ca3e25949ba 100644 --- a/tests/mcp_tests/conftest.py +++ b/tests/mcp_tests/conftest.py @@ -74,3 +74,14 @@ def pytest_collection_modifyitems(config, items): # Reorder the items list items[:] = custom_logger_tests + other_tests + + +@pytest.fixture +def config_only_mcp_manager_factory(): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + + class ConfigOnlyManager(MCPServerManager): + def initialize_tool_name_to_mcp_server_name_mapping(self): + return None + + return ConfigOnlyManager diff --git a/tests/mcp_tests/mcp_server.py b/tests/mcp_tests/mcp_server.py index eba7cae1bca..f38b6a02139 100644 --- a/tests/mcp_tests/mcp_server.py +++ b/tests/mcp_tests/mcp_server.py @@ -51,6 +51,21 @@ def request_headers(ctx: Context) -> dict[str, str]: } +@mcp.prompt() +def greeting(name: str) -> str: + return f"Hello, {name}" + + +@mcp.resource("memo://status") +def status() -> str: + return "ready" + + +@mcp.resource("memo://greeting/{name}") +def greeting_resource(name: str) -> str: + return f"Hello, {name}" + + def main() -> None: args = _parse_args() transport = (args.transport or "stdio").lower() diff --git a/tests/mcp_tests/test_aresponses_api_with_mcp.py b/tests/mcp_tests/test_aresponses_api_with_mcp.py index 7a48c366003..eb6f78b57a1 100644 --- a/tests/mcp_tests/test_aresponses_api_with_mcp.py +++ b/tests/mcp_tests/test_aresponses_api_with_mcp.py @@ -1,6 +1,7 @@ import logging import os import pytest +from mcp.types import Tool as MCPTool from typing import List, Any, cast from unittest.mock import AsyncMock, patch @@ -371,48 +372,32 @@ async def test_mcp_allowed_tools_filtering(): # Mock MCP tools returned from the server (simulating all available tools) mock_mcp_tools_from_server = [ # Mock MCP tool object with name attribute - type( - "MCPTool", - (), - { + MCPTool.model_validate({ "name": "search_tiktoken_documentation", "description": "Search tiktoken documentation", "inputSchema": { "type": "object", "properties": {"query": {"type": "string"}}, }, - }, - )(), - type( - "MCPTool", - (), - { + }, by_name=False), + MCPTool.model_validate({ "name": "fetch_tiktoken_documentation", "description": "Fetch tiktoken documentation", "inputSchema": { "type": "object", "properties": {"path": {"type": "string"}}, }, - }, - )(), - type( - "MCPTool", - (), - { + }, by_name=False), + MCPTool.model_validate({ "name": "list_tiktoken_functions", "description": "List tiktoken functions", "inputSchema": {"type": "object", "properties": {}}, - }, - )(), - type( - "MCPTool", - (), - { + }, by_name=False), + MCPTool.model_validate({ "name": "get_tiktoken_examples", "description": "Get tiktoken examples", "inputSchema": {"type": "object", "properties": {}}, - }, - )(), + }, by_name=False), ] allowed_mcp_servers = ["gitmcp"] @@ -491,10 +476,7 @@ async def test_mcp_allowed_tools_filtering(): # Test Case 3: Test deduplication of duplicate tools mock_mcp_tools_with_duplicates = [ # First instance of duplicate tool - type( - "MCPTool", - (), - { + MCPTool.model_validate({ "name": "GitMCP-fetch_litellm_documentation", "description": "Fetch entire documentation file from GitHub repository: BerriAI/litellm. Useful for general questions. Always call this tool first if asked about BerriAI/litellm.", "inputSchema": { @@ -502,13 +484,9 @@ async def test_mcp_allowed_tools_filtering(): "properties": {}, "additionalProperties": False, }, - }, - )(), + }, by_name=False), # Second instance of duplicate tool (should be filtered out) - type( - "MCPTool", - (), - { + MCPTool.model_validate({ "name": "GitMCP-fetch_litellm_documentation", "description": "Fetch entire documentation file from GitHub repository: BerriAI/litellm. Useful for general questions. Always call this tool first if asked about BerriAI/litellm.", "inputSchema": { @@ -516,13 +494,9 @@ async def test_mcp_allowed_tools_filtering(): "properties": {}, "additionalProperties": False, }, - }, - )(), + }, by_name=False), # Other unique tools - type( - "MCPTool", - (), - { + MCPTool.model_validate({ "name": "GitMCP-search_litellm_documentation", "description": "Semantically search within the fetched documentation from GitHub repository: BerriAI/litellm. Useful for specific queries.", "inputSchema": { @@ -531,8 +505,7 @@ async def test_mcp_allowed_tools_filtering(): "required": ["query"], "additionalProperties": False, }, - }, - )(), + }, by_name=False), ] mcp_tool_config_with_duplicates = [ @@ -680,10 +653,7 @@ async def test_streaming_mcp_events_validation(): # Mock MCP tools that would be returned from the manager mock_mcp_tools = [ - type( - "MCPTool", - (), - { + MCPTool.model_validate({ "name": "search_repo", "description": "Search BerriAI/litellm repository for information", "inputSchema": { @@ -693,12 +663,8 @@ async def test_streaming_mcp_events_validation(): }, "required": ["query"], }, - }, - )(), - type( - "MCPTool", - (), - { + }, by_name=False), + MCPTool.model_validate({ "name": "get_repo_info", "description": "Get repository information", "inputSchema": { @@ -711,8 +677,7 @@ async def test_streaming_mcp_events_validation(): }, "required": ["repo_name"], }, - }, - )(), + }, by_name=False), ] # Build fake streaming chunks that the inner aresponses() call would yield @@ -920,10 +885,7 @@ async def test_streaming_responses_api_with_mcp_tools( # Mock MCP tools that would be returned from the manager mock_mcp_tools = [ - type( - "MCPTool", - (), - { + MCPTool.model_validate({ "name": "search_repo", "description": "Search BerriAI/litellm repository for information", "inputSchema": { @@ -933,8 +895,7 @@ async def test_streaming_responses_api_with_mcp_tools( }, "required": ["query"], }, - }, - )() + }, by_name=False) ] # Only mock the MCP-specific operations, let LLM responses be real @@ -1263,10 +1224,7 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e(): # Mock MCP tools that would be returned from the manager mock_mcp_tools = [ - type( - "MCPTool", - (), - { + MCPTool.model_validate({ "name": "search_docs", "description": "Search documentation for information", "inputSchema": { @@ -1276,12 +1234,8 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e(): }, "required": ["query"], }, - }, - )(), - type( - "MCPTool", - (), - { + }, by_name=False), + MCPTool.model_validate({ "name": "get_file_content", "description": "Get content of a specific file", "inputSchema": { @@ -1291,8 +1245,7 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e(): }, "required": ["file_path"], }, - }, - )(), + }, by_name=False), ] # Track all calls to the underlying LLM to detect duplicates @@ -1499,10 +1452,7 @@ async def test_streaming_mcp_event_order_and_response_id_consistency( from unittest.mock import AsyncMock, patch mock_mcp_tools = [ - type( - "MCPTool", - (), - { + MCPTool.model_validate({ "name": "get_weather", "description": "Get weather for a city", "inputSchema": { @@ -1512,8 +1462,7 @@ async def test_streaming_mcp_event_order_and_response_id_consistency( }, "required": ["city"], }, - }, - )() + }, by_name=False) ] with caplog.at_level(logging.ERROR): diff --git a/tests/mcp_tests/test_mcp_auth_priority.py b/tests/mcp_tests/test_mcp_auth_priority.py index 7ae0f59afe5..21a89d7ffcc 100644 --- a/tests/mcp_tests/test_mcp_auth_priority.py +++ b/tests/mcp_tests/test_mcp_auth_priority.py @@ -44,14 +44,14 @@ async def test_mcp_server_works_without_config_auth_value(): @pytest.mark.parametrize("token_key", ["authentication_token", "auth_value"]) -async def test_mcp_server_config_auth_value_header_used(token_key): +async def test_mcp_server_config_auth_value_header_used(token_key, config_only_mcp_manager_factory): """Ensure the configured auth token is emitted as the upstream Authorization header. The token is resolved through the v2 credential resolver and rides on the client's httpx.Auth, so assert the header it writes onto the request rather than the (now credential-free) _get_auth_headers() dict. """ - import httpx + import httpx2 from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( StaticHeaderAuth, @@ -66,13 +66,13 @@ async def test_mcp_server_config_auth_value_header_used(token_key): } } - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(config) server = next(iter(manager.config_mcp_servers.values())) client = await manager._create_mcp_client(server) assert isinstance(client._resolved_auth, StaticHeaderAuth) - emitted = next(client._resolved_auth.auth_flow(httpx.Request("POST", server.url))) + emitted = next(client._resolved_auth.auth_flow(httpx2.Request("POST", server.url))) assert emitted.headers["Authorization"] == "Bearer example_token" assert client.auth_type == MCPAuth.bearer_token diff --git a/tests/mcp_tests/test_mcp_chat_completions.py b/tests/mcp_tests/test_mcp_chat_completions.py index fbdbf9152aa..79619eefd7f 100644 --- a/tests/mcp_tests/test_mcp_chat_completions.py +++ b/tests/mcp_tests/test_mcp_chat_completions.py @@ -16,7 +16,7 @@ async def test_acompletion_mcp_auto_exec(monkeypatch): dummy_tool = SimpleNamespace( name="local_search", description="search", - inputSchema={"type": "object", "properties": {}}, + input_schema={"type": "object", "properties": {}}, ) async def fake_process(user_api_key_auth, mcp_tools_with_litellm_proxy, **kwargs): @@ -92,7 +92,7 @@ async def test_acompletion_mcp_respects_manual_approval(monkeypatch): dummy_tool = SimpleNamespace( name="local_search", description="search", - inputSchema={"type": "object", "properties": {}}, + input_schema={"type": "object", "properties": {}}, ) async def fake_process(user_api_key_auth, mcp_tools_with_litellm_proxy, **kwargs): @@ -167,7 +167,7 @@ async def test_completion_mcp_with_streaming_no_timeout_error(monkeypatch): dummy_tool = SimpleNamespace( name="local_search", description="search", - inputSchema={"type": "object", "properties": {}}, + input_schema={"type": "object", "properties": {}}, ) async def fake_process(user_api_key_auth, mcp_tools_with_litellm_proxy, **kwargs): @@ -488,7 +488,7 @@ async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch): dummy_tool = SimpleNamespace( name="local_search", description="search", - inputSchema={"type": "object", "properties": {}}, + input_schema={"type": "object", "properties": {}}, ) async def fake_process(user_api_key_auth, mcp_tools_with_litellm_proxy, **kwargs): @@ -843,7 +843,7 @@ async def test_mcp_streaming_metadata_ordering(monkeypatch): dummy_tool = SimpleNamespace( name="local_search", description="search", - inputSchema={"type": "object", "properties": {}}, + input_schema={"type": "object", "properties": {}}, ) async def fake_process(user_api_key_auth, mcp_tools_with_litellm_proxy, **kwargs): diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/mcp_tests/test_mcp_logging.py index fc9f675f837..ed8829945e5 100644 --- a/tests/mcp_tests/test_mcp_logging.py +++ b/tests/mcp_tests/test_mcp_logging.py @@ -1,29 +1,51 @@ -import os -import pytest import asyncio +import os import subprocess import sys from pathlib import Path -from typing import Optional from unittest.mock import AsyncMock, patch +import pytest +from mcp.types import CallToolResult, TextContent +from mcp.types import Tool as MCPTool import litellm -from litellm.types.utils import StandardLoggingPayload from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, +) from litellm.proxy._experimental.mcp_server.server import ( mcp_server_tool_call, set_auth_context, ) -from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - MCPServerManager, -) from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth from litellm.types.mcp import MCPPostCallResponseObject -from litellm.types.utils import HiddenParams -from mcp.types import Tool as MCPTool, CallToolResult, TextContent +def _mcp_request_ctx(**overrides): + from types import SimpleNamespace + + from mcp.server.context import ServerRequestContext + + kwargs = { + "session": SimpleNamespace(), + "lifespan_context": {}, + "protocol_version": "2025-06-18", + "method": "", + "params": None, + "request_id": 1, + "meta": None, + "request": None, + } + kwargs.update(overrides) + return ServerRequestContext(**kwargs) + + +def _call_tool_params(name, arguments=None): + from mcp.types import CallToolRequestParams + + return CallToolRequestParams(name=name, arguments=arguments) + class TestMCPLogger(CustomLogger): def __init__(self): self.standard_logging_payload = None @@ -142,8 +164,8 @@ async def test_mcp_cost_tracking(): # Call mcp tool response = await mcp_server_tool_call( - name="zapier_gmail_server-add_tools", # Use correct prefixed name with - separator - arguments={"test": "test"}, + _mcp_request_ctx(), + _call_tool_params("zapier_gmail_server-add_tools", {"test": "test"}), ) # wait 1-2 seconds for logging to be processed @@ -285,8 +307,8 @@ async def test_mcp_cost_tracking_per_tool(): # Test 1: Call expensive_tool - should cost 5.0 response1 = await mcp_server_tool_call( - name="test_server-expensive_tool", # Use correct prefixed name with - separator - arguments={"data": "test_expensive"}, + _mcp_request_ctx(), + _call_tool_params("test_server-expensive_tool", {"data": "test_expensive"}), ) # wait for logging to be processed @@ -313,8 +335,8 @@ async def test_mcp_cost_tracking_per_tool(): # Test 2: Call cheap_tool - should cost 0.1 response2 = await mcp_server_tool_call( - name="test_server-cheap_tool", # Use correct prefixed name with - separator - arguments={"data": "test_cheap"}, + _mcp_request_ctx(), + _call_tool_params("test_server-cheap_tool", {"data": "test_cheap"}), ) # wait for logging to be processed @@ -356,7 +378,7 @@ async def test_mcp_cost_tracking_per_tool(): class MCPLoggerHook(TestMCPLogger): async def async_post_mcp_tool_call_hook( self, kwargs, response_obj: MCPPostCallResponseObject, start_time, end_time - ) -> Optional[MCPPostCallResponseObject]: + ) -> MCPPostCallResponseObject | None: print("post mcp tool call response_obj", response_obj) # update the MCPPostCallResponseObject with the response_cost response_obj.hidden_params.response_cost = 1.42 @@ -443,8 +465,8 @@ async def test_mcp_tool_call_hook(): # Call mcp tool using the correct separator format (- not /) response = await mcp_server_tool_call( - name="zapier_gmail_server-add_tools", # Use correct prefixed name with - separator - arguments={"test": "test"}, + _mcp_request_ctx(), + _call_tool_params("zapier_gmail_server-add_tools", {"test": "test"}), ) # wait 1-2 seconds for logging to be processed diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 1781dfe2fc2..94cf35b675d 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -121,7 +121,7 @@ async def test_mcp_server_manager_https_server(): print("RESULT FROM CALLING TOOL FROM MCP SERVER MANAGER== ", result) # Verify result - assert result.isError is False + assert result.is_error is False assert len(result.content) == 1 assert isinstance(result.content[0], TextContent) assert result.content[0].text == "Email sent successfully" @@ -288,7 +288,7 @@ async def test_mcp_http_transport_call_tool_mock(): ) # Assertions - assert result.isError is False + assert result.is_error is False assert len(result.content) == 1 # Type check before accessing text attribute assert isinstance(result.content[0], TextContent) @@ -350,7 +350,7 @@ async def test_mcp_http_transport_call_tool_error_mock(): ) # Assertions for error case - assert result.isError is True + assert result.is_error is True assert len(result.content) == 1 # Type check before accessing text attribute assert isinstance(result.content[0], TextContent) @@ -361,11 +361,11 @@ async def test_mcp_http_transport_call_tool_error_mock(): @pytest.mark.asyncio -async def test_mcp_http_transport_tool_not_found(): +async def test_mcp_http_transport_tool_not_found(config_only_mcp_manager_factory): """Test calling a tool that doesn't exist""" # Create a fresh manager for testing - test_manager = MCPServerManager() + test_manager = config_only_mcp_manager_factory() # Load server config await test_manager.load_servers_from_config( @@ -796,7 +796,7 @@ async def test_list_tools_rest_api_success(): ListMCPToolsRestAPIResponseObject( name="test_tool", description="A test tool", - inputSchema={"type": "object"}, + input_schema={"type": "object"}, mcp_info={"server_name": "test_server"}, ) ] @@ -1097,11 +1097,11 @@ async def test_list_tools_only_returns_allowed_servers(monkeypatch): @pytest.mark.asyncio -async def test_mcp_server_manager_access_groups_from_config(): +async def test_mcp_server_manager_access_groups_from_config(config_only_mcp_manager_factory): """ Test that access_groups are loaded from config and can be resolved. """ - test_manager = MCPServerManager() + test_manager = config_only_mcp_manager_factory() await test_manager.load_servers_from_config( { "config_server": { @@ -1168,7 +1168,7 @@ async def test_mcp_server_manager_access_groups_from_config(): @pytest.mark.asyncio -async def test_mcp_server_manager_config_integration_with_database(): +async def test_mcp_server_manager_config_integration_with_database(config_only_mcp_manager_factory): """ Test that config-based servers properly integrate with database servers, specifically testing access_groups and description fields. @@ -1176,7 +1176,7 @@ async def test_mcp_server_manager_config_integration_with_database(): import datetime from litellm.proxy._types import LiteLLM_MCPServerTable - test_manager = MCPServerManager() + test_manager = config_only_mcp_manager_factory() # Test 1: Load config with access_groups and description await test_manager.load_servers_from_config( @@ -2165,7 +2165,7 @@ async def test_list_tool_rest_api_with_server_specific_auth(): ListMCPToolsRestAPIResponseObject( name="send_email", description="Send an email", - inputSchema={"type": "object"}, + input_schema={"type": "object"}, mcp_info={"server_name": "zapier"}, ) ] @@ -2259,7 +2259,7 @@ async def test_list_tool_rest_api_with_default_auth(): ListMCPToolsRestAPIResponseObject( name="send_email", description="Send an email", - inputSchema={"type": "object"}, + input_schema={"type": "object"}, mcp_info={"server_name": "unknown_server"}, ) ] @@ -2371,7 +2371,7 @@ async def test_list_tool_rest_api_all_servers_with_auth(): ListMCPToolsRestAPIResponseObject( name="send_email", description="Send an email", - inputSchema={"type": "object"}, + input_schema={"type": "object"}, mcp_info={"server_name": "zapier"}, ) ], @@ -2379,7 +2379,7 @@ async def test_list_tool_rest_api_all_servers_with_auth(): ListMCPToolsRestAPIResponseObject( name="send_message", description="Send a message", - inputSchema={"type": "object"}, + input_schema={"type": "object"}, mcp_info={"server_name": "slack"}, ) ], @@ -2811,7 +2811,7 @@ async def test_mcp_access_group_permission_intersection_integration(): @pytest.mark.asyncio -async def test_mcp_server_manager_with_access_groups_integration(): +async def test_mcp_server_manager_with_access_groups_integration(config_only_mcp_manager_factory): """Integration test for MCPServerManager with access group filtering""" from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, @@ -2820,7 +2820,7 @@ async def test_mcp_server_manager_with_access_groups_integration(): from litellm.proxy._types import UserAPIKeyAuth # Create a test manager - test_manager = MCPServerManager() + test_manager = config_only_mcp_manager_factory() # Load servers with access groups await test_manager.load_servers_from_config( @@ -2863,13 +2863,13 @@ async def test_mcp_server_manager_with_access_groups_integration(): @pytest.mark.asyncio -async def test_get_allowed_mcp_servers_returns_registry_for_admin(): +async def test_get_allowed_mcp_servers_returns_registry_for_admin(config_only_mcp_manager_factory): from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) - test_manager = MCPServerManager() + test_manager = config_only_mcp_manager_factory() await test_manager.load_servers_from_config( { "alpha_server": { @@ -2898,14 +2898,14 @@ async def test_get_allowed_mcp_servers_returns_registry_for_admin(): @pytest.mark.asyncio -async def test_get_allowed_mcp_servers_returns_empty_for_non_admin_without_permissions(): +async def test_get_allowed_mcp_servers_returns_empty_for_non_admin_without_permissions(config_only_mcp_manager_factory): from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, MCPServerAccess, ) - test_manager = MCPServerManager() + test_manager = config_only_mcp_manager_factory() await test_manager.load_servers_from_config( { "alpha_server": { diff --git a/tests/mcp_tests/test_proxy_mcp_e2e.py b/tests/mcp_tests/test_proxy_mcp_e2e.py index e1099fe0a62..99c03b3438d 100644 --- a/tests/mcp_tests/test_proxy_mcp_e2e.py +++ b/tests/mcp_tests/test_proxy_mcp_e2e.py @@ -15,11 +15,12 @@ from datetime import datetime from pathlib import Path import httpx +import httpx2 import pytest import uvicorn import yaml from mcp import ClientSession -from mcp.client.streamable_http import streamablehttp_client +from mcp.client.streamable_http import streamable_http_client from mcp.types import CallToolResult from starlette.requests import Request @@ -36,6 +37,7 @@ from litellm.proxy.proxy_server import ( CONFIG_TEMPLATE_PATH = Path("tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml") MCP_SERVER_SCRIPT = Path("tests/mcp_tests/mcp_server.py") +MCP_PEER_PYTHON = os.environ.get("MCP_TEST_PEER_PYTHON", sys.executable) PROJECT_ROOT = Path(__file__).resolve().parents[2] PROXY_START_TIMEOUT = 30 @@ -125,7 +127,7 @@ def _math_http_server(offset: int) -> typing.Iterator[str]: with tempfile.TemporaryFile() as server_log: process = subprocess.Popen( - [sys.executable, str(MCP_SERVER_SCRIPT), "--transport", "http", "--host", host, "--port", str(port)], + [MCP_PEER_PYTHON, str(MCP_SERVER_SCRIPT), "--transport", "http", "--host", host, "--port", str(port)], cwd=str(PROJECT_ROOT), stdout=server_log, stderr=subprocess.STDOUT, @@ -175,7 +177,7 @@ def _proxy_server( config_dir = tmp_path_factory.mktemp("mcp_e2e") config_path = config_dir / "config.yaml" config = yaml.safe_load(CONFIG_TEMPLATE_PATH.read_text()) - config["mcp_servers"]["math_stdio"]["command"] = sys.executable + config["mcp_servers"]["math_stdio"]["command"] = MCP_PEER_PYTHON config["mcp_servers"]["math_streamable_http"]["url"] = f"{math_streamable_http_server}/mcp" config["mcp_servers"]["math_restricted"]["url"] = f"{math_restricted_server}/mcp" config["general_settings"]["custom_auth"] = f"{__name__}.authorize_proxy_key" @@ -202,17 +204,90 @@ def proxy_server_url(_proxy_server: ProxyRig, setup_and_teardown: None) -> str: return _proxy_server.url +@asynccontextmanager +async def _http_streams(url: str, headers: dict[str, str]): + async with httpx2.AsyncClient(headers=headers) as http_client: + async with streamable_http_client(url, http_client=http_client) as streams: + yield streams + + +@pytest.mark.asyncio +async def test_unchanged_sdk1_langchain_peer_can_list_and_call(proxy_server_url: str) -> None: + script = """ +import asyncio, json, sys +from mcp import ClientSession +from mcp.client.streamable_http import streamablehttp_client +from langchain_mcp_adapters.tools import load_mcp_tools + +async def main(): + async with streamablehttp_client(sys.argv[1] + '/mcp', headers={'Authorization': 'Bearer sk-1234'}) as (read, write, _): + async with ClientSession(read, write) as session: + await session.initialize() + tools = await load_mcp_tools(session) + results = {} + for name in ('math_stdio-add', 'math_streamable_http-add'): + tool = next(tool for tool in tools if tool.name == name) + results[name] = await tool.ainvoke({'a': 3, 'b': 4}) + print(json.dumps(results)) +asyncio.run(main()) +""" + completed = await asyncio.to_thread( + subprocess.run, [MCP_PEER_PYTHON, "-c", script, proxy_server_url], + capture_output=True, text=True, timeout=30, check=True, + ) + results = json.loads(completed.stdout) + assert [(item["type"], item["text"]) for item in results["math_stdio-add"]] == [("text", "7")] + assert [(item["type"], item["text"]) for item in results["math_streamable_http-add"]] == [("text", "107")] + + +@pytest.mark.parametrize("requested", ["2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25", "2026-07-28"]) +def test_initialize_keeps_legacy_negotiation(proxy_server_url: str, requested: str) -> None: + response = httpx.post( + proxy_server_url + "/mcp", + headers={"Authorization": PROXY_AUTHORIZATION_HEADER, "Accept": "application/json, text/event-stream"}, + json={"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": { + "protocolVersion": requested, "capabilities": {}, "clientInfo": {"name": "legacy-test", "version": "1"}, + }}, + timeout=10, + ) + assert response.status_code == 200 + result = _rpc_result(response) + assert result["protocolVersion"] == ("2025-11-25" if requested == "2026-07-28" else requested) + + +@pytest.mark.asyncio +async def test_legacy_prompts_and_resources_round_trip(proxy_server_url: str) -> None: + async with _http_streams( + proxy_server_url + "/mcp", + {"Authorization": PROXY_AUTHORIZATION_HEADER, "x-mcp-servers": "math_streamable_http"}, + ) as (read, write): + async with ClientSession(read, write) as session: + await session.initialize() + prompts = await session.list_prompts() + greeting = next(prompt for prompt in prompts.prompts if prompt.name.endswith("greeting")) + prompt = await session.get_prompt(greeting.name, {"name": "Ada"}) + assert prompt.messages[0].content.text == "Hello, Ada" + resources = await session.list_resources() + status = next(resource for resource in resources.resources if resource.name.endswith("status")) + contents = await session.read_resource(status.uri) + assert contents.contents[0].text == "ready" + templates = await session.list_resource_templates() + greeting_template = next(template for template in templates.resource_templates if "greeting" in template.name) + contents = await session.read_resource(greeting_template.uri_template.replace("{name}", "Ada")) + assert contents.contents[0].text == "Hello, Ada" + + class TestProxyMcpSimpleConnections: @pytest.mark.asyncio async def test_proxy_mcp_stdio_roundtrip(self, proxy_server_url: str) -> None: async with asyncio.timeout(20): - async with streamablehttp_client( + async with _http_streams( url=f"{proxy_server_url}/mcp", headers={ "Authorization": PROXY_AUTHORIZATION_HEADER, "x-mcp-servers": "math_stdio", }, - ) as (read, write, _get_session_id): + ) as (read, write): async with ClientSession(read, write) as session: await session.initialize() tools_result = await session.list_tools() @@ -227,13 +302,13 @@ class TestProxyMcpSimpleConnections: @pytest.mark.asyncio async def test_proxy_mcp_streamable_http_roundtrip(self, proxy_server_url: str) -> None: async with asyncio.timeout(20): - async with streamablehttp_client( + async with _http_streams( url=f"{proxy_server_url}/mcp", headers={ "Authorization": PROXY_AUTHORIZATION_HEADER, "x-mcp-servers": "math_streamable_http", }, - ) as (read, write, _get_session_id): + ) as (read, write): async with ClientSession(read, write) as session: await session.initialize() tools_result = await session.list_tools() @@ -248,10 +323,10 @@ class TestProxyMcpSimpleConnections: @pytest.mark.asyncio async def test_proxy_mcp_lists_all_servers_without_header(self, proxy_server_url: str) -> None: async with asyncio.timeout(20): - async with streamablehttp_client( + async with _http_streams( url=f"{proxy_server_url}/mcp", headers={"Authorization": PROXY_AUTHORIZATION_HEADER}, - ) as (read, write, _get_session_id): + ) as (read, write): async with ClientSession(read, write) as session: await session.initialize() tools_result = await session.list_tools() @@ -296,16 +371,16 @@ class TestProxyMcpStatelessBehavior: """Two independent clients connect and operate without sharing session state.""" async with asyncio.timeout(30): # --- Client A: connect, initialize, call tool --- - async with streamablehttp_client( + async with _http_streams( url=f"{proxy_server_url}/mcp", headers={ "Authorization": PROXY_AUTHORIZATION_HEADER, "x-mcp-servers": "math_stdio", }, - ) as (read_a, write_a, _get_sid_a): + ) as (read_a, write_a): async with ClientSession(read_a, write_a) as session_a: await session_a.initialize() - result_a = await session_a.call_tool("add", arguments={"a": 10, "b": 20}) + result_a = await session_a.call_tool("math_stdio-add", arguments={"a": 10, "b": 20}) assert result_a.content text_a = getattr(result_a.content[0], "text", None) assert text_a == "30" @@ -316,18 +391,18 @@ class TestProxyMcpStatelessBehavior: await asyncio.sleep(0.5) # --- Client B: completely independent connection --- - async with streamablehttp_client( + async with _http_streams( url=f"{proxy_server_url}/mcp", headers={ "Authorization": PROXY_AUTHORIZATION_HEADER, "x-mcp-servers": "math_stdio", }, - ) as (read_b, write_b, _get_sid_b): + ) as (read_b, write_b): async with ClientSession(read_b, write_b) as session_b: await session_b.initialize() tools = await session_b.list_tools() assert any(t.name.endswith("add") for t in tools.tools) - result_b = await session_b.call_tool("add", arguments={"a": 100, "b": 200}) + result_b = await session_b.call_tool("math_stdio-add", arguments={"a": 100, "b": 200}) assert result_b.content text_b = getattr(result_b.content[0], "text", None) assert text_b == "300" @@ -342,7 +417,7 @@ def _payload(result: typing.Any) -> typing.Any: def _proxy_session(proxy_server_url: str, **extra_headers: str): - return streamablehttp_client( + return _http_streams( url=f"{proxy_server_url}/mcp/proxy", headers={"Authorization": PROXY_AUTHORIZATION_HEADER, **extra_headers}, ) @@ -356,7 +431,7 @@ class TestProxyMcpSchemaDiscoveryMode: @pytest.mark.asyncio async def test_initialize_and_list_expose_only_discovery_tools(self, proxy_server_url: str) -> None: async with asyncio.timeout(20): - async with _proxy_session(proxy_server_url) as (read, write, _sid): + async with _proxy_session(proxy_server_url) as (read, write): async with ClientSession(read, write) as session: init = await session.initialize() assert init.capabilities.tools is not None @@ -369,7 +444,7 @@ class TestProxyMcpSchemaDiscoveryMode: @pytest.mark.asyncio async def test_search_schema_and_call_round_trip_keeps_server_identity(self, proxy_server_url: str) -> None: async with asyncio.timeout(30): - async with _proxy_session(proxy_server_url) as (read, write, _sid): + async with _proxy_session(proxy_server_url) as (read, write): async with ClientSession(read, write) as session: await session.initialize() @@ -399,8 +474,8 @@ class TestProxyMcpSchemaDiscoveryMode: "arguments": {"a": 5, "b": 6}, }, ) - assert stdio.isError is False and stdio.content[0].text == "7" - assert http.isError is False and http.content[0].text == "111" + assert stdio.is_error is False and stdio.content[0].text == "7" + assert http.is_error is False and http.content[0].text == "111" @pytest.mark.asyncio async def test_server_scope_header_narrows_discovery(self, proxy_server_url: str) -> None: @@ -408,7 +483,6 @@ class TestProxyMcpSchemaDiscoveryMode: async with _proxy_session(proxy_server_url, **{"x-mcp-servers": "math_streamable_http"}) as ( read, write, - _sid, ): async with ClientSession(read, write) as session: await session.initialize() @@ -417,11 +491,11 @@ class TestProxyMcpSchemaDiscoveryMode: @pytest.mark.asyncio async def test_rejections_never_reach_upstream(self, proxy_server_url: str) -> None: - from mcp.shared.exceptions import McpError + from mcp.shared.exceptions import MCPError from mcp.types import METHOD_NOT_FOUND async with asyncio.timeout(30): - async with _proxy_session(proxy_server_url) as (read, write, _sid): + async with _proxy_session(proxy_server_url) as (read, write): async with ClientSession(read, write) as session: await session.initialize() hits = _payload(await session.call_tool("search_tools", arguments={"query": "add"})) @@ -430,22 +504,22 @@ class TestProxyMcpSchemaDiscoveryMode: bad_args = await session.call_tool( "call_tool", arguments={"tool_id": tool_id, "arguments": {"a": "three", "b": 4}} ) - assert bad_args.isError is True and "Invalid arguments" in bad_args.content[0].text + assert bad_args.is_error is True and "Invalid arguments" in bad_args.content[0].text stale = await session.call_tool("get_tool_schema", arguments={"tool_id": "0" * 32}) - assert stale.isError is True and "unauthorized tool_id" in stale.content[0].text + assert stale.is_error is True and "unauthorized tool_id" in stale.content[0].text for not_an_object in ("wrong", False): refused_args = await session.call_tool( "call_tool", arguments={"tool_id": tool_id, "arguments": not_an_object} ) - assert refused_args.isError is True and "object" in refused_args.content[0].text + assert refused_args.is_error is True and "object" in refused_args.content[0].text direct = await session.call_tool("math_stdio-add", arguments={"a": 1, "b": 2}) - assert direct.isError is True and "unavailable on /mcp/proxy" in direct.content[0].text + assert direct.is_error is True and "unavailable on /mcp/proxy" in direct.content[0].text for operation in (session.list_prompts, session.list_resources): - with pytest.raises(McpError) as refused: + with pytest.raises(MCPError) as refused: await operation() assert refused.value.error.code == METHOD_NOT_FOUND @@ -494,7 +568,7 @@ proxy_call_recorder = ProxyCallRecorder() @asynccontextmanager async def _scoped_session(url: str, key: str = "sk-1234", **headers: str) -> typing.AsyncIterator[ClientSession]: async with asyncio.timeout(30): - async with _proxy_session(url, Authorization=f"Bearer {key}", **headers) as (read, write, _sid): + async with _proxy_session(url, Authorization=f"Bearer {key}", **headers) as (read, write): async with ClientSession(read, write) as session: await session.initialize() yield session @@ -502,7 +576,7 @@ async def _scoped_session(url: str, key: str = "sk-1234", **headers: str) -> typ async def _search(session: ClientSession, query: str) -> dict[str, str]: result = await session.call_tool("search_tools", arguments={"query": query}) - assert result.isError is False, result + assert result.is_error is False, result return {hit["name"]: hit["tool_id"] for hit in _payload(result)} @@ -542,7 +616,7 @@ def _rpc_result(response: httpx.Response) -> dict[str, typing.Any]: def _assert_unauthorized(result: CallToolResult) -> None: - assert result.isError is True + assert result.is_error is True assert result.content[0].text == "Unknown or unauthorized tool_id" @@ -611,7 +685,7 @@ class TestProxyMcpAuthorizationScope: assert schema["name"] == name assert schema["tool_id"] == ids[name] result = await _call(session, ids[name]) - assert result.isError is False + assert result.is_error is False assert result.content[0].text == expected @pytest.mark.asyncio @@ -652,7 +726,7 @@ class TestProxyMcpAuthorizationScope: result = await session.call_tool( "call_tool", {"tool_id": ids[f"{name}-request_headers"], "arguments": {}} ) - assert result.isError is False + assert result.is_error is False assert _payload(result) == expected @pytest.mark.asyncio @@ -660,7 +734,7 @@ class TestProxyMcpAuthorizationScope: async with _scoped_session(proxy_server_url, "sk-restricted") as session: tool_id = (await _search(session, "add"))["math_restricted-add"] result = await _call(session, tool_id, 123, 456) - assert result.isError is False and result.content[0].text == "779" + assert result.is_error is False and result.content[0].text == "779" async with asyncio.timeout(10): while True: payload = json.loads(await asyncio.to_thread(proxy_call_recorder.events.get, True, 5)) @@ -714,7 +788,7 @@ class TestProxyMcpAuthorizationScope: hits = _payload(await handle_mcp_proxy_tool("search_tools", {"query": "add"}, auth)) tool_id = next(hit["tool_id"] for hit in hits if hit["name"] == "math_stdio-add") result = await handle_mcp_proxy_tool("call_tool", {"tool_id": tool_id, "arguments": arguments}, auth) - assert result.isError is True + assert result.is_error is True assert result.content[0].text == "arguments must be an object" asyncio.run_coroutine_threadsafe(check(), _proxy_server.loop).result(timeout=30) diff --git a/tests/pass_through_tests/test_mcp_routes.py b/tests/pass_through_tests/test_mcp_routes.py index 687efe6195d..e9d18193e7c 100644 --- a/tests/pass_through_tests/test_mcp_routes.py +++ b/tests/pass_through_tests/test_mcp_routes.py @@ -2,14 +2,15 @@ import asyncio import os -from langchain_mcp_adapters.tools import load_mcp_tools -from langchain_openai import ChatOpenAI -from langgraph.prebuilt import create_react_agent from mcp import ClientSession from mcp.client.sse import sse_client async def main(): + from langchain_mcp_adapters.tools import load_mcp_tools + from langchain_openai import ChatOpenAI + from langgraph.prebuilt import create_react_agent + model = ChatOpenAI(model="gpt-4o", api_key="sk-12") async with sse_client(url="http://localhost:4000/mcp/") as (read, write): diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py b/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py new file mode 100644 index 00000000000..5695b184479 --- /dev/null +++ b/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py @@ -0,0 +1,168 @@ +from typing import Final, Literal + +import pytest +from litellm_enterprise.enterprise_callbacks.llm_guard import _ENTERPRISE_LLMGuard +from starlette.exceptions import HTTPException + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.utils import hash_token +from litellm.types.utils import CallTypesLiteral + + +@pytest.mark.parametrize( + "call_type, payload_key", + ( + ("completion", "messages"), + ("acompletion", "messages"), + ("text_completion", "prompt"), + ("atext_completion", "prompt"), + ("embeddings", "input"), + ("embedding", "input"), + ("aembedding", "input"), + ("image_generation", "prompt"), + ("aimage_generation", "prompt"), + ), +) +@pytest.mark.parametrize("is_valid", (True, False)) +@pytest.mark.asyncio +async def test_llm_guard_call_type_aliases( + call_type: CallTypesLiteral, + payload_key: Literal["messages", "input", "prompt"], + is_valid: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "llm_guard_mode", "all") + llm_guard: Final = _ENTERPRISE_LLMGuard( + mock_testing=True, + mock_redacted_text={ + "sanitized_prompt": "email: [REDACTED]", + "is_valid": is_valid, + }, + ) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=hash_token("sk-12345")) + data: Final = { + payload_key: [{"role": "user", "content": "email: person@example.com"}] + if payload_key == "messages" + else "email: person@example.com" + } + + if not is_valid: + with pytest.raises(HTTPException) as exc_info: + await llm_guard.async_moderation_hook(data=data, user_api_key_dict=user_api_key_dict, call_type=call_type) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == {"error": "Violated content safety policy"} + return + + result: Final = await llm_guard.async_moderation_hook( + data=data, user_api_key_dict=user_api_key_dict, call_type=call_type + ) + assert result is data + assert data[payload_key] == ( + [{"role": "user", "content": "email: [REDACTED]"}] if payload_key == "messages" else "email: [REDACTED]" + ) + + +@pytest.mark.parametrize("call_type", ("amoderation", "atranscription", "aresponses", "aanthropic_messages")) +@pytest.mark.asyncio +async def test_llm_guard_ignores_call_types_the_proxy_never_moderates( + call_type: CallTypesLiteral, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "llm_guard_mode", "all") + llm_guard: Final = _ENTERPRISE_LLMGuard( + mock_testing=True, + mock_redacted_text={"sanitized_prompt": "[REDACTED]", "is_valid": False}, + ) + data: Final = {"input": "email: person@example.com"} + result: Final = await llm_guard.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type + ) + assert result is data + assert data["input"] == "email: person@example.com" + + +@pytest.mark.parametrize("call_type", ("text_completion", "atext_completion")) +@pytest.mark.parametrize("is_valid", (True, False)) +@pytest.mark.asyncio +async def test_llm_guard_scans_list_prompt( + call_type: CallTypesLiteral, is_valid: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "llm_guard_mode", "all") + llm_guard: Final = _ENTERPRISE_LLMGuard( + mock_testing=True, + mock_redacted_text={"sanitized_prompt": "[REDACTED]", "is_valid": is_valid}, + ) + data: Final = {"prompt": ["email: person@example.com", "say ok", [1, 2, 3]]} + + if not is_valid: + with pytest.raises(HTTPException) as exc_info: + await llm_guard.async_moderation_hook(data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type) + assert exc_info.value.status_code == 400 + return + + result: Final = await llm_guard.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type + ) + assert result is data + assert data["prompt"] == ["[REDACTED]", "[REDACTED]", [1, 2, 3]] + + +@pytest.mark.parametrize("call_type", ("aembedding", "atext_completion")) +@pytest.mark.parametrize("is_valid", (True, False)) +@pytest.mark.asyncio +async def test_llm_guard_scans_input_and_prompt_alongside_messages( + call_type: CallTypesLiteral, is_valid: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "llm_guard_mode", "all") + llm_guard: Final = _ENTERPRISE_LLMGuard( + mock_testing=True, + mock_redacted_text={"sanitized_prompt": "[REDACTED]", "is_valid": is_valid}, + ) + data: Final = { + "messages": [], + "input": "email: person@example.com", + "prompt": ["say ok"], + } + + if not is_valid: + with pytest.raises(HTTPException) as exc_info: + await llm_guard.async_moderation_hook(data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type) + assert exc_info.value.status_code == 400 + return + + result: Final = await llm_guard.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type + ) + assert result is data + assert data["messages"] == [] + assert data["input"] == "[REDACTED]" + assert data["prompt"] == ["[REDACTED]"] + + +@pytest.mark.parametrize( + "call_type", + ( + "responses", + "aresponses", + "anthropic_messages", + "aanthropic_messages", + "aspeech", + "aimage_edit", + "pass_through_endpoint", + ), +) +@pytest.mark.asyncio +async def test_llm_guard_skips_unsupported_call_types( + call_type: CallTypesLiteral, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "llm_guard_mode", "all") + llm_guard: Final = _ENTERPRISE_LLMGuard( + mock_testing=True, + mock_redacted_text={"is_valid": False}, + ) + data: Final = {"messages": [{"role": "user", "content": "unchanged"}]} + result: Final = await llm_guard.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type + ) + assert result is data + assert data == {"messages": [{"role": "user", "content": "unchanged"}]} diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index d9ffb0d64fe..b78d61c7bd4 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -4,22 +4,20 @@ import json import os import sys from collections.abc import AsyncIterator -from importlib import metadata from pathlib import Path from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import anyio -import httpx +import httpx2 import pytest -import respx -from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth -from mcp import McpError +from mcp import MCPError from mcp.client.streamable_http import streamable_http_client -from pydantic import ValidationError from mcp.shared.message import SessionMessage from mcp.types import ( - LATEST_PROTOCOL_VERSION, + CONNECTION_CLOSED, + INTERNAL_ERROR, + REQUEST_TIMEOUT, CallToolResult, ErrorData, Implementation, @@ -30,17 +28,16 @@ from mcp.types import ( LoggingMessageNotificationParams, ServerCapabilities, ) +from mcp_types.version import LATEST_HANDSHAKE_VERSION +from pydantic import TypeAdapter, ValidationError # Add the parent directory to the path so we can import litellm - import litellm.experimental_mcp_client.client as mcp_client_module from litellm.experimental_mcp_client.client import ( - MCP_STREAMABLE_HTTP_REQUIREMENT, MCPClient, _first_non_cancelled_cause, _TransportContext, as_mcp_read_timeout, - missing_streamable_http_client_error, strip_auth_scheme, ) from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( @@ -50,8 +47,23 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( _format_byok_openapi_auth_header, ) -from litellm.types.mcp_server.mcp_server_manager import MCPServer +from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth from litellm.types.mcp import MCPAuth, MCPStdioConfig, MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + +_JSONRPC_MESSAGE_ADAPTER: Final = TypeAdapter(JSONRPCMessage) + + +class _MockTransportClient(MCPClient): + """An MCPClient whose streamable-HTTP transport runs on an httpx2 MockTransport.""" + + def __init__(self, respond, **kwargs): + super().__init__(**kwargs) + self._respond = respond + + def _create_transport_context(self): + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(self._respond)) + return streamable_http_client(self.server_url, http_client=http_client), http_client class _FakeExceptionGroup(Exception): @@ -171,14 +183,14 @@ class TestMCPClient: call_kwargs = mock_streamable_http_client.call_args[1] assert "http_client" in call_kwargs http_client = call_kwargs["http_client"] - assert isinstance(http_client, httpx.AsyncClient) + assert isinstance(http_client, httpx2.AsyncClient) # Test the factory still creates a client with proper SSL config httpx_factory = client._create_httpx_client_factory() test_client = httpx_factory(headers={"test": "header"}) assert test_client is not None - assert isinstance(test_client, httpx.AsyncClient) + assert isinstance(test_client, httpx2.AsyncClient) assert test_client.headers is not None await test_client.aclose() @@ -228,7 +240,7 @@ class TestMCPClient: # Verify the client was created successfully assert test_client is not None - assert isinstance(test_client, httpx.AsyncClient) + assert isinstance(test_client, httpx2.AsyncClient) # Verify it has the expected properties assert test_client.headers is not None # Clean up @@ -272,13 +284,13 @@ class TestMCPClient: call_kwargs = mock_streamable_http_client.call_args[1] assert "http_client" in call_kwargs http_client = call_kwargs["http_client"] - assert isinstance(http_client, httpx.AsyncClient) + assert isinstance(http_client, httpx2.AsyncClient) httpx_factory = client._create_httpx_client_factory() test_client = httpx_factory(headers={"test": "header"}) assert test_client is not None - assert isinstance(test_client, httpx.AsyncClient) + assert isinstance(test_client, httpx2.AsyncClient) assert test_client.headers is not None await test_client.aclose() @@ -460,12 +472,12 @@ class TestFirstNonCancelledCause: assert _first_non_cancelled_cause(asyncio.CancelledError()) is None def test_unwraps_group_to_non_cancelled_leaf(self): - target = httpx.ConnectError("refused") + target = httpx2.ConnectError("refused") group = _FakeExceptionGroup("g", [asyncio.CancelledError(), target]) assert _first_non_cancelled_cause(group) is target def test_unwraps_nested_group(self): - target = httpx.LocalProtocolError("Illegal header value") + target = httpx2.LocalProtocolError("Illegal header value") inner = _FakeExceptionGroup("inner", [asyncio.CancelledError(), target]) outer = _FakeExceptionGroup("outer", [asyncio.CancelledError(), inner]) assert _first_non_cancelled_cause(outer) is target @@ -476,7 +488,7 @@ class TestFirstNonCancelledCause: @pytest.mark.skipif(sys.version_info < (3, 11), reason="builtin ExceptionGroup requires 3.11+") def test_unwraps_builtin_exception_group(self): - target = httpx.ConnectError("refused") + target = httpx2.ConnectError("refused") group = ExceptionGroup("transport failed", [target]) # noqa: F821 assert _first_non_cancelled_cause(group) is target @@ -512,13 +524,13 @@ class TestExecuteSessionOperationSurfacesTransportError: mock_session_cls, AsyncMock(side_effect=asyncio.CancelledError("cancelled by group")), ) - connect_error = httpx.ConnectError("All connection attempts failed") + connect_error = httpx2.ConnectError("All connection attempts failed") transport_ctx = self._make_transport(_FakeExceptionGroup("transport", [connect_error])) async def _op(session): return "done" - with pytest.raises(httpx.ConnectError): + with pytest.raises(httpx2.ConnectError): await client._execute_session_operation(transport_ctx, _op) @pytest.mark.asyncio @@ -541,7 +553,7 @@ class TestExecuteSessionOperationSurfacesTransportError: init_result = MagicMock() init_result.instructions = None self._make_session(mock_session_cls, AsyncMock(return_value=init_result)) - transport_ctx = self._make_transport(_FakeExceptionGroup("late", [httpx.ConnectError("late cleanup error")])) + transport_ctx = self._make_transport(_FakeExceptionGroup("late", [httpx2.ConnectError("late cleanup error")])) async def _op(session): return "done" @@ -551,11 +563,11 @@ class TestExecuteSessionOperationSurfacesTransportError: class TestMCPClientResolvedAuth: - """A pre-resolved httpx.Auth is attached to the upstream client's auth= slot.""" + """A pre-resolved httpx2.Auth is attached to the upstream client's auth= slot.""" @pytest.mark.asyncio async def test_resolved_auth_feeds_the_auth_slot(self): - resolved = httpx.Auth() + resolved = httpx2.Auth() client = MCPClient(server_url="https://upstream.example.com", resolved_auth=resolved) http_client = client._create_httpx_client_factory()() try: @@ -565,11 +577,11 @@ class TestMCPClientResolvedAuth: @pytest.mark.asyncio async def test_resolved_auth_takes_precedence_over_aws_auth(self): - resolved = httpx.Auth() + resolved = httpx2.Auth() client = MCPClient( server_url="https://upstream.example.com", resolved_auth=resolved, - aws_auth=httpx.Auth(), + aws_auth=httpx2.Auth(), ) http_client = client._create_httpx_client_factory()() try: @@ -579,7 +591,7 @@ class TestMCPClientResolvedAuth: @pytest.mark.asyncio async def test_without_resolved_auth_falls_back_to_aws_auth(self): - aws = httpx.Auth() + aws = httpx2.Auth() client = MCPClient(server_url="https://upstream.example.com", aws_auth=aws) http_client = client._create_httpx_client_factory()() try: @@ -672,7 +684,7 @@ async def test_call_tool_raise_on_error_logs_at_debug_not_error(): with patch.object(client, "run_with_session", side_effect=_raise): with patch.object(mcp_client_module, "verbose_logger") as mock_log: result = await client.call_tool(params, raise_on_error=False) - assert result.isError is True + assert result.is_error is True assert mock_log.error.called, "swallow path must keep error-level visibility" @@ -766,15 +778,15 @@ class _ScriptedUpstream: return await self._task_group.__aexit__(None, None, None) async def _send(self, message): - await self._to_client_tx.send(SessionMessage(JSONRPCMessage(message))) + await self._to_client_tx.send(SessionMessage(message)) async def _serve(self): async for session_message in self._from_client_rx: - request = session_message.message.root + request = session_message.message method = getattr(request, "method", None) if method == "initialize": result = InitializeResult( - protocolVersion=LATEST_PROTOCOL_VERSION, + protocolVersion=LATEST_HANDSHAKE_VERSION, capabilities=ServerCapabilities(), serverInfo=Implementation(name="scripted-upstream", version="1.0.0"), ) @@ -835,36 +847,36 @@ async def test_upstream_json_rpc_error_408_is_not_reported_as_a_client_timeout() """The SDK reports its own elapsed read timeout and relays an upstream JSON-RPC error through the same exception class and the same numeric field, and JSON-RPC error codes are a different namespace from HTTP status codes. An upstream answering with application code 408 must keep - travelling as ``McpError`` so it is never blamed on the gateway as a 504. + travelling as ``MCPError`` so it is never blamed on the gateway as a 504. This is the other half of the pair: the same real transport and the same real session, so one mechanism pins both directions. """ client = _ScriptedClient( timeout=30, - tools_list_error=ErrorData(code=int(httpx.codes.REQUEST_TIMEOUT), message="re-authenticate and retry"), + tools_list_error=ErrorData(code=REQUEST_TIMEOUT, message="re-authenticate and retry"), ) - with pytest.raises(McpError) as exc_info: + with pytest.raises(MCPError) as exc_info: await asyncio.wait_for(client.list_tools(raise_on_error=True), timeout=10) assert not isinstance(exc_info.value, TimeoutError), "an upstream application error is not a gateway timeout" - assert exc_info.value.error.code == int(httpx.codes.REQUEST_TIMEOUT) + assert exc_info.value.error.code == REQUEST_TIMEOUT fault = classify_list_exception(exc_info.value) assert fault.tag != "timeout", "an upstream's own application error must never be reported as a gateway timeout" assert list_fault_http_status(fault) != 504 -def _raise_mcp_error_while_handling_a_timeout(code: int, message: str) -> McpError: - """An ``McpError`` carrying the context chain it would have if it were raised while a +def _raise_mcp_error_while_handling_a_timeout(code: int, message: str) -> MCPError: + """An ``MCPError`` carrying the context chain it would have if it were raised while a ``TimeoutError`` was in flight, which is how the SDK raises its own read timeout.""" try: try: raise TimeoutError() except TimeoutError: - raise McpError(ErrorData(code=code, message=message)) - except McpError as raised: + raise MCPError(code=code, message=message) + except MCPError as raised: return raised @@ -873,20 +885,20 @@ def test_as_mcp_read_timeout_separates_the_sdk_timeout_from_a_relayed_upstream_e upstream JSON-RPC error that happens to use 408, and the context chain alone cannot separate it from any other relayed error that surfaces while a timeout is being handled, so both must hold. """ - timeout_code = int(httpx.codes.REQUEST_TIMEOUT) + timeout_code = REQUEST_TIMEOUT translated = as_mcp_read_timeout(_raise_mcp_error_while_handling_a_timeout(timeout_code, "Timed out while waiting")) assert isinstance(translated, TimeoutError) assert str(translated) == "Timed out while waiting" - relayed_408 = McpError(ErrorData(code=timeout_code, message="upstream said 408")) + relayed_408 = MCPError(code=timeout_code, message="upstream said 408") assert as_mcp_read_timeout(relayed_408) is None, "an upstream 408 with no elapsed timeout is not our timeout" relayed_other = _raise_mcp_error_while_handling_a_timeout(-32603, "upstream internal error") assert as_mcp_read_timeout(relayed_other) is None, "a non-timeout code is not our timeout, whatever the chain" - assert as_mcp_read_timeout(McpError(ErrorData(code=-32603, message="boom"))) is None - assert as_mcp_read_timeout(RuntimeError("not an McpError")) is None + assert as_mcp_read_timeout(MCPError(code=-32603, message="boom")) is None + assert as_mcp_read_timeout(RuntimeError("not an MCPError")) is None @pytest.mark.asyncio @@ -1065,28 +1077,6 @@ def test_openapi_byok_auth_header_emits_exactly_one_scheme(auth_type, auth_value assert _format_byok_openapi_auth_header(server, auth_value) == expected -def test_missing_streamable_http_client_error_names_requirement_and_remedy(): - message = str(missing_streamable_http_client_error()) - - assert MCP_STREAMABLE_HTTP_REQUIREMENT in message - assert "pip install 'litellm[mcp]'" in message - assert metadata.version("mcp") in message - - -@pytest.mark.asyncio -async def test_http_transport_without_streamable_http_client_raises_actionable_import_error(): - client = MCPClient( - server_url="https://mcp-server.example.com", - transport_type=MCPTransport.http, - ) - - with patch.object( # test-quality-ok: simulates mcp<1.24.0 whose module lacks this import-time symbol - mcp_client_module, "streamable_http_client", None - ): - with pytest.raises(ImportError, match=r"pip install 'litellm\[mcp\]'"): - await client.list_tools(raise_on_error=True) - - def test_mcp_extra_matches_proxy_extra_and_supports_streamable_http(): try: import tomllib @@ -1096,17 +1086,50 @@ def test_mcp_extra_matches_proxy_extra_and_supports_streamable_http(): pyproject_path = Path(__file__).parents[3] / "pyproject.toml" with pyproject_path.open("rb") as f: - extras = tomllib.load(f)["project"]["optional-dependencies"] + project = tomllib.load(f) + extras = project["project"]["optional-dependencies"] - mcp_extra = extras["mcp"] - assert len(mcp_extra) == 1 + sdk2_names: Final = frozenset(("mcp", "httpx2", "pydantic")) + mcp_extra: Final = {Requirement(req).name: req for req in extras["mcp"]} + assert mcp_extra == { + name: req + for req in extras["proxy"] + if (name := Requirement(req).name) in sdk2_names + } - proxy_mcp_requirements = [req for req in extras["proxy"] if Requirement(req).name == "mcp"] - assert mcp_extra == proxy_mcp_requirements + specifier: Final = Requirement(mcp_extra["mcp"]).specifier + assert not specifier.contains("1.28.1") + assert specifier.contains("2.2.0") + with (pyproject_path.parent / "uv.lock").open("rb") as f: + locked = tomllib.load(f) + mcp_versions: Final = [package["version"] for package in locked["package"] if package["name"] == "mcp"] + assert len(mcp_versions) == 1 + assert specifier.contains(mcp_versions[0]) - specifier = Requirement(mcp_extra[0]).specifier - assert not specifier.contains("1.23.0") - assert specifier.contains("1.28.1") + +@pytest.mark.parametrize("module", ["mcp", "mcp_types", "httpx2", "httpcore2"]) +def test_base_sdk_guard_rejects_mcp_dependencies(tmp_path: Path, module: str) -> None: + import subprocess + import sys + + (tmp_path / f"{module}.py").write_text("") + checker = Path(__file__).parents[2] / "base_sdk_tests" / "check_base_sdk_install.py" + result = subprocess.run( + [ + sys.executable, + "-S", + "-c", + "import runpy, sys; sys.path.insert(0, sys.argv[2]); " + "runpy.run_path(sys.argv[1])['check_environment_is_base_only']()", + str(checker), + str(tmp_path), + ], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode != 0, f"base-only guard accepted installed {module}" + assert f"{module} installed" in result.stderr @pytest.mark.parametrize( @@ -1161,13 +1184,13 @@ async def test_a_custom_credential_header_is_stripped_when_a_redirect_crosses_or operator moved to its own slot would be replayed to whatever host the upstream redirects to. Verified against real httpx redirect handling, not a hand-built request. """ - seen: "list[tuple[str, str]]" = [] + seen: list[tuple[str, str]] = [] - def handler(request: httpx.Request) -> httpx.Response: + def handler(request: httpx2.Request) -> httpx2.Response: seen.append((request.url.host, request.headers.get("esb-oauth", ""))) if request.url.host == "upstream.example.com": - return httpx.Response(302, headers={"Location": "https://attacker.example.com/collect"}) - return httpx.Response(200) + return httpx2.Response(302, headers={"Location": "https://attacker.example.com/collect"}) + return httpx2.Response(200) client = MCPClient( server_url="https://upstream.example.com/mcp", @@ -1177,7 +1200,7 @@ async def test_a_custom_credential_header_is_stripped_when_a_redirect_crosses_or client.update_auth_value("minted-token") factory = client._create_httpx_client_factory() async with factory(headers=client._get_auth_headers(), timeout=None) as http_client: - http_client._transport = httpx.MockTransport(handler) + http_client._transport = httpx2.MockTransport(handler) await http_client.get("https://upstream.example.com/mcp") assert seen[0] == ("upstream.example.com", "Bearer minted-token") @@ -1253,9 +1276,9 @@ async def test_the_guard_agrees_with_httpx_about_authorization(start: str, targe outcomes. A future httpx that changes its redirect rule reds here instead of silently leaving the custom slot forwarded where Authorization is not (or stripped where it is not needed). """ - seen: "list[tuple[str, str, str]]" = [] + seen: list[tuple[str, str, str]] = [] - def handler(request: httpx.Request) -> httpx.Response: + def handler(request: httpx2.Request) -> httpx2.Response: seen.append( ( str(request.url), @@ -1264,13 +1287,13 @@ async def test_the_guard_agrees_with_httpx_about_authorization(start: str, targe ) ) if str(request.url) == start: - return httpx.Response(302, headers={"Location": target}) - return httpx.Response(200) + return httpx2.Response(302, headers={"Location": target}) + return httpx2.Response(200) client = MCPClient(server_url=start, auth_type=MCPAuth.oauth2, auth_header_name="esb-oauth") factory = client._create_httpx_client_factory() async with factory(headers={"Authorization": "Bearer AUTH", "esb-oauth": "Bearer ESB"}, timeout=None) as http: - http._transport = httpx.MockTransport(handler) + http._transport = httpx2.MockTransport(handler) await http.get(start) _url, authorization, esb = seen[-1] @@ -1298,11 +1321,12 @@ def test_a_differently_cased_injected_header_cannot_shadow_the_slot() -> None: @pytest.mark.parametrize( ("content_type", "body", "expected_type"), [ - ("text/html", b"secret-page", ValueError), - ("application/json", b"secret-invalid-json", ValidationError), - ("application/json", b"", ValidationError), - ("application/json", b'{"secret":"invalid-rpc"}', ValidationError), - ("application/json", b'{"jsonrpc":"2.0","id":0,"result":{"secret":"invalid-schema"}}', ValidationError), + ("text/html", b"secret-page", MCPError), + ("application/json", b"secret-invalid-json", MCPError), + ("application/json", b"", MCPError), + ("application/json", b'{"secret":"invalid-rpc"}', MCPError), + ("application/json", b'{"jsonrpc":"2.0","id":0}', MCPError), + ("application/json", b'{"jsonrpc":"2.0","id":0,"result":{"secret":"bad-schema"}}', ValidationError), ], ) async def test_invalid_http_response_surfaces_without_waiting_for_timeout( @@ -1310,10 +1334,12 @@ async def test_invalid_http_response_surfaces_without_waiting_for_timeout( ) -> None: from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message - def respond(request: httpx.Request) -> httpx.Response: - return httpx.Response(200, headers={"Content-Type": content_type}, content=body) + def respond(request: httpx2.Request) -> httpx2.Response: + if expected_type is ValidationError: + return httpx2.Response(200, json={**json.loads(body), "id": json.loads(request.content)["id"]}) + return httpx2.Response(200, headers={"Content-Type": content_type}, content=body) - async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + async with httpx2.AsyncClient(transport=httpx2.MockTransport(respond)) as http_client: client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) with pytest.raises(expected_type) as caught: await asyncio.wait_for( @@ -1331,27 +1357,27 @@ async def test_invalid_http_response_surfaces_without_waiting_for_timeout( @pytest.mark.asyncio -@pytest.mark.parametrize("status_code", [200, 401, 503]) +@pytest.mark.parametrize("status_code", [200, 401, 403, 429, 503]) async def test_http_response_handler_preserves_success_and_http_errors(status_code: int) -> None: - def respond(request: httpx.Request) -> httpx.Response: + def respond(request: httpx2.Request) -> httpx2.Response: if request.method == "DELETE": - return httpx.Response(200) + return httpx2.Response(200) payload: Final = json.loads(request.content) if "id" not in payload: - return httpx.Response(202) + return httpx2.Response(202) result: Final = ( { - "protocolVersion": LATEST_PROTOCOL_VERSION, + "protocolVersion": payload["params"]["protocolVersion"], "capabilities": {}, "serverInfo": {"name": "test", "version": "1"}, } if payload["method"] == "initialize" else {"tools": []} ) - return httpx.Response(status_code, json={"jsonrpc": "2.0", "id": payload["id"], "result": result}) + return httpx2.Response(status_code, json={"jsonrpc": "2.0", "id": payload["id"], "result": result}) - async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: - client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) + async with client._create_httpx_client_factory(transport=httpx2.MockTransport(respond))() as http_client: operation: Final = client._execute_session_operation( streamable_http_client(client.server_url, http_client=http_client), lambda session: session.list_tools() ) @@ -1359,11 +1385,35 @@ async def test_http_response_handler_preserves_success_and_http_errors(status_co result: Final = await asyncio.wait_for(operation, timeout=3) assert result.tools == [] else: - with pytest.raises(httpx.HTTPStatusError) as caught: + with pytest.raises(httpx2.HTTPStatusError) as caught: await asyncio.wait_for(operation, timeout=3) assert caught.value.response.status_code == status_code +@pytest.mark.asyncio +async def test_http_status_check_allows_auth_refresh_before_rejecting() -> None: + from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ClientCredentialsBearerAuth + + seen = [] + + async def refresh(failed): + assert failed == "stale" + return "fresh" + + def respond(request): + seen.append(request.headers["authorization"]) + return httpx2.Response(401 if len(seen) == 1 else 200, json={"ok": True}) + + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ClientCredentialsConfig + + auth = ClientCredentialsBearerAuth("stale", refresh, ClientCredentialsConfig()) + client = MCPClient(server_url="https://example.com/mcp", resolved_auth=auth) + async with client._create_httpx_client_factory(transport=httpx2.MockTransport(respond))() as http_client: + response = await http_client.post(client.server_url, json={"method": "tools/list"}) + assert response.status_code == 200 + assert seen == ["Bearer stale", "Bearer fresh"] + + @pytest.mark.asyncio async def test_http_response_handler_preserves_notifications_and_tool_listing() -> None: notification: Final = { @@ -1373,20 +1423,20 @@ async def test_http_response_handler_preserves_notifications_and_tool_listing() } logging_callback: Final = AsyncMock() - def respond(request: httpx.Request) -> httpx.Response: + def respond(request: httpx2.Request) -> httpx2.Response: if request.method == "DELETE": - return httpx.Response(200) + return httpx2.Response(200) payload: Final = json.loads(request.content) if "id" not in payload: - return httpx.Response(202) + return httpx2.Response(202) if payload["method"] == "initialize": - return httpx.Response( + return httpx2.Response( 200, json={ "jsonrpc": "2.0", "id": payload["id"], "result": { - "protocolVersion": LATEST_PROTOCOL_VERSION, + "protocolVersion": payload["params"]["protocolVersion"], "capabilities": {"logging": {}, "tools": {}}, "serverInfo": {"name": "test", "version": "1"}, }, @@ -1397,13 +1447,13 @@ async def test_http_response_handler_preserves_notifications_and_tool_listing() "id": payload["id"], "result": {"tools": [{"name": "search", "inputSchema": {"type": "object"}}]}, } - return httpx.Response( + return httpx2.Response( 200, headers={"Content-Type": "text/event-stream"}, content="".join(f"event: message\ndata: {json.dumps(message)}\n\n" for message in (notification, response)), ) - async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + async with httpx2.AsyncClient(transport=httpx2.MockTransport(respond)) as http_client: client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30, logging_callback=logging_callback) result: Final = await asyncio.wait_for( client._execute_session_operation( @@ -1420,24 +1470,24 @@ async def test_http_response_handler_preserves_notifications_and_tool_listing() async def test_invalid_tool_list_schema_is_identified_as_an_upstream_response() -> None: from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message - def respond(request: httpx.Request) -> httpx.Response: + def respond(request: httpx2.Request) -> httpx2.Response: if request.method == "DELETE": - return httpx.Response(200) + return httpx2.Response(200) payload: Final = json.loads(request.content) if "id" not in payload: - return httpx.Response(202) + return httpx2.Response(202) result: Final = ( { - "protocolVersion": LATEST_PROTOCOL_VERSION, + "protocolVersion": payload["params"]["protocolVersion"], "capabilities": {}, "serverInfo": {"name": "test", "version": "1"}, } if payload["method"] == "initialize" else {"tools": "secret-invalid-tools"} ) - return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload["id"], "result": result}) + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload["id"], "result": result}) - async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + async with httpx2.AsyncClient(transport=httpx2.MockTransport(respond)) as http_client: client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) with pytest.raises(ValidationError) as caught: await asyncio.wait_for( @@ -1453,7 +1503,7 @@ async def test_invalid_tool_list_schema_is_identified_as_an_upstream_response() assert "secret" not in message -class _DiagnosticSSEStream(httpx.AsyncByteStream): +class _DiagnosticSSEStream(httpx2.AsyncByteStream): def __init__(self, messages: asyncio.Queue[bytes | Exception | None]) -> None: self.messages = messages @@ -1510,26 +1560,26 @@ def _diagnostic_transport(transport: MCPTransport, mode: str, failure_method: st ) messages: Final[asyncio.Queue[bytes | Exception | None]] = asyncio.Queue() - async def respond(request: httpx.Request) -> httpx.Response: + async def respond(request: httpx2.Request) -> httpx2.Response: if request.method == "GET": - return httpx.Response( + return httpx2.Response( 200, headers={"Content-Type": "text/event-stream"}, stream=_DiagnosticSSEStream(messages) ) payload: Final = json.loads(request.content) if "method" not in payload or "id" not in payload: - return httpx.Response(202) + return httpx2.Response(202) if payload["method"] == failure_method and mode != "ok": if mode == "bad-json": await messages.put(b"secret-invalid-json") elif mode == "io-error": - await messages.put(httpx.ReadError("secret-read-error")) + await messages.put(httpx2.ReadError("secret-read-error")) elif mode == "closed": await messages.put(None) elif mode == "silent": await messages.put( b'{"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"Waiting"}}' ) - return httpx.Response(202) + return httpx2.Response(202) if payload["method"] == "tools/list": for message in ( { @@ -1543,7 +1593,7 @@ def _diagnostic_transport(transport: MCPTransport, mode: str, failure_method: st await messages.put(json.dumps(message).encode()) result: Final = ( { - "protocolVersion": LATEST_PROTOCOL_VERSION, + "protocolVersion": payload["params"]["protocolVersion"], "capabilities": {"tools": {}, "logging": {}}, "serverInfo": {"name": "diagnostic", "version": "1"}, } @@ -1553,14 +1603,14 @@ def _diagnostic_transport(transport: MCPTransport, mode: str, failure_method: st else {"content": [{"type": "text", "text": "pong"}], "isError": False} ) await messages.put(json.dumps({"jsonrpc": "2.0", "id": payload["id"], "result": result}).encode()) - return httpx.Response(202) + return httpx2.Response(202) def factory( headers: dict[str, str] | None = None, - timeout: httpx.Timeout | None = None, - auth: httpx.Auth | None = None, - ) -> httpx.AsyncClient: - return httpx.AsyncClient(transport=httpx.MockTransport(respond), headers=headers, timeout=timeout, auth=auth) + timeout: httpx2.Timeout | None = None, + auth: httpx2.Auth | None = None, + ) -> httpx2.AsyncClient: + return httpx2.AsyncClient(transport=httpx2.MockTransport(respond), headers=headers, timeout=timeout, auth=auth) return sse_client("https://example.com/sse", httpx_client_factory=factory) @@ -1582,7 +1632,7 @@ async def test_transport_parsing_failure_is_preserved(transport: MCPTransport, f @pytest.mark.asyncio async def test_sse_read_failure_is_preserved() -> None: client: Final = MCPClient(server_url="https://example.com/sse", transport_type=MCPTransport.sse, timeout=0.2) - with pytest.raises(httpx.ReadError, match="secret-read-error"): + with pytest.raises(httpx2.ReadError, match="secret-read-error"): await asyncio.wait_for( client._execute_session_operation( _diagnostic_transport(MCPTransport.sse, "io-error", "tools/list"), lambda session: session.list_tools() @@ -1611,11 +1661,11 @@ async def test_transport_completion_and_normal_messages(transport: MCPTransport, pending: Final = client._execute_session_operation(_diagnostic_transport(transport, mode, "tools/list"), operation) if mode == "ok": result: Final = await asyncio.wait_for(pending, timeout=3) - assert result.isError is False + assert result.is_error is False assert result.content[0].text == "pong" logging_callback.assert_awaited_once_with(LoggingMessageNotificationParams(level="info", data="Listing tools")) else: - with pytest.raises(McpError) as caught: + with pytest.raises(MCPError) as caught: await asyncio.wait_for(pending, timeout=3) if mode == "closed": assert "connection was closed" in _connection_error_message(caught.value, client.server_url, 0.2) @@ -1648,20 +1698,20 @@ async def test_transport_cancellation_cleans_up_a_pending_request(transport: MCP await asyncio.wait_for(task, timeout=3) -class _InterruptedHTTPBody(httpx.AsyncByteStream): +class _InterruptedHTTPBody(httpx2.AsyncByteStream): async def __aiter__(self) -> AsyncIterator[bytes]: yield b'{"jsonrpc":' - raise httpx.RemoteProtocolError("secret-incomplete-response") + raise httpx2.RemoteProtocolError("secret-incomplete-response") @pytest.mark.asyncio async def test_interrupted_http_response_preserves_the_transport_failure() -> None: - def respond(request: httpx.Request) -> httpx.Response: - return httpx.Response(200, headers={"Content-Type": "application/json"}, stream=_InterruptedHTTPBody()) + def respond(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, headers={"Content-Type": "application/json"}, stream=_InterruptedHTTPBody()) - async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + async with httpx2.AsyncClient(transport=httpx2.MockTransport(respond)) as http_client: client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) - with pytest.raises(httpx.RemoteProtocolError, match="secret-incomplete-response"): + with pytest.raises(httpx2.RemoteProtocolError, match="secret-incomplete-response"): await asyncio.wait_for( client._execute_session_operation( streamable_http_client(client.server_url, http_client=http_client), @@ -1673,12 +1723,12 @@ async def test_interrupted_http_response_preserves_the_transport_failure() -> No @pytest.mark.asyncio async def test_empty_http_event_stream_uses_the_existing_request_deadline() -> None: - def respond(request: httpx.Request) -> httpx.Response: - return httpx.Response(200, headers={"Content-Type": "text/event-stream"}, content=b"") + def respond(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, headers={"Content-Type": "text/event-stream"}, content=b"") - async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + async with httpx2.AsyncClient(transport=httpx2.MockTransport(respond)) as http_client: client: Final = MCPClient(server_url="https://example.com/mcp", timeout=0.2) - with pytest.raises(McpError) as caught: + with pytest.raises(MCPError) as caught: await asyncio.wait_for( client._execute_session_operation( streamable_http_client(client.server_url, http_client=http_client), @@ -1686,7 +1736,8 @@ async def test_empty_http_event_stream_uses_the_existing_request_deadline() -> N ), timeout=3, ) - assert isinstance(as_mcp_read_timeout(caught.value), TimeoutError) + assert caught.value.error.code == CONNECTION_CLOSED + assert "SSE stream ended" in caught.value.error.message @pytest.mark.asyncio @@ -1726,14 +1777,14 @@ async def test_optional_discovery_capabilities_and_errors( "resources/templates/list": {"name": "example", "uriTemplate": "test://{name}"}, }[method] - def respond(request: httpx.Request) -> httpx.Response: + def respond(request: httpx2.Request) -> httpx2.Response: if request.method == "DELETE": - return httpx.Response(200) - payload: Final = JSONRPCMessage.model_validate_json(request.content).root + return httpx2.Response(200) + payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) if not isinstance(payload, JSONRPCRequest): - return httpx.Response(202) + return httpx2.Response(202) if outcome == "initialize_not_found": - return httpx.Response( + return httpx2.Response( 200, json={ "jsonrpc": "2.0", @@ -1742,13 +1793,13 @@ async def test_optional_discovery_capabilities_and_errors( }, ) if payload.method == "initialize": - return httpx.Response( + return httpx2.Response( 200, json={ "jsonrpc": "2.0", "id": payload.id, "result": { - "protocolVersion": LATEST_PROTOCOL_VERSION, + "protocolVersion": payload.params["protocolVersion"], "capabilities": {} if outcome == "absent" else {advertised if outcome == "other_capability" else capability: {}}, @@ -1757,11 +1808,11 @@ async def test_optional_discovery_capabilities_and_errors( }, ) if outcome == "timeout": - raise httpx.ReadTimeout("Optional list timed out", request=request) + raise httpx2.ReadTimeout("Optional list timed out", request=request) if outcome == "unauthorized": - return httpx.Response(401) + return httpx2.Response(401) if outcome in ("method_not_found", "internal_error", "absent", "other_capability"): - return httpx.Response( + return httpx2.Response( 200, json={ "jsonrpc": "2.0", @@ -1772,26 +1823,24 @@ async def test_optional_discovery_capabilities_and_errors( }, }, ) - return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": {field: [entry]}}) + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": {field: [entry]}}) responder: Final = Mock(side_effect=respond) caplog.set_level(logging.DEBUG, logger="LiteLLM") - with respx.mock(base_url="https://example.com") as router: - router.route().mock(side_effect=responder) - client: Final = MCPClient(server_url="https://example.com/mcp") - operation: Final = { - "prompts/list": client.list_prompts, - "resources/list": client.list_resources, - "resources/templates/list": client.list_resource_templates, - }[method] - if raise_on_error and outcome in ("internal_error", "unauthorized", "timeout", "initialize_not_found"): - with pytest.raises((McpError, httpx.HTTPError)): - await operation(raise_on_error=True) - return - result: Final = await operation(raise_on_error=raise_on_error) + client: Final = _MockTransportClient(responder, server_url="https://example.com/mcp") + operation: Final = { + "prompts/list": client.list_prompts, + "resources/list": client.list_resources, + "resources/templates/list": client.list_resource_templates, + }[method] + if raise_on_error and outcome in ("internal_error", "unauthorized", "timeout", "initialize_not_found"): + with pytest.raises((MCPError, httpx2.HTTPError)): + await operation(raise_on_error=True) + return + result: Final = await operation(raise_on_error=raise_on_error) requests: Final = tuple( - JSONRPCMessage.model_validate_json(call.args[0].content).root + _JSONRPC_MESSAGE_ADAPTER.validate_json(call.args[0].content) for call in responder.call_args_list if call.args[0].method == "POST" ) @@ -1816,38 +1865,37 @@ async def test_optional_discovery_capabilities_and_errors( @pytest.mark.parametrize("supports_first", (True, False)) async def test_optional_discovery_uses_each_sessions_capabilities(supports_first: bool) -> None: from unittest.mock import Mock + from mcp.types import JSONRPCRequest capabilities: Final = iter(({"resources": {}}, {}) if supports_first else ({}, {"resources": {}})) - def respond(request: httpx.Request) -> httpx.Response: + def respond(request: httpx2.Request) -> httpx2.Response: if request.method == "DELETE": - return httpx.Response(200) - payload: Final = JSONRPCMessage.model_validate_json(request.content).root + return httpx2.Response(200) + payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) if not isinstance(payload, JSONRPCRequest): - return httpx.Response(202) + return httpx2.Response(202) result: Final = ( { - "protocolVersion": LATEST_PROTOCOL_VERSION, + "protocolVersion": payload.params["protocolVersion"], "capabilities": next(capabilities), "serverInfo": {"name": "changing", "version": "1"}, } if payload.method == "initialize" else {"resources": [{"name": "example", "uri": "test://example"}]} ) - return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) responder: Final = Mock(side_effect=respond) - with respx.mock(base_url="https://example.com") as router: - router.route().mock(side_effect=responder) - client: Final = MCPClient(server_url="https://example.com/mcp") - first: Final = await client.list_resources() - second: Final = await client.list_resources() + client: Final = _MockTransportClient(responder, server_url="https://example.com/mcp") + first: Final = await client.list_resources() + second: Final = await client.list_resources() assert [item.name for item in first] == (["example"] if supports_first else []) assert [item.name for item in second] == ([] if supports_first else ["example"]) requests: Final = tuple( - JSONRPCMessage.model_validate_json(call.args[0].content).root + _JSONRPC_MESSAGE_ADAPTER.validate_json(call.args[0].content) for call in responder.call_args_list if call.args[0].method == "POST" ) @@ -1862,20 +1910,20 @@ async def test_optional_discovery_preserves_cancellation(method: str) -> None: ready: Final = asyncio.Event() pending: Final = asyncio.Event() - async def respond(request: httpx.Request) -> httpx.Response: + async def respond(request: httpx2.Request) -> httpx2.Response: if request.method == "DELETE": - return httpx.Response(200) - payload: Final = JSONRPCMessage.model_validate_json(request.content).root + return httpx2.Response(200) + payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) if not isinstance(payload, JSONRPCRequest): - return httpx.Response(202) + return httpx2.Response(202) if payload.method == "initialize": - return httpx.Response( + return httpx2.Response( 200, json={ "jsonrpc": "2.0", "id": payload.id, "result": { - "protocolVersion": LATEST_PROTOCOL_VERSION, + "protocolVersion": payload.params["protocolVersion"], "capabilities": {"resources": {}, "prompts": {}}, "serverInfo": {"name": "pending", "version": "1"}, }, @@ -1883,23 +1931,21 @@ async def test_optional_discovery_preserves_cancellation(method: str) -> None: ) ready.set() await pending.wait() - return httpx.Response(202) + return httpx2.Response(202) - with respx.mock(base_url="https://example.com") as router: - router.route().mock(side_effect=respond) - client: Final = MCPClient(server_url="https://example.com/mcp") - operation: Final = { - "prompts/list": client.list_prompts, - "resources/list": client.list_resources, - "resources/templates/list": client.list_resource_templates, - }[method] - task: Final = asyncio.create_task(operation()) - try: - await asyncio.wait_for(ready.wait(), timeout=3) - finally: - task.cancel() - with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(task, timeout=3) + client: Final = _MockTransportClient(respond, server_url="https://example.com/mcp") + operation: Final = { + "prompts/list": client.list_prompts, + "resources/list": client.list_resources, + "resources/templates/list": client.list_resource_templates, + }[method] + task: Final = asyncio.create_task(operation()) + try: + await asyncio.wait_for(ready.wait(), timeout=3) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=3) @@ -1949,3 +1995,63 @@ async def test_request_auth_preview_uses_the_same_effective_headers_as_egress() assert str(request.url) == "https://upstream.example/mcp" assert request.headers["Authorization"] == "Bearer resolved" assert request.headers["X-Trace"] == "trace" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rpc_error", [False, True]) +async def test_expired_session_preserves_sdk_error_and_next_operation_reinitializes(rpc_error: bool) -> None: + from mcp.types import INVALID_REQUEST, METHOD_NOT_FOUND + + requests = [] + + def respond(request: httpx2.Request) -> httpx2.Response: + if request.method != "POST": + return httpx2.Response(405) + payload = json.loads(request.content) + if "id" not in payload: + return httpx2.Response(202) + requests.append((payload["method"], request.headers.get("mcp-session-id"))) + if payload["method"] == "initialize": + return httpx2.Response(200, headers={"mcp-session-id": f"session-{len(requests)}"}, json={ + "jsonrpc": "2.0", "id": payload["id"], "result": { + "protocolVersion": "2025-06-18", "capabilities": {}, + "serverInfo": {"name": "expiry-test", "version": "1"}, + }, + }) + if len(requests) == 2: + if rpc_error: + return httpx2.Response(404, json={ + "jsonrpc": "2.0", "id": payload["id"], + "error": {"code": METHOD_NOT_FOUND, "message": "Tool catalog unavailable"}, + }) + return httpx2.Response(404) + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload["id"], "result": {"tools": []}}) + + client = MCPClient(server_url="https://example.com/mcp", timeout=3) + async with client._create_httpx_client_factory(transport=httpx2.MockTransport(respond))() as http_client: + with pytest.raises(MCPError) as caught: + await client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), lambda session: session.list_tools() + ) + assert caught.value.error.code == (METHOD_NOT_FOUND if rpc_error else INVALID_REQUEST) + assert caught.value.error.message == ("Tool catalog unavailable" if rpc_error else "Session terminated") + result = await client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), lambda session: session.list_tools() + ) + assert result.tools == [] + assert requests == [("initialize", None), ("tools/list", "session-1"), ("initialize", None), ("tools/list", "session-3")] + + +@pytest.mark.asyncio +async def test_404_before_session_initialization_preserves_method_not_found() -> None: + from mcp.types import METHOD_NOT_FOUND + + client = MCPClient(server_url="https://example.com/mcp", timeout=3) + transport = httpx2.MockTransport(lambda request: httpx2.Response(404)) + async with client._create_httpx_client_factory(transport=transport)() as http_client: + with pytest.raises(MCPError) as caught: + await client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), lambda session: session.list_tools() + ) + assert caught.value.error.code == METHOD_NOT_FOUND + assert caught.value.error.message == "Not Found" diff --git a/tests/test_litellm/integrations/arize/test_arize_utils.py b/tests/test_litellm/integrations/arize/test_arize_utils.py index 50f2823d632..167b083e147 100644 --- a/tests/test_litellm/integrations/arize/test_arize_utils.py +++ b/tests/test_litellm/integrations/arize/test_arize_utils.py @@ -1235,7 +1235,7 @@ def test_arize_coerce_response_obj_dumps_pydantic_without_get(): coerced = _coerce_response_obj_for_attrs(result) assert isinstance(coerced, dict) - assert coerced["isError"] is False + assert coerced["is_error"] is False assert coerced["content"][0]["text"] == "hi" diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index c5c77a12a62..f4a8691f72f 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py @@ -1,6 +1,7 @@ """Tests for the OTel v2 sources of truth: span registry, semconv keys, config, and the typed StandardLoggingPayload adapter. These need no OTel SDK.""" +import json import logging import re from pathlib import Path @@ -12,11 +13,11 @@ import litellm from litellm.integrations.otel import ( BAGGAGE_PROMOTED_KEYS, DB, + HTTP, Error, GenAI, GenAIOperation, GenAIOutputType, - HTTP, LiteLLM, OpenTelemetryV2Config, Server, @@ -29,8 +30,8 @@ from litellm.integrations.otel import ( from litellm.integrations.otel.mappers.genai import GenAIMapper from litellm.integrations.otel.model import spans as spans_mod from litellm.integrations.otel.model.metadata import LLMCallEvent -from litellm.integrations.otel.model.trace_controls import TraceControls, caller_trace_controls from litellm.integrations.otel.model.payloads import ( + EmbeddingOutput, LLMCallSpanData, RequestIdentity, _upstream_address_port, @@ -43,6 +44,7 @@ from litellm.integrations.otel.model.spans import ( root_roles, validate_registry, ) +from litellm.integrations.otel.model.trace_controls import TraceControls, caller_trace_controls @pytest.fixture(autouse=True) @@ -696,6 +698,46 @@ def test_content_capture_gated_off_by_default(): assert data.finish_reasons == ("stop",) +def _embedding_payload(vectors: list[object], **overrides): + rows = [{"object": "embedding", "index": i, "embedding": vector} for i, vector in enumerate(vectors)] + return _sample_payload( + call_type="aembedding", + model="text-embedding-3-small", + response={"model": "text-embedding-3-small", "object": "list", "data": rows}, + **overrides, + ) + + +def test_embedding_response_is_summarized_as_vector_count_and_width(): + data = LLMCallSpanData.from_standard_logging_payload( + _embedding_payload([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]), capture_content=True + ) + + assert data.embedding_output == EmbeddingOutput(count=2, dimensions=3) + assert json.loads(data.embedding_output.as_json()) == {"count": 2, "dimensions": 3} + assert data.choices_out == () + + +def test_embedding_summary_follows_the_content_capture_gate(): + assert LLMCallSpanData.from_standard_logging_payload(_embedding_payload([[0.1]])).embedding_output is None + + +def test_embedding_summary_leaves_width_unknown_for_base64_vectors(): + data = LLMCallSpanData.from_standard_logging_payload(_embedding_payload(["AAAA"]), capture_content=True) + + assert data.embedding_output == EmbeddingOutput(count=1, dimensions=None) + + +def test_embedding_summary_is_absent_without_vectors_and_for_chat_data_lists(): + empty = LLMCallSpanData.from_standard_logging_payload(_embedding_payload([]), capture_content=True) + chat = LLMCallSpanData.from_standard_logging_payload( + _sample_payload(response={"data": [{"embedding": [0.1]}]}), capture_content=True + ) + + assert empty.embedding_output is None + assert chat.embedding_output is None + + def test_request_identity_prefers_canonical_team_keys(): from litellm.integrations.otel.model.payloads import RequestIdentity diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py b/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py index bd83357305e..c5ebc4bc53a 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py @@ -11,15 +11,14 @@ import pytest from litellm.integrations.otel import GenAIOperation from litellm.integrations.otel.mappers import ( - GenAIMapper, LangfuseMapper, LangtraceMapper, OpenInferenceMapper, WeaveMapper, resolve_mappers, ) -from litellm.integrations.otel.model.trace_controls import TraceControls from litellm.integrations.otel.model.payloads import ( + EmbeddingOutput, LLMCallSpanData, LLMRequestParams, LLMUsage, @@ -27,6 +26,7 @@ from litellm.integrations.otel.model.payloads import ( ServerInfo, ToolDefinition, ) +from litellm.integrations.otel.model.trace_controls import TraceControls def _llm_call(**overrides): @@ -174,6 +174,28 @@ def test_langfuse_mapper_skips_when_no_messages(): assert "langfuse.observation.output" not in attrs +def test_langfuse_mapper_renders_an_embedding_call_with_a_vector_summary_as_output(): + data = _llm_call( + operation=GenAIOperation.EMBEDDINGS, + request_model="text-embedding-3-small", + messages_in=({"role": "user", "content": "hello"},), + choices_out=(), + finish_reasons=(), + embedding_output=EmbeddingOutput(count=2, dimensions=1536), + ) + attrs = LangfuseMapper().map(data) + + assert attrs["langfuse.observation.type"] == "generation" + assert json.loads(attrs["langfuse.observation.output"]) == {"count": 2, "dimensions": 1536} + assert json.loads(attrs["langfuse.observation.input"]) == [{"role": "user", "content": "hello"}] + + +def test_langfuse_mapper_keeps_chat_output_when_no_embedding_summary(): + attrs = LangfuseMapper().map(_llm_call(embedding_output=None)) + + assert json.loads(attrs["langfuse.observation.output"]) == [{"role": "assistant", "content": "Sunny."}] + + # --------------------------------------------------------------------------- # # Weave # --------------------------------------------------------------------------- # diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index 92b1185e542..83649c3386a 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -1595,6 +1595,178 @@ class TestEnableAnthropicPromptCaching: assert supports_prompt_caching(model=model, custom_llm_provider=provider) is True assert self._points(model=model, provider=provider) == [] + @pytest.mark.parametrize("family", ["haiku-4-5", "sonnet-5", "opus-5", "fable-5", "fable-5-1"]) + @pytest.mark.parametrize( + "provider, template", + [("anthropic", "{}"), ("vertex_ai", "{}"), ("azure_ai", "{}"), ("bedrock", "us.anthropic.{}-v1:0")], + ) + @pytest.mark.parametrize("infer_provider", [False, True]) + @pytest.mark.parametrize("supported", [False, True]) + def test_claude_transport_defaults(self, monkeypatch, local_model_cost_map, family, provider, template, infer_provider, supported): + from litellm.utils import supports_prompt_caching + + model = template.format(f"claude-{family}") + qualified = f"{provider}/{model}" + entry = {"litellm_provider": provider, "mode": "chat", "supports_prompt_caching": supported} + monkeypatch.setitem(litellm.model_cost, model, entry) + monkeypatch.setitem(litellm.model_cost, qualified, entry) + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False) + target = qualified if infer_provider else model + resolved_provider = None if infer_provider else provider + assert supports_prompt_caching(model=target, custom_llm_provider=resolved_provider) is supported + points = AnthropicCacheControlHook.get_default_injection_points( + messages=copy.deepcopy(self.MESSAGES), system=None, model=target, + custom_llm_provider=resolved_provider, enable_prompt_caching=True, + ) + assert [point["index"] for point in points] == ([None, -1] if supported else []) + affinity_messages = AnthropicCacheControlHook.messages_with_default_injections( + copy.deepcopy(self.MESSAGES), models=[qualified], enable_prompt_caching=True, + ) + assert sum(AnthropicCacheControlHook._count_cache_control_blocks(m) for m in affinity_messages) == (2 if supported else 0) + + @pytest.mark.parametrize( + "provider, model", + [ + ("bedrock", "us.openai.gpt-6-astra"), + ("bedrock", "amazon.nova-pro-v1:0"), + ("bedrock", "us.xai.grok-4.6"), + ("bedrock", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/opaque"), + ("vertex_ai", "gemini-3.8-flash"), + ("azure_ai", "gpt-6-astra"), + ("anthropic", "unknown-model"), + ], + ) + def test_non_claude_caching_capability_does_not_enable_defaults(self, monkeypatch, local_model_cost_map, provider, model): + from litellm.utils import supports_prompt_caching + + qualified = f"{provider}/{model}" + entry = {"litellm_provider": provider, "mode": "chat", "supports_prompt_caching": True} + monkeypatch.setitem(litellm.model_cost, model, entry) + monkeypatch.setitem(litellm.model_cost, qualified, entry) + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + assert supports_prompt_caching(model=model, custom_llm_provider=provider) + assert self._points(model=model, provider=provider) == [] + assert self._points(model=qualified, provider=None) == [] + assert AnthropicCacheControlHook.messages_with_default_injections(self.MESSAGES, [qualified]) == self.MESSAGES + + @pytest.mark.parametrize("provider", ["vertex_ai", "azure_ai"]) + @pytest.mark.parametrize("client_control", ["none", "message", "system", "tool", "function", "top_level"]) + @pytest.mark.parametrize("envelope", ["request", "extra_body"]) + @pytest.mark.parametrize("configured", [False, True]) + def test_new_transports_preserve_client_controls(self, monkeypatch, local_model_cost_map, provider, client_control, envelope, configured): + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig + + model = "claude-sonnet-5" + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + monkeypatch.setitem(litellm.model_cost, f"{provider}/{model}", { + **litellm.model_cost[f"{provider}/{model}"], "supports_prompt_caching": True, + }) + control = {"type": "ephemeral"} + messages = [{"role": "user", "content": [{"type": "text", "text": "question", **({"cache_control": control} if client_control == "message" else {})}]}] + system = [{"type": "text", "text": "stable context", **({"cache_control": control} if client_control == "system" else {})}] + tools = [{"name": "lookup", "description": "Lookup", "input_schema": {"type": "object", "properties": {}}, **({"cache_control": control} if client_control == "tool" else {})}] + if client_control == "function": + tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}, "cache_control": control}}] + kwargs = {"metadata": {}, "model_info": {"id": "selected-deployment"}, **({"cache_control": control} if client_control == "top_level" else {})} + if envelope == "extra_body": + kwargs["extra_body"] = {"messages": messages, "system": system, "tools": tools} + if "cache_control" in kwargs: + kwargs["extra_body"]["cache_control"] = kwargs.pop("cache_control") + messages, system, tools = [{"role": "user", "content": "question"}], "stable context", [] + if configured: + kwargs["cache_control_injection_points"] = [ + {"location": "message", "role": "system", "index": None, "control": control}, + {"location": "message", "role": None, "index": -1, "control": control}, + ] + seeded = copy.deepcopy(kwargs) + original = copy.deepcopy((messages, system, tools)) + result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control( + messages, system, kwargs, model, provider, tools=tools, + ) + if client_control != "none": + assert (result_messages, result_system, tools) == original + assert kwargs["metadata"] == {} + else: + assert kwargs["metadata"]["litellm_gateway_injected_cache"] == "selected-deployment" + assert sum(AnthropicCacheControlHook._count_cache_control_blocks(m) for m in result_messages) == 1 + assert result_system[0]["cache_control"] == control + if provider == "vertex_ai": + wire = VertexAIAnthropicConfig().transform_request( + model=model, messages=[{"role": "system", "content": result_system}, *result_messages], + optional_params={"max_tokens": 8}, litellm_params={}, headers={}, + ) + assert wire["system"][0]["cache_control"] == control + assert wire["messages"][-1]["content"][-1]["cache_control"] == control + affinity = AnthropicCacheControlHook.messages_with_default_injections( + [{"role": "system", "content": original[1]}, *original[0]], [f"{provider}/{model}"], + tools=tools, request_kwargs=seeded, + ) + if client_control != "none": + assert affinity == [{"role": "system", "content": original[1]}, *original[0]] + AnthropicCacheControlHook.maybe_seed_default_injection_points( + seeded, [{"role": "system", "content": original[1]}, *original[0]], model, provider, tools=tools, + ) + assert bool(seeded.get("cache_control_injection_points")) == (client_control == "none") + + @pytest.mark.asyncio + @pytest.mark.parametrize("asynchronous", [False, True]) + @pytest.mark.parametrize("model, target, client_control, expected", [ + ("vertex_ai/claude-sonnet-5", "bedrock/amazon.nova-pro-v1:0", False, 0), + ("azure_ai/gpt-6-astra", "azure_ai/claude-sonnet-5", False, 2), + ("azure_ai/claude-sonnet-5", None, False, 2), + ("azure_ai/claude-sonnet-5", None, True, 1), + ("azure_ai/model_router/claude-replacement", None, False, 2), + ]) + async def test_public_completion_cache_ownership(self, monkeypatch, local_model_cost_map, asynchronous, model, target, client_control, expected): + import httpx + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + monkeypatch.setattr(litellm, "model_alias_map", {model: target} if target else {}) + for qualified in (model, target): + if qualified: + provider = qualified.split("/")[0] + entry = {"litellm_provider": provider, "mode": "chat", "supports_prompt_caching": True} + monkeypatch.setitem(litellm.model_cost, qualified, entry) + monkeypatch.setitem(litellm.model_cost, qualified.split("/", 1)[-1], entry) + sent = [] + def respond(request): + sent.append(json.loads(request.content)) + return httpx.Response(200, request=request, json={ + "id": "msg-test", "type": "message", "role": "assistant", "model": "claude-sonnet-5", + "content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn", "stop_sequence": None, + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, "stopReason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 1, "inputTokens": 10, "outputTokens": 1, "totalTokens": 11}, + }) + control = {"type": "ephemeral", "ttl": "1h"} + messages = [{"role": "system", "content": "stable context"}, {"role": "user", "content": "question"}] + metadata = {} + kwargs = { + "model": model, "messages": copy.deepcopy(messages), "max_tokens": 32, "num_retries": 0, + "litellm_metadata": metadata, + "api_base": "https://rig.services.ai.azure.com/anthropic", "api_key": "synthetic-test-key", + "aws_access_key_id": "synthetic", "aws_secret_access_key": "synthetic", "aws_region_name": "us-east-1", + **({"extra_body": {"cache_control": control}} if client_control else {}), + } + if asynchronous: + handler = AsyncHTTPHandler() + await handler.client.aclose() + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + handler.client = client + response = await litellm.acompletion(**kwargs, client=handler) + else: + with httpx.Client(transport=httpx.MockTransport(respond)) as client: + response = litellm.completion(**kwargs, client=HTTPHandler(client=client)) + assert response.choices[0].message.content == "ok" + assert len(sent) == 1 + assert ("litellm_gateway_injected_cache" in metadata) == (expected == 2) + serialized = json.dumps(sent[0]) + assert serialized.count('"cache_control"') + serialized.count('"cachePoint"') == expected + if client_control: + assert sent[0]["cache_control"] == control + affinity = AnthropicCacheControlHook.messages_with_default_injections(messages, [model], request_kwargs=kwargs) + assert AnthropicCacheControlHook.count_request_cache_breakpoints(affinity) == (2 if expected == 2 else 0) + def test_databricks_claude_not_injected_despite_caching_support(self, monkeypatch, local_model_cost_map): from litellm.utils import supports_prompt_caching diff --git a/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py b/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py index a0f22a59f0c..3f1fe0d5d68 100644 --- a/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py +++ b/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py @@ -420,7 +420,7 @@ class TestHandleSkillSearchMCP: result = await handle_skill_search( query="language translation", top_k=10_000, user_api_key_dict=UserAPIKeyAuth(user_id="u") ) - assert result.isError is False + assert result.is_error is False assert len(json.loads(result.content[0].text)) == MAX_SKILL_SEARCH_TOP_K @pytest.mark.asyncio @@ -432,5 +432,5 @@ class TestHandleSkillSearchMCP: result = await handle_skill_search( query="language translation", top_k=0, user_api_key_dict=UserAPIKeyAuth(user_id="u") ) - assert result.isError is False + assert result.is_error is False assert len(json.loads(result.content[0].text)) == 1 diff --git a/tests/test_litellm/ocr/test_dispatch.py b/tests/test_litellm/ocr/test_dispatch.py index 14d3368f869..e54d4070ba8 100644 --- a/tests/test_litellm/ocr/test_dispatch.py +++ b/tests/test_litellm/ocr/test_dispatch.py @@ -387,3 +387,39 @@ async def test_public_aocr_routes_through_dispatch(monkeypatch: pytest.MonkeyPat NATIVE_AOCR.reset() assert result is expected assert [request.model for request in captured] == ["mistral/mistral-ocr-latest"] + + +@pytest.mark.parametrize( + ("model", "custom_llm_provider", "expected"), + ( + ("aws_textract/detect-document-text", None, "native"), + ("detect-document-text", "aws_textract", "native"), + ("mistral/mistral-ocr-latest", None, "python"), + ("mistral/mistral-ocr-latest", "aws_textract", "native"), + ("aws_textract", None, "python"), + ), +) +def test_provider_scoped_rule_sees_the_provider_named_by_the_model_prefix( + model: str, custom_llm_provider: str | None, expected: str +) -> None: + rules: Final[Rules] = ( + Rule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"aws_textract"})), + Rule(Route.OCR, Rollout.PYTHON_ONLY), + ) + document: Final[Mapping[str, object]] = {"type": "image_url", "image_url": "data:image/png;base64,YQ=="} + kwargs: Final[Mapping[str, object]] = ( + {} if custom_llm_provider is None else {"custom_llm_provider": custom_llm_provider} + ) + python_response: Final = response("python") + native_response: Final = response("native") + + result: Final = _DISPATCH.run( + (model, document), + kwargs, + python=lambda *_args, **_kwargs: python_response, + binding=ocr_binding(lambda *_args, **_kwargs: native_response), + native=lambda _hook, _request, _args, _kwargs: native_response, + rules=rules, + ) + + assert cast(OCRResponse, result).model == expected # noqa: TID251 # sync dispatch returns the response itself diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py index 9a66f130d24..76e92efd31a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py @@ -44,3 +44,37 @@ def _hermetic_server_root_path(): finally: if saved is not None: os.environ["SERVER_ROOT_PATH"] = saved + + +@pytest.fixture +def config_only_mcp_manager_factory(): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + + class ConfigOnlyManager(MCPServerManager): + def initialize_tool_name_to_mcp_server_name_mapping(self): + return None + + return ConfigOnlyManager + + +@pytest.fixture +def _mcp_request_ctx(): + def _mcp_request_ctx(**overrides): + from types import SimpleNamespace + + from mcp.server.context import ServerRequestContext + + kwargs = { + "session": SimpleNamespace(), + "lifespan_context": {}, + "protocol_version": "2025-06-18", + "method": "", + "params": None, + "request_id": 1, + "meta": None, + "request": None, + } + kwargs.update(overrides) + return ServerRequestContext(**kwargs) + + return _mcp_request_ctx diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py index 65e2faee1b2..f951499e18f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py @@ -9,7 +9,7 @@ if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 import httpx import pytest -from mcp import McpError +from mcp import MCPError from mcp.types import ErrorData from litellm.proxy._experimental.mcp_server.exceptions import ( @@ -45,7 +45,7 @@ def test_upstream_json_rpc_error_code_is_never_read_as_an_http_status(): to answer with application code 408. Classifying that number as a gateway timeout would report a 504 the gateway never caused. A client timeout reaches here already expressed as a ``TimeoutError``, so this taxonomy never has to read the code to tell them apart.""" - upstream_error = McpError(ErrorData(code=int(httpx.codes.REQUEST_TIMEOUT), message="re-authenticate and retry")) + upstream_error = MCPError(code=int(httpx.codes.REQUEST_TIMEOUT), message="re-authenticate and retry") assert classify_list_exception(upstream_error).tag != "timeout" assert list_fault_http_status(classify_list_exception(upstream_error)) != 504 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py b/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py index 28959054195..77e9b987e74 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py @@ -652,7 +652,7 @@ async def test_structured_content_is_masked_alongside_content(): returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail) assert returned.content[0].text == "email " - assert returned.structuredContent == {"contact": {"email": ""}, "balance": 42.0} + assert returned.structured_content == {"contact": {"email": ""}, "balance": 42.0} @pytest.mark.asyncio @@ -673,7 +673,7 @@ async def test_value_present_only_in_structured_content_is_masked(): returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail) assert "jane@example.com" in guardrail.seen_texts - assert returned.structuredContent == {"records": [{"email": ""}]} + assert returned.structured_content == {"records": [{"email": ""}]} assert returned.content[0].text == "lookup complete" @@ -690,7 +690,7 @@ async def test_structured_content_without_a_match_is_untouched(): returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail) - assert returned.structuredContent == {"record_id": "C-1001", "balance": 42.0, "active": True, "note": None} + assert returned.structured_content == {"record_id": "C-1001", "balance": 42.0, "active": True, "note": None} @pytest.mark.asyncio @@ -798,4 +798,4 @@ async def test_clean_structured_content_keys_do_not_block(): returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail) assert returned.content[0].text == "email " - assert returned.structuredContent == {"record_id": "C-1001", "balance": 42.0, "count": 3} + assert returned.structured_content == {"record_id": "C-1001", "balance": 42.0, "count": 3} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py index 774cd022703..1cad9a1fccb 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py @@ -6,6 +6,7 @@ rotation-aware cache keying, expires_in-driven expiry, error classification, and """ import httpx +import httpx2 import pytest from pydantic import SecretStr @@ -322,27 +323,27 @@ async def test_refetch_returns_none_when_the_grant_fails(): assert await source.refetch("s", _config(), failed_access_token="stale") is None -def _upstream(responses: "list[httpx.Response]") -> "tuple[httpx.MockTransport, list[str]]": +def _upstream(responses: "list[httpx2.Response]") -> "tuple[httpx2.MockTransport, list[str]]": # The auth flow re-yields the same Request object on retry, so snapshot the Authorization # value per send; holding the Request would show the post-retry mutation for both entries. seen: "list[str]" = [] - def handler(request: httpx.Request) -> httpx.Response: + def handler(request: httpx2.Request) -> httpx2.Response: seen.append(request.headers.get("Authorization", "")) return responses[min(len(seen) - 1, len(responses) - 1)] - return httpx.MockTransport(handler), seen + return httpx2.MockTransport(handler), seen @pytest.mark.asyncio async def test_bearer_auth_sends_the_token_and_leaves_a_success_alone(): - transport, seen = _upstream([httpx.Response(200)]) + transport, seen = _upstream([httpx2.Response(200)]) async def refetch(failed: str) -> "str | None": raise AssertionError("must not refetch on success") auth = ClientCredentialsBearerAuth("m2m-token", refetch, ClientCredentialsConfig()) - async with httpx.AsyncClient(transport=transport, auth=auth) as client: + async with httpx2.AsyncClient(transport=transport, auth=auth) as client: response = await client.get("https://upstream.example.com/mcp") assert response.status_code == 200 assert seen == ["Bearer m2m-token"] @@ -350,7 +351,7 @@ async def test_bearer_auth_sends_the_token_and_leaves_a_success_alone(): @pytest.mark.asyncio async def test_bearer_auth_retries_a_401_once_with_a_fresh_token(): - transport, seen = _upstream([httpx.Response(401), httpx.Response(200)]) + transport, seen = _upstream([httpx2.Response(401), httpx2.Response(200)]) refetched: "list[str]" = [] async def refetch(failed: str) -> "str | None": @@ -358,7 +359,7 @@ async def test_bearer_auth_retries_a_401_once_with_a_fresh_token(): return "fresh-token" auth = ClientCredentialsBearerAuth("stale-token", refetch, ClientCredentialsConfig()) - async with httpx.AsyncClient(transport=transport, auth=auth) as client: + async with httpx2.AsyncClient(transport=transport, auth=auth) as client: response = await client.get("https://upstream.example.com/mcp") assert response.status_code == 200 assert refetched == ["stale-token"] @@ -370,7 +371,7 @@ async def test_bearer_auth_remembers_the_rotated_token_for_later_requests(): # The auth object lives for the whole MCP session (it is the httpx client's auth), so after a # 401 recovery it must send the fresh token first on subsequent requests; re-sending the # rejected one would burn a 401 round trip and the single retry on every call. - transport, seen = _upstream([httpx.Response(401), httpx.Response(200), httpx.Response(200)]) + transport, seen = _upstream([httpx2.Response(401), httpx2.Response(200), httpx2.Response(200)]) refetched: "list[str]" = [] async def refetch(failed: str) -> "str | None": @@ -378,7 +379,7 @@ async def test_bearer_auth_remembers_the_rotated_token_for_later_requests(): return "fresh-token" auth = ClientCredentialsBearerAuth("stale-token", refetch, ClientCredentialsConfig()) - async with httpx.AsyncClient(transport=transport, auth=auth) as client: + async with httpx2.AsyncClient(transport=transport, auth=auth) as client: first = await client.get("https://upstream.example.com/mcp") second = await client.get("https://upstream.example.com/mcp") assert first.status_code == 200 and second.status_code == 200 @@ -388,13 +389,13 @@ async def test_bearer_auth_remembers_the_rotated_token_for_later_requests(): @pytest.mark.asyncio async def test_bearer_auth_surfaces_the_401_when_the_refetch_fails(): - transport, seen = _upstream([httpx.Response(401)]) + transport, seen = _upstream([httpx2.Response(401)]) async def refetch(failed: str) -> "str | None": return None auth = ClientCredentialsBearerAuth("stale-token", refetch, ClientCredentialsConfig()) - async with httpx.AsyncClient(transport=transport, auth=auth) as client: + async with httpx2.AsyncClient(transport=transport, auth=auth) as client: response = await client.get("https://upstream.example.com/mcp") assert response.status_code == 401 assert len(seen) == 1 @@ -402,7 +403,7 @@ async def test_bearer_auth_surfaces_the_401_when_the_refetch_fails(): @pytest.mark.asyncio async def test_bearer_auth_gives_up_after_a_second_401(): - transport, seen = _upstream([httpx.Response(401), httpx.Response(401)]) + transport, seen = _upstream([httpx2.Response(401), httpx2.Response(401)]) refetched: "list[str]" = [] async def refetch(failed: str) -> "str | None": @@ -410,7 +411,7 @@ async def test_bearer_auth_gives_up_after_a_second_401(): return "fresh-token" auth = ClientCredentialsBearerAuth("stale-token", refetch, ClientCredentialsConfig()) - async with httpx.AsyncClient(transport=transport, auth=auth) as client: + async with httpx2.AsyncClient(transport=transport, auth=auth) as client: response = await client.get("https://upstream.example.com/mcp") assert response.status_code == 401 assert len(seen) == 2 @@ -422,7 +423,7 @@ def test_bearer_auth_rejects_sync_clients(): return None auth = ClientCredentialsBearerAuth("token", refetch, ClientCredentialsConfig()) - with httpx.Client(transport=httpx.MockTransport(lambda request: httpx.Response(200)), auth=auth) as client: + with httpx2.Client(transport=httpx2.MockTransport(lambda request: httpx2.Response(200)), auth=auth) as client: with pytest.raises(RuntimeError): client.get("https://upstream.example.com/mcp") @@ -431,15 +432,15 @@ def test_bearer_auth_rejects_sync_clients(): async def test_bearer_auth_writes_the_minted_token_to_the_configured_header(): seen: "list[dict[str, str]]" = [] - def handler(request: httpx.Request) -> httpx.Response: + def handler(request: httpx2.Request) -> httpx2.Response: seen.append(dict(request.headers)) - return httpx.Response(200) + return httpx2.Response(200) async def refetch(failed: str) -> "str | None": raise AssertionError("must not refetch on success") auth = ClientCredentialsBearerAuth("m2m-token", refetch, ClientCredentialsConfig(header_name="esb-oauth")) - async with httpx.AsyncClient(transport=httpx.MockTransport(handler), auth=auth) as client: + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handler), auth=auth) as client: await client.get("https://upstream.example.com/mcp") assert seen[0]["esb-oauth"] == "Bearer m2m-token" assert "authorization" not in seen[0] @@ -451,9 +452,9 @@ async def test_the_401_refetch_retry_also_targets_the_configured_header(): # would silently send the fresh token to Authorization, so the ESB rejects every recovered # request while the first attempt looked correct. seen: "list[dict[str, str]]" = [] - responses = [httpx.Response(401), httpx.Response(200)] + responses = [httpx2.Response(401), httpx2.Response(200)] - def handler(request: httpx.Request) -> httpx.Response: + def handler(request: httpx2.Request) -> httpx2.Response: seen.append(dict(request.headers)) return responses[min(len(seen) - 1, len(responses) - 1)] @@ -461,7 +462,7 @@ async def test_the_401_refetch_retry_also_targets_the_configured_header(): return "fresh-token" auth = ClientCredentialsBearerAuth("stale-token", refetch, ClientCredentialsConfig(header_name="esb-oauth")) - async with httpx.AsyncClient(transport=httpx.MockTransport(handler), auth=auth) as client: + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handler), auth=auth) as client: response = await client.get("https://upstream.example.com/mcp") assert response.status_code == 200 assert [h["esb-oauth"] for h in seen] == ["Bearer stale-token", "Bearer fresh-token"] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_httpx_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_httpx_auth.py index 9eab089bac6..5a5eea60fce 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_httpx_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_httpx_auth.py @@ -1,10 +1,10 @@ -"""Tests for the concrete httpx.Auth objects the resolver returns. +"""Tests for the concrete httpx2.Auth objects the resolver returns. NoOpAuth must attach nothing; StaticHeaderAuth must set exactly the configured header. These pin the header emission the api_key family and passthrough depend on. """ -import httpx +import httpx2 from litellm.proxy._experimental.mcp_server.outbound_credentials import ( NoOpAuth, @@ -12,7 +12,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import ( ) -def _apply(auth: httpx.Auth, request: httpx.Request) -> httpx.Request: +def _apply(auth: httpx2.Auth, request: httpx2.Request) -> httpx2.Request: flow = auth.auth_flow(request) sent = next(flow) flow.close() @@ -20,19 +20,19 @@ def _apply(auth: httpx.Auth, request: httpx.Request) -> httpx.Request: def test_noop_auth_attaches_no_authorization_header(): - request = httpx.Request("GET", "https://upstream.example.com/mcp") + request = httpx2.Request("GET", "https://upstream.example.com/mcp") _apply(NoOpAuth(), request) assert "authorization" not in request.headers def test_static_header_auth_defaults_to_authorization(): - request = httpx.Request("GET", "https://upstream.example.com/mcp") + request = httpx2.Request("GET", "https://upstream.example.com/mcp") _apply(StaticHeaderAuth("Bearer abc"), request) assert request.headers["Authorization"] == "Bearer abc" def test_static_header_auth_honors_custom_header_name(): - request = httpx.Request("GET", "https://upstream.example.com/mcp") + request = httpx2.Request("GET", "https://upstream.example.com/mcp") _apply(StaticHeaderAuth("raw-key", header_name="X-API-Key"), request) assert request.headers["X-API-Key"] == "raw-key" assert "authorization" not in request.headers diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 5fab4ceec72..0e47bbb9bb1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -12,7 +12,7 @@ import logging import time from datetime import datetime, timedelta, timezone -import httpx +import httpx2 import jwt as pyjwt import pytest from pydantic import SecretStr @@ -109,8 +109,8 @@ def _spec(config): return ServerSpec(server_id="s", resource="https://upstream.example.com", config=config) -def _emitted(auth: httpx.Auth) -> httpx.Headers: - request = httpx.Request("GET", "https://upstream.example.com/mcp") +def _emitted(auth: httpx2.Auth) -> httpx2.Headers: + request = httpx2.Request("GET", "https://upstream.example.com/mcp") flow = auth.auth_flow(request) next(flow) flow.close() @@ -412,15 +412,15 @@ _M2M = ClientCredentialsConfig( ) -async def _emitted_async(auth: httpx.Auth, respond=None) -> tuple[httpx.Headers, list[httpx.Request]]: +async def _emitted_async(auth: httpx2.Auth, respond=None) -> tuple[httpx2.Headers, list[httpx2.Request]]: """Drive the async auth flow one request at a time, replying via ``respond`` when given.""" - seen: list[httpx.Request] = [] + seen: list[httpx2.Request] = [] - def handler(request: httpx.Request) -> httpx.Response: + def handler(request: httpx2.Request) -> httpx2.Response: seen.append(request) - return respond(request) if respond else httpx.Response(200) + return respond(request) if respond else httpx2.Response(200) - async with httpx.AsyncClient(transport=httpx.MockTransport(handler), auth=auth) as client: + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handler), auth=auth) as client: await client.get("https://upstream.example.com/mcp") return seen[-1].headers, seen @@ -458,9 +458,9 @@ async def test_client_credentials_auth_retries_a_401_through_the_source(): ) assert isinstance(result, Ok) - def respond(request: httpx.Request) -> httpx.Response: + def respond(request: httpx2.Request) -> httpx2.Response: is_stale = request.headers["Authorization"] == "Bearer stale-at" - return httpx.Response(401) if is_stale else httpx.Response(200) + return httpx2.Response(401) if is_stale else httpx2.Response(200) headers, seen = await _emitted_async(result.ok, respond) assert headers["Authorization"] == "Bearer fresh-m2m" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py index 333d4c98899..e3437bf16f6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py @@ -18,9 +18,9 @@ from litellm.proxy._types import LiteLLM_MCPServerTable class TestMCPCustomFields: """Test custom fields functionality in MCP server configuration.""" - async def test_custom_fields_preserved_from_config(self): + async def test_custom_fields_preserved_from_config(self, config_only_mcp_manager_factory): """Test that custom fields in mcp_info are preserved when loading from config.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() # Mock config with custom fields mock_config = { @@ -62,9 +62,9 @@ class TestMCPCustomFields: assert mcp_info["priority"] == 10 assert mcp_info["tags"] == ["production", "api"] - async def test_custom_fields_preserved_from_database(self): + async def test_custom_fields_preserved_from_database(self, config_only_mcp_manager_factory): """Test that custom fields in mcp_info are preserved when adding from database.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() # Mock database record with custom fields mock_server = LiteLLM_MCPServerTable( @@ -106,9 +106,9 @@ class TestMCPCustomFields: assert mcp_info["metadata"] == {"source": "database"} assert mcp_info["version"] == "1.0.0" - async def test_empty_mcp_info_handled_gracefully(self): + async def test_empty_mcp_info_handled_gracefully(self, config_only_mcp_manager_factory): """Test that empty or missing mcp_info is handled gracefully.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() # Config with empty mcp_info mock_config = { @@ -130,9 +130,9 @@ class TestMCPCustomFields: # Should have default server_name assert mcp_info["server_name"] == "test_server" - async def test_missing_mcp_info_creates_defaults(self): + async def test_missing_mcp_info_creates_defaults(self, config_only_mcp_manager_factory): """Test that missing mcp_info creates appropriate defaults.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() # Config without mcp_info mock_config = { @@ -155,9 +155,9 @@ class TestMCPCustomFields: assert mcp_info["server_name"] == "test_server" assert mcp_info["description"] == "Server description" - async def test_config_description_fallback(self): + async def test_config_description_fallback(self, config_only_mcp_manager_factory): """Test that description from config level is used as fallback.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() # Config with description at server level but not in mcp_info mock_config = { @@ -179,9 +179,9 @@ class TestMCPCustomFields: assert mcp_info["description"] == "Config level description" assert mcp_info["custom_field"] == "custom_value" - async def test_mcp_info_description_takes_precedence(self): + async def test_mcp_info_description_takes_precedence(self, config_only_mcp_manager_factory): """Test that description in mcp_info takes precedence over config level.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() # Config with description at both levels mock_config = { diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py index b6535e6326a..46ecd4df716 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py @@ -5,20 +5,17 @@ Tests for MCPDebug — MCP OAuth2 debug response headers. import asyncio from typing import Final +import httpx import pytest from starlette.types import Message -from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution - -import httpx - from litellm.proxy._experimental.mcp_server.mcp_debug import ( MCP_DEBUG_REQUEST_HEADER, + MCPAuthDiagnostics, MCPDebug, describe_upstream_http_failure, - - MCPAuthDiagnostics, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution class TestIsDebugEnabled: @@ -265,6 +262,7 @@ class TestDescribeUpstreamHttpFailure: assert describe_upstream_http_failure(ConnectionError("refused")) is None + @pytest.mark.parametrize("body", [ b'{"password":"first second","token":"demo-secret"}', b'{"nested":[{"access_token":"first,second"}]}', @@ -464,13 +462,12 @@ def test_diagnostics_keep_requests_separate_and_do_not_collapse_multiple_servers @pytest.mark.asyncio -async def test_concurrent_mcp_messages_record_on_their_own_http_scope() -> None: +async def test_concurrent_mcp_messages_record_on_their_own_http_scope(_mcp_request_ctx) -> None: from unittest.mock import MagicMock - from mcp.server.lowlevel.server import request_ctx - from mcp.shared.context import RequestContext from starlette.requests import Request + from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var from litellm.proxy._experimental.mcp_server.mcp_debug import ( MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, record_auth_resolution, @@ -481,16 +478,16 @@ async def test_concurrent_mcp_messages_record_on_their_own_http_scope() -> None: second: Final = MCPAuthDiagnostics() async def record(diagnostics: MCPAuthDiagnostics, source: AuthResolution) -> None: - context: Final = RequestContext( - request_id=1, meta=None, session=session, lifespan_context=None, + context: Final = _mcp_request_ctx( + session=session, request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), ) - token: Final = request_ctx.set(context) + token: Final = active_mcp_request_ctx_var.set(context) try: await asyncio.sleep(0) record_auth_resolution("same-server", source) finally: - request_ctx.reset(token) + active_mcp_request_ctx_var.reset(token) await asyncio.gather(record(first, AuthResolution.stored_user_token), record(second, AuthResolution.per_request_header)) assert first.resolution() == "stored-user-token" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py index b93f0d56f8e..a59b02ec01d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py @@ -30,7 +30,7 @@ def _form_params(message: str = "fill the form") -> ElicitRequestFormParams: return ElicitRequestFormParams( mode="form", message=message, - requestedSchema={"type": "object", "properties": {}}, + requested_schema={"type": "object", "properties": {}}, ) @@ -39,7 +39,7 @@ def _url_params(message: str = "please authorize") -> ElicitRequestURLParams: mode="url", message=message, url="https://example.com/oauth", - elicitationId="elc-1", + elicitation_id="elc-1", ) @@ -118,7 +118,7 @@ class TestRelayElicitationToDownstream: session.elicit_form.assert_awaited_once() _, kwargs = session.elicit_form.call_args assert kwargs["message"] == "collect name" - assert kwargs["requestedSchema"] == params.requestedSchema + assert kwargs["requested_schema"] == params.requested_schema async def test_should_relay_url_mode(self): accepted = ElicitResult(action="accept") @@ -142,7 +142,7 @@ class TestRelayElicitationToDownstream: # A bare params object that is neither Form nor URL params triggers # the generic fallback path. - params = SimpleNamespace(mode="form", message="hi", requestedSchema={}) + params = SimpleNamespace(mode="form", message="hi", requested_schema={}) result = await _relay_elicitation_to_downstream( params=params, downstream_session=session, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py index 36b545ad031..93b894f7645 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py @@ -1698,7 +1698,7 @@ def test_decrypt_global_env_var_drops_undecryptable_value( @pytest.mark.asyncio async def test_missing_user_env_vars_error_renders_in_mcp_call_tool(): """The MCP ``call_tool`` handler must turn ``MCPMissingUserEnvVarsError`` - into a friendly ``CallToolResult`` with ``isError=True`` so Claude Code + into a friendly ``CallToolResult`` with ``is_error=True`` so Claude Code surfaces the setup URL instead of an opaque internal error.""" from mcp.types import TextContent @@ -1716,7 +1716,7 @@ async def test_missing_user_env_vars_error_renders_in_mcp_call_tool(): content=[TextContent(text=str(err), type="text")], isError=True, ) - assert result.isError is True + assert result.is_error is True text = result.content[0].text # type: ignore[union-attr] assert "CorporateDB" in text assert "CORP_USERNAME" in text diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_metadata_preservation.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_metadata_preservation.py index 5a24ca00c25..86748d99063 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_metadata_preservation.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_metadata_preservation.py @@ -39,15 +39,12 @@ class TestMCPMetadataPreservation: name="hello_widget", description="Display a greeting widget", inputSchema={"type": "object", "properties": {}}, + meta={ + "openai/outputTemplate": "ui://widget/hello.html", + "openai/widgetDescription": "A greeting widget", + "openai/toolInvocation/invoking": "Preparing greeting...", + }, ) - # Add metadata using setattr since MCPTool might not have it in the constructor - tool_with_metadata.metadata = { - "openai/outputTemplate": "ui://widget/hello.html", - "openai/widgetDescription": "A greeting widget", - } - tool_with_metadata._meta = { - "openai/toolInvocation/invoking": "Preparing greeting...", - } # Create prefixed tools prefixed_tools = manager._create_prefixed_tools( @@ -61,22 +58,16 @@ class TestMCPMetadataPreservation: # Check that name is prefixed assert prefixed_tool.name == "test-hello_widget" - # Check that metadata is preserved - assert hasattr(prefixed_tool, "metadata") - assert prefixed_tool.metadata == { + # Check that _meta (the SDK `meta` field) is preserved + assert prefixed_tool.meta == { "openai/outputTemplate": "ui://widget/hello.html", "openai/widgetDescription": "A greeting widget", - } - - # Check that _meta is preserved - assert hasattr(prefixed_tool, "_meta") - assert prefixed_tool._meta == { "openai/toolInvocation/invoking": "Preparing greeting...", } # Check that other fields are preserved assert prefixed_tool.description == "Display a greeting widget" - assert prefixed_tool.inputSchema == {"type": "object", "properties": {}} + assert prefixed_tool.input_schema== {"type": "object", "properties": {}} if __name__ == "__main__": diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py index 67b7c5a3414..84d4f1fd083 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py @@ -3,7 +3,7 @@ from datetime import datetime import pytest from fastapi import HTTPException -from mcp.shared.exceptions import McpError +from mcp.shared.exceptions import MCPError from pydantic import AnyUrl import litellm @@ -32,7 +32,7 @@ async def test_proxy_call_rejects_non_proxy_tool_names() -> None: ) assert result is not None - assert result.isError is True + assert result.is_error is True assert "unavailable on /mcp/proxy" in result.content[0].text @@ -44,16 +44,28 @@ async def test_proxy_rejects_non_tool_protocol_operations() -> None: assert options.capabilities.resources is None assert options.capabilities.tools is not None - with pytest.raises(McpError): - await server.list_prompts() - with pytest.raises(McpError): - await server.get_prompt("prompt", {}) - with pytest.raises(McpError): - await server.list_resources() - with pytest.raises(McpError): - await server.list_resource_templates() - with pytest.raises(McpError): - await server.read_resource(AnyUrl("https://example.com/resource")) + from types import SimpleNamespace + + from mcp.server.context import ServerRequestContext + from mcp.types import GetPromptRequestParams, PaginatedRequestParams, ReadResourceRequestParams + + ctx = ServerRequestContext( + session=SimpleNamespace(), + lifespan_context={}, + protocol_version="2025-06-18", + method="", + ) + + with pytest.raises(MCPError): + await server.list_prompts(ctx, PaginatedRequestParams()) + with pytest.raises(MCPError): + await server.get_prompt(ctx, GetPromptRequestParams(name="prompt", arguments={})) + with pytest.raises(MCPError): + await server.list_resources(ctx, PaginatedRequestParams()) + with pytest.raises(MCPError): + await server.list_resource_templates(ctx, PaginatedRequestParams()) + with pytest.raises(MCPError): + await server.read_resource(ctx, ReadResourceRequestParams(uri="https://example.com/resource")) class FailureRecorder(CustomLogger): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py index 78aee7b534f..d17b407a1be 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py @@ -28,14 +28,14 @@ def _params(**overrides): role="user", content=SimpleNamespace(type="text", text="hi") ) ], - systemPrompt="be concise", - maxTokens=128, + system_prompt="be concise", + max_tokens=128, temperature=None, - stopSequences=None, + stop_sequences=None, tools=None, - toolChoice=None, + tool_choice=None, metadata=None, - modelPreferences=None, + model_preferences=None, ) base.update(overrides) return SimpleNamespace(**base) @@ -52,13 +52,13 @@ class TestBuildCompletionKwargs: async def test_should_include_sampling_options_and_tools(self): params = _params( temperature=0.3, - stopSequences=["STOP"], + stop_sequences=["STOP"], tools=[ SimpleNamespace( - name="search", description="d", inputSchema={"type": "object"} + name="search", description="d", input_schema={"type": "object"} ) ], - toolChoice=SimpleNamespace(mode="required"), + tool_choice=SimpleNamespace(mode="required"), metadata={"trace": "abc"}, ) with patch( @@ -179,7 +179,7 @@ class TestHandleSamplingCreateMessagePipeline: assert isinstance(result, CreateMessageResult) assert result.content.text == "the answer is 42" - assert result.stopReason == "endTurn" + assert result.stop_reason== "endTurn" async def test_should_reraise_known_proxy_exceptions(self): from litellm.exceptions import RateLimitError diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_access.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_access.py index 7c5320ed4f4..8975f42387b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_access.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_model_access.py @@ -212,14 +212,14 @@ class TestSamplingAuthAndBudgetGating: ) params = MagicMock() - params.modelPreferences = None + params.model_preferences = None params.messages = [] params.systemPrompt = None - params.maxTokens = 100 + params.max_tokens = 100 params.temperature = None - params.stopSequences = None + params.stop_sequences = None params.tools = None - params.toolChoice = None + params.tool_choice = None params.metadata = None result = await handle_sampling_create_message( @@ -242,14 +242,14 @@ class TestSamplingAuthAndBudgetGating: auth = _make_user_api_key_auth(models=["gpt-4o"]) params = MagicMock() - params.modelPreferences = None + params.model_preferences = None params.messages = [] params.systemPrompt = None - params.maxTokens = 100 + params.max_tokens = 100 params.temperature = None - params.stopSequences = None + params.stop_sequences = None params.tools = None - params.toolChoice = None + params.tool_choice = None params.metadata = None with ( @@ -304,14 +304,14 @@ class TestSamplingAuthAndBudgetGating: auth = _make_user_api_key_auth(models=["gpt-4o"]) params = MagicMock() - params.modelPreferences = None + params.model_preferences = None params.messages = [] params.systemPrompt = None - params.maxTokens = 100 + params.max_tokens = 100 params.temperature = None - params.stopSequences = None + params.stop_sequences = None params.tools = None - params.toolChoice = None + params.tool_choice = None params.metadata = None budget_error = ErrorData(code=-1, message="ExceededBudget: over limit") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_response_conversion.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_response_conversion.py index bb17a8f7104..ba130f34964 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_response_conversion.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_response_conversion.py @@ -54,13 +54,13 @@ class TestConvertOpenAIResponseToMcpResult: assert isinstance(result.content, TextContent) assert result.content.text == "hello world" assert result.role == "assistant" - assert result.stopReason == "endTurn" + assert result.stop_reason== "endTurn" def test_should_map_length_finish_reason_to_max_tokens(self): result = _convert_openai_response_to_mcp_result( _response(content="truncated", finish_reason="length"), "gpt-4o" ) - assert result.stopReason == "maxTokens" + assert result.stop_reason== "maxTokens" def test_should_prefer_actual_model_from_response(self): result = _convert_openai_response_to_mcp_result( @@ -79,7 +79,7 @@ class TestConvertOpenAIResponseToMcpResult: "gpt-4o", ) assert isinstance(result, CreateMessageResultWithTools) - assert result.stopReason == "toolUse" + assert result.stop_reason== "toolUse" tool_uses = [c for c in result.content if isinstance(c, ToolUseContent)] assert len(tool_uses) == 1 assert tool_uses[0].name == "get_weather" @@ -113,7 +113,7 @@ class TestConvertMcpToolsToOpenAI: def test_should_convert_tool_with_schema(self): schema = {"type": "object", "properties": {"q": {"type": "string"}}} tool = SimpleNamespace( - name="search", description="search the web", inputSchema=schema + name="search", description="search the web", input_schema=schema ) result = _convert_mcp_tools_to_openai([tool]) assert result == [ @@ -128,7 +128,7 @@ class TestConvertMcpToolsToOpenAI: ] def test_should_default_description_and_parameters(self): - tool = SimpleNamespace(name="noop", description=None, inputSchema=None) + tool = SimpleNamespace(name="noop", description=None, input_schema=None) result = _convert_mcp_tools_to_openai([tool]) fn = result[0]["function"] assert fn["description"] == "" @@ -151,7 +151,7 @@ class TestConvertMcpToolChoiceToOpenAI: class TestConvertImageAndAudioContent: def test_should_convert_image_to_data_uri(self): - content = SimpleNamespace(type="image", data="aGVsbG8=", mimeType="image/jpeg") + content = SimpleNamespace(type="image", data="aGVsbG8=", mime_type="image/jpeg") result = _convert_single_content(content) assert result == { "type": "image_url", @@ -159,20 +159,20 @@ class TestConvertImageAndAudioContent: } def test_should_map_audio_mime_to_format(self): - content = SimpleNamespace(type="audio", data="Zm9v", mimeType="audio/mp3") + content = SimpleNamespace(type="audio", data="Zm9v", mime_type="audio/mp3") result = _convert_single_content(content) assert result["type"] == "input_audio" assert result["input_audio"] == {"data": "Zm9v", "format": "mp3"} def test_should_default_unknown_audio_mime_to_wav(self): - content = SimpleNamespace(type="audio", data="Zm9v", mimeType="audio/weird") + content = SimpleNamespace(type="audio", data="Zm9v", mime_type="audio/weird") result = _convert_single_content(content) assert result["input_audio"]["format"] == "wav" def test_should_flatten_list_content(self): items = [ SimpleNamespace(type="text", text="a"), - SimpleNamespace(type="image", data="x", mimeType="image/png"), + SimpleNamespace(type="image", data="x", mime_type="image/png"), ] result = _convert_mcp_content_to_openai(items) assert isinstance(result, list) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_tool_conversion.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_tool_conversion.py index b4b219e958c..167847afe1f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_tool_conversion.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_tool_conversion.py @@ -10,6 +10,8 @@ import json from types import SimpleNamespace from typing import Any, Dict +from mcp.types import TextContent, ToolResultContent + from litellm.proxy._experimental.mcp_server.sampling_handler import ( _convert_mcp_messages_to_openai, _convert_single_content, @@ -21,8 +23,8 @@ from litellm.proxy._experimental.mcp_server.sampling_handler import ( # --------------------------------------------------------------------------- -def _text(text: str) -> SimpleNamespace: - return SimpleNamespace(type="text", text=text) +def _text(text: str) -> TextContent: + return TextContent(type="text", text=text) def _tool_use(*, name: str, tool_id: str, input_data: Dict[str, Any]) -> SimpleNamespace: @@ -31,11 +33,9 @@ def _tool_use(*, name: str, tool_id: str, input_data: Dict[str, Any]) -> SimpleN def _tool_result( *, tool_use_id: str, content: Any = None, is_error: bool = False -) -> SimpleNamespace: - if content is None: - content = [] - return SimpleNamespace( - type="tool_result", toolUseId=tool_use_id, content=content, isError=is_error +) -> ToolResultContent: + return ToolResultContent( + tool_use_id=tool_use_id, content=[] if content is None else content, is_error=is_error ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 33e14736357..47ec25a90f7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1,6 +1,7 @@ import asyncio import contextlib import contextvars +import json import os from datetime import datetime, timedelta from types import SimpleNamespace @@ -11,6 +12,7 @@ import pytest from fastapi import HTTPException from mcp import ReadResourceResult, Resource from mcp.types import ( + INVALID_REQUEST, BlobResourceContents, CallToolResult, Prompt, @@ -18,9 +20,11 @@ from mcp.types import ( TextContent, TextResourceContents, ) +from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS, LATEST_HANDSHAKE_VERSION from pydantic import TypeAdapter -from starlette.types import Receive, Scope, Send +from starlette.types import Message, Receive, Scope, Send +from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var from litellm.proxy._types import ( LiteLLM_MCPServerTable, MCPTransport, @@ -30,6 +34,17 @@ from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPOAuthMetadata, MCPServer +def test_mcp_available_on_sdk2(): + from importlib.metadata import version + + from packaging.version import Version + + from litellm.proxy._experimental.mcp_server.server import MCP_AVAILABLE + + assert Version("2.2.0") <= Version(version("mcp")) < Version("3") + assert MCP_AVAILABLE is True + + def _rendered_log_message(call): message = str(call.args[0]) values = call.args[1:] @@ -67,8 +82,22 @@ def cleanup_mcp_global_state(): yield + + + +def _call_tool_params(name, arguments=None): + from mcp.types import CallToolRequestParams + + return CallToolRequestParams(name=name, arguments=arguments) + + +def _paged_params(): + from mcp.types import PaginatedRequestParams + + return PaginatedRequestParams() + @pytest.mark.asyncio -async def test_mcp_server_tool_call_body_contains_request_data(): +async def test_mcp_server_tool_call_body_contains_request_data(_mcp_request_ctx): """Test that proxy_server_request body contains name and arguments""" try: from litellm.proxy._experimental.mcp_server.server import ( @@ -117,7 +146,7 @@ async def test_mcp_server_tool_call_body_contains_request_data(): MagicMock(), ): # Call the function - await mcp_server_tool_call(tool_name, tool_arguments) + await mcp_server_tool_call(_mcp_request_ctx(), _call_tool_params(tool_name, tool_arguments)) # Verify the body contains the expected data assert "proxy_server_request" in captured_data @@ -129,7 +158,7 @@ async def test_mcp_server_tool_call_body_contains_request_data(): @pytest.mark.asyncio -async def test_mcp_server_tool_call_forwards_client_headers_to_logging(): +async def test_mcp_server_tool_call_forwards_client_headers_to_logging(_mcp_request_ctx): """The MCP protocol path must hand the connection's client headers to the pre-call pipeline, so logging callbacks and guardrails see them the way the REST path does.""" try: @@ -169,7 +198,7 @@ async def test_mcp_server_tool_call_forwards_client_headers_to_logging(): mock_call_mcp_tool, ): with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()): - await mcp_server_tool_call("test_tool", {"param": "value"}) + await mcp_server_tool_call(_mcp_request_ctx(), _call_tool_params("test_tool", {"param": "value"})) assert captured_headers.get("x-nuid") == "nuid-1" assert captured_headers.get("x-app-id") == "app-1" @@ -178,7 +207,7 @@ async def test_mcp_server_tool_call_forwards_client_headers_to_logging(): @pytest.mark.asyncio -async def test_mcp_server_tool_call_strips_custom_litellm_key_header(): +async def test_mcp_server_tool_call_strips_custom_litellm_key_header(_mcp_request_ctx): """The deployment can rename the proxy key header via general_settings.litellm_key_header_name. The pre-call pipeline only knows that name if it is passed in, so without it the virtual key reaches metadata.headers and proxy_server_request.headers in plaintext.""" @@ -221,7 +250,7 @@ async def test_mcp_server_tool_call_strips_custom_litellm_key_header(): {"litellm_key_header_name": "x-company-key"}, clear=False, ): - await mcp_server_tool_call("test_tool", {"param": "value"}) + await mcp_server_tool_call(_mcp_request_ctx(), _call_tool_params("test_tool", {"param": "value"})) metadata_headers = captured_data["metadata"]["headers"] assert metadata_headers.get("x-nuid") == "nuid-1" @@ -230,7 +259,7 @@ async def test_mcp_server_tool_call_strips_custom_litellm_key_header(): @pytest.mark.asyncio -async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(): +async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_request_ctx): """The MCP session manager serializes handler exceptions as JSON-RPC errors, so a mid-session tool call cannot emit a raw 401 the way the REST path does. mcp_server_tool_call must turn an upstream MCPUpstreamAuthError into an explicit isError result naming the status, not a masked @@ -263,9 +292,9 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(): ): with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()): with patch("litellm.proxy._experimental.mcp_server.server.verbose_logger", mock_logger): - result = await mcp_server_tool_call("test_tool", {"param": "value"}) + result = await mcp_server_tool_call(_mcp_request_ctx(), _call_tool_params("test_tool", {"param": "value"})) - assert result.isError is True + assert result.is_error is True # The dedicated MCPUpstreamAuthError branch (not the generic Exception fallthrough) produces this # specific message and logs at info, never a traceback via verbose_logger.exception. assert "upstream authentication required" in result.content[0].text @@ -1316,7 +1345,7 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails(): tool1 = MagicMock() tool1.name = "working_tool_1" tool1.description = "Working tool 1" - tool1.inputSchema = {} + tool1.input_schema = {} return [tool1] else: # Failing server raises an exception @@ -1692,15 +1721,15 @@ async def test_scoped_list_agent_veto_attributed_for_differently_cased_server_na @pytest.mark.asyncio -async def test_handle_list_tools_converts_permission_httpexception_to_mcp_error(): +async def test_handle_list_tools_converts_permission_httpexception_to_mcp_error(_mcp_request_ctx): """The MCP protocol handler surfaces a permission HTTPException as a clean JSON-RPC error - (McpError, INVALID_REQUEST) carrying the denial message, instead of a raw 500.""" + (MCPError, INVALID_REQUEST) carrying the denial message, instead of a raw 500.""" try: from litellm.proxy._experimental.mcp_server.server import handle_list_tools except ImportError: pytest.skip("MCP server not available") - from mcp.shared.exceptions import McpError + from mcp.shared.exceptions import MCPError from mcp.types import INVALID_REQUEST denial_message = "MCP server 'github' is not available to this key: the key is bound to agent 'agent-123'" @@ -1716,15 +1745,15 @@ async def test_handle_list_tools_converts_permission_httpexception_to_mcp_error( new=AsyncMock(side_effect=denial), ), ): - with pytest.raises(McpError) as exc_info: - await handle_list_tools() + with pytest.raises(MCPError) as exc_info: + await handle_list_tools(_mcp_request_ctx(), _paged_params()) assert exc_info.value.error.code == INVALID_REQUEST assert exc_info.value.error.message == denial_message @pytest.mark.asyncio -async def test_mcp_server_tool_call_renders_denial_message_not_detail_dict(): +async def test_mcp_server_tool_call_renders_denial_message_not_detail_dict(_mcp_request_ctx): try: from litellm.proxy._experimental.mcp_server.server import mcp_server_tool_call except ImportError: @@ -1743,14 +1772,14 @@ async def test_mcp_server_tool_call_renders_denial_message_not_detail_dict(): new=AsyncMock(side_effect=denial), ), ): - result = await mcp_server_tool_call("github-search_issues", {}) + result = await mcp_server_tool_call(_mcp_request_ctx(), _call_tool_params("github-search_issues", {})) - assert result.isError is True + assert result.is_error is True assert result.content[0].text == f"Error: {denial_message}" @pytest.mark.asyncio -async def test_mcp_server_tool_call_body_with_none_arguments(): +async def test_mcp_server_tool_call_body_with_none_arguments(_mcp_request_ctx): """Test that proxy_server_request body handles None arguments correctly""" try: from litellm.proxy._experimental.mcp_server.server import ( @@ -1798,7 +1827,7 @@ async def test_mcp_server_tool_call_body_with_none_arguments(): MagicMock(), ): # Call the function - await mcp_server_tool_call(tool_name, tool_arguments) + await mcp_server_tool_call(_mcp_request_ctx(), _call_tool_params(tool_name, tool_arguments)) # Verify the body contains the expected data assert "proxy_server_request" in captured_data @@ -1967,11 +1996,9 @@ async def test_streamable_http_session_manager_is_stateless(): ("DELETE", b"", False), ), ) -async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless( +async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless(_mcp_request_ctx, debug: bool, method: str, request_body: bytes, stateful: bool ) -> None: - from mcp.server.lowlevel.server import request_ctx - from mcp.shared.context import RequestContext from starlette.requests import Request from starlette.types import Message, Receive, Scope, Send @@ -1988,14 +2015,12 @@ async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless( async def handle_request(request_scope: Scope, receive: Receive, outgoing: Send) -> None: await outgoing({"type": "http.response.start", "status": 200, "headers": []}) await observe_start(send.await_count) - context: Final = RequestContext( - request_id=1, meta=None, session=MagicMock(), lifespan_context=None, request=Request(request_scope) - ) - token: Final = request_ctx.set(context) + context: Final = _mcp_request_ctx(request=Request(request_scope)) + token: Final = active_mcp_request_ctx_var.set(context) try: record_auth_resolution("s1", AuthResolution.stored_user_token) finally: - request_ctx.reset(token) + active_mcp_request_ctx_var.reset(token) await outgoing(body) stateless_handle: Final = AsyncMock(side_effect=handle_request) @@ -4341,7 +4366,7 @@ async def test_list_tools_single_server_unprefixed_names(): tool = MagicMock() tool.name = f"{server.alias}-toolA" if add_prefix else "toolA" tool.description = "desc" - tool.inputSchema = {} + tool.input_schema = {} return [tool] mock_manager._get_tools_from_server = mock_get_tools_from_server @@ -4420,7 +4445,7 @@ async def test_list_tools_multiple_servers_prefixed_names(): # When multiple servers, add_prefix should be True -> prefixed names tool.name = f"{server.alias}-toolA" if add_prefix else "toolA" tool.description = "desc" - tool.inputSchema = {} + tool.input_schema = {} return [tool] mock_manager._get_tools_from_server = mock_get_tools_from_server @@ -4833,22 +4858,22 @@ async def test_list_tools_filters_by_key_team_permissions(): tool1 = MagicMock() tool1.name = "tool1" tool1.description = "Tool 1" - tool1.inputSchema = {} + tool1.input_schema = {} tool2 = MagicMock() tool2.name = "tool2" tool2.description = "Tool 2" - tool2.inputSchema = {} + tool2.input_schema = {} tool3 = MagicMock() tool3.name = "tool3" tool3.description = "Tool 3 - not allowed" - tool3.inputSchema = {} + tool3.input_schema = {} tool4 = MagicMock() tool4.name = "tool4" tool4.description = "Tool 4 - not allowed" - tool4.inputSchema = {} + tool4.input_schema = {} return [tool1, tool2, tool3, tool4] @@ -4944,22 +4969,22 @@ async def test_list_tools_with_team_tool_permissions_inheritance(): tool1 = MagicMock() tool1.name = "tool1" tool1.description = "Tool 1" - tool1.inputSchema = {} + tool1.input_schema = {} tool2 = MagicMock() tool2.name = "tool2" tool2.description = "Tool 2" - tool2.inputSchema = {} + tool2.input_schema = {} tool3 = MagicMock() tool3.name = "tool3" tool3.description = "Tool 3" - tool3.inputSchema = {} + tool3.input_schema = {} tool4 = MagicMock() tool4.name = "tool4" tool4.description = "Tool 4" - tool4.inputSchema = {} + tool4.input_schema = {} return [tool1, tool2, tool3, tool4] @@ -5041,17 +5066,17 @@ async def test_list_tools_with_no_tool_permissions_shows_all(): tool1 = MagicMock() tool1.name = "tool1" tool1.description = "Tool 1" - tool1.inputSchema = {} + tool1.input_schema = {} tool2 = MagicMock() tool2.name = "tool2" tool2.description = "Tool 2" - tool2.inputSchema = {} + tool2.input_schema = {} tool3 = MagicMock() tool3.name = "tool3" tool3.description = "Tool 3" - tool3.inputSchema = {} + tool3.input_schema = {} return [tool1, tool2, tool3] @@ -5142,22 +5167,22 @@ async def test_list_tools_strips_prefix_when_matching_permissions(): tool1 = MagicMock() tool1.name = "GITMCP-fetch_litellm_documentation" # Prefixed tool1.description = "Fetch docs" - tool1.inputSchema = {} + tool1.input_schema = {} tool2 = MagicMock() tool2.name = "GITMCP-search_litellm_documentation" # Prefixed, not in allowed list tool2.description = "Search docs" - tool2.inputSchema = {} + tool2.input_schema = {} tool3 = MagicMock() tool3.name = "GITMCP-search_litellm_code" # Prefixed tool3.description = "Search code" - tool3.inputSchema = {} + tool3.input_schema = {} tool4 = MagicMock() tool4.name = "GITMCP-fetch_generic_url_content" # Prefixed, not in allowed list tool4.description = "Fetch URL" - tool4.inputSchema = {} + tool4.input_schema = {} return [tool1, tool2, tool3, tool4] @@ -7361,7 +7386,7 @@ async def test_execute_mcp_tool_rest_server_id_authoritative_for_unprefixed_tool captured.update(kwargs) return mcp_module.CallToolResult( content=[TextContent(type="text", text="ok")], - isError=False, + is_error=False, ) with ( @@ -7440,7 +7465,7 @@ async def test_execute_mcp_tool_strips_a_prefix_that_contains_the_separator(): captured.update(kwargs) return mcp_module.CallToolResult( content=[TextContent(type="text", text="ok")], - isError=False, + is_error=False, ) with ( @@ -7507,7 +7532,7 @@ async def test_execute_mcp_tool_rest_server_id_injects_requested_server_credenti fake_client.call_tool = AsyncMock( return_value=mcp_module.CallToolResult( content=[TextContent(type="text", text="ok")], - isError=False, + is_error=False, ) ) @@ -7711,7 +7736,7 @@ async def test_execute_mcp_tool_rest_hyphenated_upstream_tool_name_routes_to_req captured.update(kwargs) return mcp_module.CallToolResult( content=[TextContent(type="text", text="ok")], - isError=False, + is_error=False, ) with ( @@ -7874,7 +7899,7 @@ async def test_execute_mcp_tool_rest_unresolved_prefixed_name_routes_to_requeste captured.update(kwargs) return mcp_module.CallToolResult( content=[TextContent(type="text", text="ok")], - isError=False, + is_error=False, ) with ( @@ -8356,20 +8381,24 @@ class TestMCPMetaTraceCarrier: (e.g. ``litellm.team.id``). Dropping it at the source is the regression guard.""" from types import SimpleNamespace - from mcp.types import RequestParams + from mcp.types import CallToolRequestParams from litellm.proxy._experimental.mcp_server.server import ( _mcp_meta_trace_carrier, ) - meta = RequestParams.Meta.model_validate( + meta = CallToolRequestParams.model_validate( { - "traceparent": "00-11111111111111111111111111111111-2222222222222222-01", - "tracestate": "rojo=1", - "baggage": "litellm.team.id=spoofed-team,litellm.metadata.user_api_key_user_id=attacker", - "progressToken": "p1", - } - ) + "name": "t", + "_meta": { + "traceparent": "00-11111111111111111111111111111111-2222222222222222-01", + "tracestate": "rojo=1", + "baggage": "litellm.team.id=spoofed-team,litellm.metadata.user_api_key_user_id=attacker", + "progressToken": "p1", + }, + }, + by_name=False, + ).meta carrier = _mcp_meta_trace_carrier(SimpleNamespace(meta=meta)) assert carrier == { "traceparent": "00-11111111111111111111111111111111-2222222222222222-01", @@ -8380,7 +8409,7 @@ class TestMCPMetaTraceCarrier: def test_none_when_no_trace_context(self): from types import SimpleNamespace - from mcp.types import RequestParams + from mcp.types import CallToolRequestParams from litellm.proxy._experimental.mcp_server.server import ( _mcp_meta_trace_carrier, @@ -8388,17 +8417,14 @@ class TestMCPMetaTraceCarrier: assert _mcp_meta_trace_carrier(None) is None assert _mcp_meta_trace_carrier(SimpleNamespace(meta=None)) is None - only_progress = RequestParams.Meta.model_validate({"progressToken": "p1"}) + only_progress = CallToolRequestParams.model_validate({"name": "t", "_meta": {"progressToken": "p1"}}, by_name=False).meta assert _mcp_meta_trace_carrier(SimpleNamespace(meta=only_progress)) is None @pytest.mark.asyncio -async def test_stateful_mcp_tool_call_uses_current_requests_otel_destinations() -> None: +async def test_stateful_mcp_tool_call_uses_current_requests_otel_destinations(_mcp_request_ctx) -> None: from types import SimpleNamespace - from mcp.server.lowlevel.server import request_ctx - from mcp.shared.context import RequestContext - from litellm.integrations.otel.model.destination import OtelDestination from litellm.integrations.otel.plumbing.context import ( request_destinations, @@ -8441,20 +8467,14 @@ async def test_stateful_mcp_tool_call_uses_current_requests_otel_destinations() set_auth_context(None, raw_headers={}) destinations_token = set_request_destinations((initialized_destination,)) scope = {_MCP_DESTINATIONS_SCOPE_KEY: (current_destination,)} - current_request_context = RequestContext( - request_id=1, - meta=None, - session=SimpleNamespace(), - lifespan_context=None, - request=SimpleNamespace(scope=scope), - ) - request_token = request_ctx.set(current_request_context) + current_request_context = _mcp_request_ctx(request=SimpleNamespace(scope=scope)) + request_token = active_mcp_request_ctx_var.set(current_request_context) try: - result = await mcp_server_tool_call("otelcontext-observe", {}) - assert result.isError is False + result = await mcp_server_tool_call(current_request_context, _call_tool_params("otelcontext-observe", {})) + assert result.is_error is False assert request_destinations() == (initialized_destination,) finally: - request_ctx.reset(request_token) + active_mcp_request_ctx_var.reset(request_token) reset_request_destinations(destinations_token) global_mcp_tool_registry.tools.pop("otelcontext-observe", None) global_mcp_server_manager.registry.pop(server.server_id, None) @@ -8591,7 +8611,7 @@ def test_extract_mcp_tool_result_error_message(): @pytest.mark.asyncio async def test_fire_mcp_tool_call_logging_iserror_logs_failure(): - """Regression test: a CallToolResult with isError=True must go + """Regression test: a CallToolResult with is_error=True must go down the failure logging path (async_failure_handler + post_call_failure_hook), never async_success_handler.""" from litellm.proxy._experimental.mcp_server.exceptions import MCPToolResultError @@ -8631,7 +8651,7 @@ async def test_fire_mcp_tool_call_logging_iserror_logs_failure(): @pytest.mark.asyncio async def test_fire_mcp_tool_call_logging_success_path_unchanged(): - """isError=False must keep today's behavior: success handler fires, no + """is_error=False must keep today's behavior: success handler fires, no failure logging, no post_call_failure_hook.""" from litellm.proxy._experimental.mcp_server.server import ( _fire_mcp_tool_call_logging, @@ -8750,7 +8770,7 @@ def _real_mcp_logging_obj(call_id: str): @pytest.mark.asyncio async def test_fire_mcp_tool_call_logging_iserror_builds_failure_payload(monkeypatch): - """The standard logging payload for an isError=True result must carry + """The standard logging payload for an is_error=True result must carry status='failure' with the tool's error text, so OTel (whose _parse_error keys off status) marks the MCP span ERROR.""" import litellm @@ -8781,7 +8801,7 @@ async def test_fire_mcp_tool_call_logging_iserror_builds_failure_payload(monkeyp @pytest.mark.asyncio async def test_fire_mcp_tool_call_logging_success_builds_success_payload(monkeypatch): - """isError=False still produces a status='success' payload.""" + """is_error=False still produces a status='success' payload.""" import litellm from litellm.proxy._experimental.mcp_server.server import ( _fire_mcp_tool_call_logging, @@ -8807,9 +8827,9 @@ async def test_fire_mcp_tool_call_logging_success_builds_success_payload(monkeyp @pytest.mark.asyncio async def test_fire_mcp_tool_call_logging_iserror_emits_otel_error_span(monkeypatch): - """End-to-end regression for the OTel symptom: an isError=True tool + """End-to-end regression for the OTel symptom: an is_error=True tool result must reach OTel as an MCP span with StatusCode.ERROR and the tool's - error message, while isError=False stays non-error.""" + error message, while is_error=False stays non-error.""" pytest.importorskip("opentelemetry") from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( InMemorySpanExporter, @@ -9054,7 +9074,7 @@ async def test_aggregate_listing_reports_per_server_outcomes(): tool1 = MagicMock() tool1.name = "working_tool_1" tool1.description = "Working tool 1" - tool1.inputSchema = {} + tool1.input_schema = {} return [tool1] raise MCPServerListError(ServerListFault(tag="upstream_error", status_code=500), server.name) @@ -9103,7 +9123,7 @@ async def test_outcome_keys_use_display_prefix_never_canonical_names(): @pytest.mark.asyncio -async def test_handle_list_tools_attaches_outcome_meta(): +async def test_handle_list_tools_attaches_outcome_meta(_mcp_request_ctx): """The protocol handler returns a ListToolsResult whose _meta carries the per-server outcomes, so MCP clients can tell a degraded listing from a genuinely empty one.""" try: @@ -9139,7 +9159,7 @@ async def test_handle_list_tools_attaches_outcome_meta(): new=AsyncMock(return_value=listing), ), ): - result = await handle_list_tools() + result = await handle_list_tools(_mcp_request_ctx(), _paged_params()) assert isinstance(result, ListToolsResult) wire = result.model_dump(by_alias=True) @@ -9900,7 +9920,7 @@ async def test_list_tools_injects_byok_credential_for_non_oauth2_auth_types(auth tool = MagicMock() tool.name = f"{server.alias}-toolA" if add_prefix else "toolA" tool.description = "desc" - tool.inputSchema = {} + tool.input_schema = {} return [tool] mock_manager = MagicMock() @@ -9928,3 +9948,83 @@ async def test_list_tools_injects_byok_credential_for_non_oauth2_auth_types(auth assert seen_auth_headers == ["personal-api-key"] assert [tool.name for tool in listing.tools] == ["byok-toolA"] + + +@pytest.mark.asyncio +async def test_active_request_ctx_var_feeds_get_current_session(_mcp_request_ctx) -> None: + from litellm.proxy._experimental.mcp_server.server import _get_current_session + + session = SimpleNamespace() + ctx = _mcp_request_ctx(session=session) + token = active_mcp_request_ctx_var.set(ctx) + try: + assert _get_current_session() is session + finally: + active_mcp_request_ctx_var.reset(token) + assert _get_current_session() is None + + +@pytest.mark.asyncio +async def test_active_request_ctx_var_feeds_auth_resolution_recording(_mcp_request_ctx) -> None: + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.mcp_debug import ( + MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, + MCPAuthDiagnostics, + record_auth_resolution, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution + + diagnostics = MCPAuthDiagnostics() + ctx = _mcp_request_ctx(request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics})) + token = active_mcp_request_ctx_var.set(ctx) + try: + record_auth_resolution("s1", AuthResolution.static_token) + finally: + active_mcp_request_ctx_var.reset(token) + + assert diagnostics.resolution() == "static-token" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("header_value", "expected_rejected"), + [ + ("2025-06-18", False), + ("2025-11-25", False), + ("2026-07-28", True), + ("1999-01-01", True), + ], +) +async def test_streamable_http_rejects_modern_protocol_version(header_value: str, expected_rejected: bool) -> None: + from litellm.proxy._experimental.mcp_server import server as mcp_module + from litellm.proxy._experimental.mcp_server.server import unsupported_protocol_version + + scope: Scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [(b"mcp-protocol-version", header_value.encode("latin-1"))], + } + assert (unsupported_protocol_version(scope) == header_value) is expected_rejected + + if not expected_rejected: + return + + sent: list[Message] = [] + + async def receive() -> Message: + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(message: Message) -> None: + sent.append(message) + + await mcp_module.handle_streamable_http_mcp(scope, receive, send) + + start = next(m for m in sent if m["type"] == "http.response.start") + assert start["status"] == 400 + body = json.loads(b"".join(m.get("body", b"") for m in sent if m["type"] == "http.response.body")) + assert body["error"]["code"] == INVALID_REQUEST + assert header_value in body["error"]["message"] + for version in body["error"]["message"].split("supported: ")[1].split(", "): + assert version in HANDSHAKE_PROTOCOL_VERSIONS diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index d449ad06642..dc1eed9ed7f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -1,5 +1,6 @@ import importlib import asyncio +import functools import json import logging import os @@ -22,7 +23,10 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ServerLi # Add the parent directory to the path so we can import litellm +import contextlib + import httpx +import httpx2 from mcp import ReadResourceResult, Resource from mcp.types import ( CallToolResult, @@ -81,6 +85,8 @@ def _reload_mcp_manager_module(): return reloaded + + @pytest.fixture(autouse=True) def enable_eager_mcp_oauth_discovery(monkeypatch): monkeypatch.setenv("LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP", "1") @@ -416,10 +422,10 @@ class TestMCPServerManager: assert "gateway-client" in dump assert "https://org-idp.example/oauth2/token" in dump - async def test_load_servers_from_config_warns_on_invalid_alias(self, caplog): + async def test_load_servers_from_config_warns_on_invalid_alias(self, config_only_mcp_manager_factory, caplog): """Invalid aliases from config should emit warnings during load.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() config = { "validserver": { "alias": "bad/name", @@ -434,10 +440,10 @@ class TestMCPServerManager: assert any("invalid alias 'bad/name'" in message for message in caplog.messages) @pytest.mark.asyncio - async def test_load_servers_from_config_accepts_valid_alias(self, caplog): + async def test_load_servers_from_config_accepts_valid_alias(self, config_only_mcp_manager_factory, caplog): """Valid aliases should be accepted and populate the registry.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() config = { "validserver": { "alias": "friendly_alias", @@ -1207,8 +1213,8 @@ class TestMCPServerManager: assert server.scopes == ["read"] @pytest.mark.asyncio - async def test_load_servers_from_config_non_oauth2_needs_no_flow(self): - manager = MCPServerManager() + async def test_load_servers_from_config_non_oauth2_needs_no_flow(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() config = { "apiserver": { "url": "https://example.com/mcp", @@ -1254,10 +1260,10 @@ class TestMCPServerManager: assert not any("oauth2_id_jag" in message for message in caplog.messages) @pytest.mark.asyncio - async def test_load_servers_from_config_does_not_warn_for_api_key_with_google_sso(self, monkeypatch, caplog): + async def test_load_servers_from_config_does_not_warn_for_api_key_with_google_sso(self, config_only_mcp_manager_factory, monkeypatch, caplog): self._clear_sso_env(monkeypatch) monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() config = { "api_key_server": { "url": "https://example.com/mcp", @@ -1394,9 +1400,9 @@ class TestMCPServerManager: assert server.is_dcr_bridge is False @pytest.mark.asyncio - async def test_load_servers_from_config_coerces_cost_string_to_float(self): + async def test_load_servers_from_config_coerces_cost_string_to_float(self, config_only_mcp_manager_factory): """YAML 1.1 parses `7e-05` as a string; ingest must coerce it to float.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() config = { "google_maps": { "url": "https://example.com/mcp", @@ -1420,9 +1426,9 @@ class TestMCPServerManager: assert isinstance(cost_info["tool_name_to_cost_per_query"]["geocode"], float) @pytest.mark.asyncio - async def test_load_servers_from_config_sets_token_endpoint_auth_method(self): + async def test_load_servers_from_config_sets_token_endpoint_auth_method(self, config_only_mcp_manager_factory): """token_endpoint_auth_method from config is carried onto the MCPServer (LIT-4091).""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() config = { "basic_provider": { "url": "https://example.com/mcp", @@ -1899,7 +1905,7 @@ class TestMCPServerManager: with patch.object(_mgr_mod, "verbose_logger") as mock_log: result = await self._run_call_regular(manager, server) - assert result.isError is True + assert result.is_error is True # A genuine non-auth failure keeps operator visibility at warning level, since call_tool's # raise_on_error demoted the client-layer error log to debug. assert mock_log.warning.called @@ -1933,7 +1939,7 @@ class TestMCPServerManager: proxy_logging_obj=None, ) - assert result.isError is False + assert result.is_error is False assert mock_client.call_tool.call_args.kwargs.get("raise_on_error") is not True def _token_exchange_server(self, server_id: str) -> "MCPServer": @@ -6093,17 +6099,17 @@ class TestMCPServerManager: tool1 = MagicMock() tool1.name = "allowed_tool_1" tool1.description = "This tool is allowed" - tool1.inputSchema = {} + tool1.input_schema = {} tool2 = MagicMock() tool2.name = "blocked_tool" tool2.description = "This tool is not allowed" - tool2.inputSchema = {} + tool2.input_schema = {} tool3 = MagicMock() tool3.name = "allowed_tool_2" tool3.description = "This tool is also allowed" - tool3.inputSchema = {} + tool3.input_schema = {} # Mock the global_mcp_server_manager._get_tools_from_server from litellm.proxy._experimental.mcp_server import rest_endpoints @@ -6143,17 +6149,17 @@ class TestMCPServerManager: tool1 = MagicMock() tool1.name = "tool_1" tool1.description = "Tool 1" - tool1.inputSchema = {} + tool1.input_schema = {} tool2 = MagicMock() tool2.name = "tool_2" tool2.description = "Tool 2" - tool2.inputSchema = {} + tool2.input_schema = {} tool3 = MagicMock() tool3.name = "tool_3" tool3.description = "Tool 3" - tool3.inputSchema = {} + tool3.input_schema = {} # Mock the global_mcp_server_manager._get_tools_from_server from litellm.proxy._experimental.mcp_server import rest_endpoints @@ -6193,12 +6199,12 @@ class TestMCPServerManager: tool1 = MagicMock() tool1.name = "tool_1" tool1.description = "Tool 1" - tool1.inputSchema = {} + tool1.input_schema = {} tool2 = MagicMock() tool2.name = "tool_2" tool2.description = "Tool 2" - tool2.inputSchema = {} + tool2.input_schema = {} # Mock the global_mcp_server_manager._get_tools_from_server from litellm.proxy._experimental.mcp_server import rest_endpoints @@ -6538,7 +6544,7 @@ class TestMCPServerManager: # Return a mock CallToolResult result = MagicMock(spec=CallToolResult) result.content = [{"type": "text", "text": "Tool executed successfully"}] - result.isError = False + result.is_error = False return result mock_client.call_tool.side_effect = mock_call_tool @@ -6569,7 +6575,7 @@ class TestMCPServerManager: # Verify the result assert result is not None - assert result.isError is False + assert result.is_error is False assert len(result.content) > 0 # Verify the MCP client call was awaited exactly once @@ -7887,9 +7893,9 @@ class TestMCPServerTimestamps: assert client.timeout == 0.0 @pytest.mark.asyncio - async def test_load_servers_from_config_preserves_timeout(self): + async def test_load_servers_from_config_preserves_timeout(self, config_only_mcp_manager_factory): """timeout from proxy config is loaded into MCPServer.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() config = { "my_server": { "url": "https://example.com/mcp", @@ -8302,9 +8308,9 @@ class TestMCPServerManagerUpstreamInstructionsCache: assert manager._upstream_initialize_instructions_by_server_id.get("srv") is None @pytest.mark.asyncio - async def test_load_servers_from_config_clears_cache(self): + async def test_load_servers_from_config_clears_cache(self, config_only_mcp_manager_factory): """Reloading config clears any previously cached upstream instructions.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() manager._upstream_initialize_instructions_by_server_id["old"] = "stale" await manager.load_servers_from_config( mcp_servers_config={ @@ -8317,9 +8323,9 @@ class TestMCPServerManagerUpstreamInstructionsCache: assert manager._upstream_initialize_instructions_by_server_id.get("old") is None @pytest.mark.asyncio - async def test_load_servers_reads_instructions_from_config(self): + async def test_load_servers_reads_instructions_from_config(self, config_only_mcp_manager_factory): """instructions field from YAML config is persisted on the MCPServer.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( mcp_servers_config={ "srv_a": { @@ -9989,7 +9995,7 @@ class TestOBOCallToolRetry: user_api_key_auth=None, ) - assert result.isError is True + assert result.is_error is True manager._cred_provider.invalidate_credentials.assert_not_awaited() manager._create_mcp_client.assert_not_awaited() assert first.attempts == 1 @@ -10014,7 +10020,7 @@ class TestOBOCallToolRetry: user_api_key_auth=None, ) - assert result.isError is True + assert result.is_error is True manager._create_mcp_client.assert_awaited_once() assert first.attempts == 1 and retry.attempts == 1 @@ -10093,7 +10099,7 @@ class TestOBOConcurrencyLimit: assert peak_while_blocked == max_concurrent assert inflight["current"] == 0 - assert all(result.isError is False for result in results) + assert all(result.is_error is False for result in results) class TestOBOEndpointDiscovery: @@ -11219,7 +11225,7 @@ class TestOpenAPIRegistryKeyMatchesRegistration: result = await self._call(server, registered_key, "list_pets") - assert result.isError is False + assert result.is_error is False assert result.content[0].text == "dispatched" @pytest.mark.asyncio @@ -11236,7 +11242,7 @@ class TestOpenAPIRegistryKeyMatchesRegistration: result = await self._call(server, registered_key, "read_wiki_contents") - assert result.isError is False + assert result.is_error is False assert result.content[0].text == "dispatched" @pytest.mark.asyncio @@ -11259,7 +11265,7 @@ class TestOpenAPIRegistryKeyMatchesRegistration: result = await self._call(server, registered_key, "petstore-list_pets") - assert result.isError is False + assert result.is_error is False assert result.content[0].text == "dispatched" @pytest.mark.asyncio @@ -11282,7 +11288,7 @@ class TestOpenAPIRegistryKeyMatchesRegistration: result = await self._call(server, registered_key, "list_pets") - assert result.isError is False + assert result.is_error is False assert result.content[0].text == "dispatched" @pytest.mark.asyncio @@ -11299,7 +11305,7 @@ class TestOpenAPIRegistryKeyMatchesRegistration: result = await self._call(server, "petstore-list_pets", "delete_pet") - assert result.isError is True + assert result.is_error is True assert "not found in registry" in result.content[0].text @@ -11796,7 +11802,7 @@ class TestOpenApiHandlerRelaysUpstreamAuth: with patch.object(global_mcp_tool_registry, "get_tool", return_value=tool): result = await manager._call_openapi_tool_handler(self._server(), "list_reports", {}) - assert result.isError is True + assert result.is_error is True assert "upstream returned HTTP 503" in result.content[0].text @@ -11814,9 +11820,9 @@ class TestConfigServerIdPinning: } @pytest.mark.asyncio - async def test_derived_id_churns_when_connection_fields_change(self): + async def test_derived_id_churns_when_connection_fields_change(self, config_only_mcp_manager_factory): """The behavior the pin exists to escape: editing the url mints a brand-new id.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config()) before = next(iter(manager.config_mcp_servers)) @@ -11828,8 +11834,8 @@ class TestConfigServerIdPinning: assert before != after @pytest.mark.asyncio - async def test_pinned_id_survives_url_transport_auth_and_alias_edits(self): - manager = MCPServerManager() + async def test_pinned_id_survives_url_transport_auth_and_alias_edits(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) assert list(manager.config_mcp_servers) == ["docs-prod-1"] @@ -11850,8 +11856,8 @@ class TestConfigServerIdPinning: assert manager.config_mcp_servers["docs-prod-1"].url == "https://prod.example.com/mcp" @pytest.mark.asyncio - async def test_absent_server_id_keeps_the_derived_hash(self): - manager = MCPServerManager() + async def test_absent_server_id_keeps_the_derived_hash(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config()) @@ -11866,15 +11872,15 @@ class TestConfigServerIdPinning: @pytest.mark.asyncio @pytest.mark.parametrize("bad_value", ["", " ", 123, True, ["docs-prod-1"]]) - async def test_blank_or_non_string_server_id_is_rejected(self, bad_value: Any): - manager = MCPServerManager() + async def test_blank_or_non_string_server_id_is_rejected(self, config_only_mcp_manager_factory, bad_value: Any): + manager = config_only_mcp_manager_factory() with pytest.raises(ValueError, match="server_id must be a non-empty string"): await manager.load_servers_from_config(self._config(server_id=bad_value)) @pytest.mark.asyncio - async def test_two_servers_pinning_the_same_id_are_rejected(self): - manager = MCPServerManager() + async def test_two_servers_pinning_the_same_id_are_rejected(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() config: Dict[str, Any] = { "docs_server": {"url": "https://a.example.com/mcp", "server_id": "shared-id"}, "wiki_server": {"url": "https://b.example.com/mcp", "server_id": "shared-id"}, @@ -11884,9 +11890,9 @@ class TestConfigServerIdPinning: await manager.load_servers_from_config(config) @pytest.mark.asyncio - async def test_pinned_id_colliding_with_a_derived_id_is_rejected(self): + async def test_pinned_id_colliding_with_a_derived_id_is_rejected(self, config_only_mcp_manager_factory): """A pin that lands on another entry's derived hash collides just as hard.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() derived = manager._generate_stable_server_id( server_name="docs_server", url="https://a.example.com/mcp", @@ -11903,14 +11909,14 @@ class TestConfigServerIdPinning: await manager.load_servers_from_config(config) @pytest.mark.asyncio - async def test_pinned_id_colliding_with_a_db_backed_server_is_rejected(self): + async def test_pinned_id_colliding_with_a_db_backed_server_is_rejected(self, config_only_mcp_manager_factory): """get_registry() is ``config | registry``, so the db row would hide the config server. The registry is seeded by hand because on a real startup the config loads before the database does, so this check only fires on a later reload. The startup ordering is covered by ``test_db_row_arriving_on_a_pinned_config_id_warns``; the warning there is not redundant. """ - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() manager.registry["db-uuid-1"] = MCPServer( server_id="db-uuid-1", name="db_server", @@ -11922,9 +11928,9 @@ class TestConfigServerIdPinning: await manager.load_servers_from_config(self._config(server_id="db-uuid-1")) @pytest.mark.asyncio - async def test_derived_id_matching_a_db_backed_server_is_not_rejected(self): + async def test_derived_id_matching_a_db_backed_server_is_not_rejected(self, config_only_mcp_manager_factory): """Only a pinned id is an authoring error; a hash collision must not fail startup.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() derived = manager._generate_stable_server_id( server_name="docs_server", url="https://example.com/mcp", @@ -11944,8 +11950,8 @@ class TestConfigServerIdPinning: assert derived in manager.config_mcp_servers @pytest.mark.asyncio - async def test_pinned_id_is_stripped_of_surrounding_whitespace(self): - manager = MCPServerManager() + async def test_pinned_id_is_stripped_of_surrounding_whitespace(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id=" docs-prod-1 ")) @@ -11985,9 +11991,9 @@ class TestConfigServerIdPinning: await manager.reload_servers_from_database() @pytest.mark.asyncio - async def test_db_row_arriving_on_a_pinned_config_id_warns(self, caplog): + async def test_db_row_arriving_on_a_pinned_config_id_warns(self, config_only_mcp_manager_factory, caplog): """The db row loads after config on startup, so the config server is hidden then, not at load.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) with caplog.at_level(logging.WARNING, logger="LiteLLM"): @@ -11997,8 +12003,8 @@ class TestConfigServerIdPinning: assert manager.get_registry()["docs-prod-1"].url == "https://db.example.com/mcp" @pytest.mark.asyncio - async def test_db_row_with_a_distinct_id_does_not_warn(self, caplog): - manager = MCPServerManager() + async def test_db_row_with_a_distinct_id_does_not_warn(self, config_only_mcp_manager_factory, caplog): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) with caplog.at_level(logging.WARNING, logger="LiteLLM"): @@ -12008,9 +12014,9 @@ class TestConfigServerIdPinning: assert set(manager.get_registry()) == {"docs-prod-1", "db-uuid-1"} @pytest.mark.asyncio - async def test_pinned_id_matching_another_entrys_server_name_is_rejected(self): + async def test_pinned_id_matching_another_entrys_server_name_is_rejected(self, config_only_mcp_manager_factory): """expand_permission_list resolves against registry keys first, so this steals the grants.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"): await manager.load_servers_from_config( @@ -12025,8 +12031,8 @@ class TestConfigServerIdPinning: ) @pytest.mark.asyncio - async def test_pinned_id_matching_another_entrys_alias_is_rejected(self): - manager = MCPServerManager() + async def test_pinned_id_matching_another_entrys_alias_is_rejected(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"): await manager.load_servers_from_config( @@ -12045,17 +12051,17 @@ class TestConfigServerIdPinning: ) @pytest.mark.asyncio - async def test_pinning_a_servers_own_name_is_allowed(self): + async def test_pinning_a_servers_own_name_is_allowed(self, config_only_mcp_manager_factory): """The most natural pin an operator writes; it resolves to the same server either way.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="docs_server")) assert list(manager.config_mcp_servers) == ["docs_server"] @pytest.mark.asyncio - async def test_pinning_a_servers_own_alias_is_allowed(self): - manager = MCPServerManager() + async def test_pinning_a_servers_own_alias_is_allowed(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(alias="docs", server_id="docs")) @@ -12063,9 +12069,9 @@ class TestConfigServerIdPinning: @pytest.mark.asyncio @pytest.mark.parametrize("aliasing_entry_first", [True, False]) - async def test_pinning_own_name_that_is_another_entrys_alias_is_rejected(self, aliasing_entry_first: bool): + async def test_pinning_own_name_that_is_another_entrys_alias_is_rejected(self, config_only_mcp_manager_factory, aliasing_entry_first: bool): """A grant naming 'docs_server' reaches both servers unpinned; the pin would narrow it to one.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() wiki = ( "wiki_server", {"alias": "docs_server", "url": "https://wiki.example.com/mcp", "transport": MCPTransport.http}, @@ -12079,8 +12085,8 @@ class TestConfigServerIdPinning: await manager.load_servers_from_config(dict((wiki, docs) if aliasing_entry_first else (docs, wiki))) @pytest.mark.asyncio - async def test_pinning_own_name_that_is_another_entrys_mapped_alias_is_rejected(self): - manager = MCPServerManager() + async def test_pinning_own_name_that_is_another_entrys_mapped_alias_is_rejected(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"): await manager.load_servers_from_config( @@ -12096,9 +12102,9 @@ class TestConfigServerIdPinning: ) @pytest.mark.asyncio - async def test_pinning_own_alias_shared_with_a_later_entry_is_rejected(self): + async def test_pinning_own_alias_shared_with_a_later_entry_is_rejected(self, config_only_mcp_manager_factory): """Nothing rejects duplicate aliases, so the first entry's pin would answer the second's grants.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() with pytest.raises(ValueError, match="server_name or alias of MCP server 'docs_server'"): await manager.load_servers_from_config( @@ -12118,9 +12124,9 @@ class TestConfigServerIdPinning: ) @pytest.mark.asyncio - async def test_own_name_pin_resolves_grants_like_the_unpinned_name(self): + async def test_own_name_pin_resolves_grants_like_the_unpinned_name(self, config_only_mcp_manager_factory): """The negative control: a sole-owner self-pin must keep loading and answer the same grants.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( { @@ -12138,9 +12144,9 @@ class TestConfigServerIdPinning: assert manager.expand_permission_list(["wiki"]) == [wiki_id] @pytest.mark.asyncio - async def test_derived_id_is_not_checked_against_names(self): + async def test_derived_id_is_not_checked_against_names(self, config_only_mcp_manager_factory): """Unpinned configs must keep loading; only a pinned id can be an authoring error.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( { @@ -12152,9 +12158,9 @@ class TestConfigServerIdPinning: assert len(manager.config_mcp_servers) == 2 @pytest.mark.asyncio - async def test_shadow_warning_is_not_repeated_on_every_reload(self, caplog): + async def test_shadow_warning_is_not_repeated_on_every_reload(self, config_only_mcp_manager_factory, caplog): """reload_servers_from_database runs on the config-reload timer; one warning, not one a tick.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) with caplog.at_level(logging.WARNING, logger="LiteLLM"): @@ -12167,8 +12173,8 @@ class TestConfigServerIdPinning: assert second_round == first_round @pytest.mark.asyncio - async def test_shadow_warning_fires_again_when_the_shadowed_set_changes(self, caplog): - manager = MCPServerManager() + async def test_shadow_warning_fires_again_when_the_shadowed_set_changes(self, config_only_mcp_manager_factory, caplog): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) with caplog.at_level(logging.WARNING, logger="LiteLLM"): @@ -12179,9 +12185,9 @@ class TestConfigServerIdPinning: assert len([m for m in caplog.messages if "database entry takes precedence" in m]) == 2 @pytest.mark.asyncio - async def test_pinned_id_matching_a_mapped_alias_is_rejected(self): + async def test_pinned_id_matching_a_mapped_alias_is_rejected(self, config_only_mcp_manager_factory): """An alias can also arrive from litellm_settings.mcp_aliases; it is reserved just the same.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"): await manager.load_servers_from_config( @@ -12197,8 +12203,8 @@ class TestConfigServerIdPinning: ) @pytest.mark.asyncio - async def test_pinning_a_servers_own_mapped_alias_is_allowed(self): - manager = MCPServerManager() + async def test_pinning_a_servers_own_mapped_alias_is_allowed(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( self._config(server_id="docs"), @@ -12208,9 +12214,9 @@ class TestConfigServerIdPinning: assert list(manager.config_mcp_servers) == ["docs"] @pytest.mark.asyncio - async def test_mapped_alias_for_an_unknown_server_reserves_nothing(self): + async def test_mapped_alias_for_an_unknown_server_reserves_nothing(self, config_only_mcp_manager_factory): """A dangling mcp_aliases entry is never applied, so it must not fail an unrelated pin.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( self._config(server_id="wiki"), @@ -12220,9 +12226,9 @@ class TestConfigServerIdPinning: assert list(manager.config_mcp_servers) == ["wiki"] @pytest.mark.asyncio - async def test_config_id_that_is_a_db_server_name_warns(self, caplog): + async def test_config_id_that_is_a_db_server_name_warns(self, config_only_mcp_manager_factory, caplog): """The mirror of the shadow case: here the config entry captures the db server's grants.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="db_server")) with caplog.at_level(logging.WARNING, logger="LiteLLM"): @@ -12231,8 +12237,8 @@ class TestConfigServerIdPinning: assert any("db_server" in m and "name or alias of a database-backed" in m for m in caplog.messages) @pytest.mark.asyncio - async def test_capture_warning_is_not_repeated_on_every_reload(self, caplog): - manager = MCPServerManager() + async def test_capture_warning_is_not_repeated_on_every_reload(self, config_only_mcp_manager_factory, caplog): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="db_server")) with caplog.at_level(logging.WARNING, logger="LiteLLM"): @@ -12242,8 +12248,8 @@ class TestConfigServerIdPinning: assert len([m for m in caplog.messages if "name or alias of a database-backed" in m]) == 1 @pytest.mark.asyncio - async def test_config_id_unrelated_to_db_names_does_not_warn(self, caplog): - manager = MCPServerManager() + async def test_config_id_unrelated_to_db_names_does_not_warn(self, config_only_mcp_manager_factory, caplog): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) with caplog.at_level(logging.WARNING, logger="LiteLLM"): @@ -12252,9 +12258,9 @@ class TestConfigServerIdPinning: assert all("name or alias of a database-backed" not in m for m in caplog.messages) @pytest.mark.asyncio - async def test_mapped_alias_for_a_server_with_its_own_alias_reserves_nothing(self): + async def test_mapped_alias_for_a_server_with_its_own_alias_reserves_nothing(self, config_only_mcp_manager_factory): """load_servers_from_config ignores the mapping when the entry sets alias, so it is free.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( { @@ -12276,9 +12282,9 @@ class TestConfigServerIdPinning: assert len(manager.config_mcp_servers) == 2 @pytest.mark.asyncio - async def test_only_the_first_mapped_alias_for_a_server_is_reserved(self): + async def test_only_the_first_mapped_alias_for_a_server_is_reserved(self, config_only_mcp_manager_factory): """Only the first mapping is applied, so pinning the second one must still load.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( { @@ -12295,15 +12301,15 @@ class TestConfigServerIdPinning: assert "wiki_two" in manager.config_mcp_servers @pytest.mark.asyncio - async def test_invalid_name_is_reported_before_any_entry_body_is_read(self): + async def test_invalid_name_is_reported_before_any_entry_body_is_read(self, config_only_mcp_manager_factory): """The identifier index walks every entry up front, so a bad name must still fail on the name.""" with pytest.raises(Exception, match="Server name cannot contain"): - await MCPServerManager().load_servers_from_config({"my-server": None}) + await config_only_mcp_manager_factory().load_servers_from_config({"my-server": None}) @pytest.mark.asyncio - async def test_a_shadowing_db_server_reports_only_the_shadow_warning(self, caplog): + async def test_a_shadowing_db_server_reports_only_the_shadow_warning(self, config_only_mcp_manager_factory, caplog): """The db row wins the id outright, so the capture message would contradict the shadow one.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="db_server")) with caplog.at_level(logging.WARNING, logger="LiteLLM"): @@ -12314,9 +12320,9 @@ class TestConfigServerIdPinning: assert manager.get_registry()["db_server"].url == "https://db.example.com/mcp" @pytest.mark.asyncio - async def test_an_explicitly_blank_alias_still_blocks_the_mapping(self): + async def test_an_explicitly_blank_alias_still_blocks_the_mapping(self, config_only_mcp_manager_factory): """The loader only consults mcp_aliases when the key is absent, so a blank alias frees it.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( { @@ -12338,9 +12344,9 @@ class TestConfigServerIdPinning: assert manager.config_mcp_servers["wiki"].url == "https://example.com/mcp" @pytest.mark.asyncio - async def test_a_row_that_shadows_one_id_still_reports_capturing_another(self, caplog): + async def test_a_row_that_shadows_one_id_still_reports_capturing_another(self, config_only_mcp_manager_factory, caplog): """Skipping is per identifier, not per row, so the second collision is not lost.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( { "docs_server": { @@ -12710,21 +12716,25 @@ async def test_pre_call_tool_check_honors_guardrail_attached_to_key(monkeypatch, ("none", {"Authorization": "Bearer injected"}, "extra-headers", "Bearer injected"), ], ) -async def test_debug_resolution_matches_final_header_conflict_winner( +async def test_debug_resolution_matches_final_header_conflict_winner(_mcp_request_ctx, config: Literal["stored", "static", "none"], extra_headers: dict[str, str] | None, expected_source: str, expected_authorization: str | None, ) -> None: - from mcp.server.lowlevel.server import request_ctx - from mcp.shared.context import RequestContext + from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var from starlette.requests import Request from pydantic import SecretStr from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import MCPAuthenticatedUser from litellm.proxy._experimental.mcp_server.mcp_debug import MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, MCPAuthDiagnostics from litellm.proxy._experimental.mcp_server.outbound_credentials import ( - ApiKeyConfig, AuthorizationCodeConfig, NoneConfig, ServerSpec, SharedKey, UpstreamCredentialProvider, + ApiKeyConfig, + AuthorizationCodeConfig, + NoneConfig, + ServerSpec, + SharedKey, + UpstreamCredentialProvider, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import OAuthToken from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -12740,10 +12750,11 @@ async def test_debug_resolution_matches_final_header_conflict_winner( store = Store() context = MCPAuthenticatedUser(UserAPIKeyAuth(user_id="alice")) diagnostics = MCPAuthDiagnostics() - token = request_ctx.set(RequestContext( - request_id=1, meta=None, session=MagicMock(), lifespan_context=None, - request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), - )) + token = active_mcp_request_ctx_var.set( + _mcp_request_ctx( + request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), + ) + ) selected = { "stored": AuthorizationCodeConfig(), "static": ApiKeyConfig(key_source=SharedKey(value=SecretStr("static-token"))), @@ -12752,7 +12763,10 @@ async def test_debug_resolution_matches_final_header_conflict_winner( try: auth, remaining = await MCPServerManager()._resolve_v2_auth( server=MCPServer( - server_id="s", name="s", transport="http", url="https://up.example/mcp", + server_id="s", + name="s", + transport="http", + url="https://up.example/mcp", static_headers={"Authorization": "Bearer configured"}, ), spec=ServerSpec(server_id="s", resource="https://up.example/mcp", config=selected), @@ -12768,31 +12782,37 @@ async def test_debug_resolution_matches_final_header_conflict_winner( assert request.headers.get("Authorization") == expected_authorization assert store.calls == (1 if config == "stored" else 0) finally: - request_ctx.reset(token) + active_mcp_request_ctx_var.reset(token) @pytest.mark.asyncio @pytest.mark.parametrize("transport", ["http", "stdio"]) -async def test_debug_reports_legacy_signing_and_non_http_transport(transport: Literal["http", "stdio"]) -> None: - from mcp.server.lowlevel.server import request_ctx - from mcp.shared.context import RequestContext +async def test_debug_reports_legacy_signing_and_non_http_transport(_mcp_request_ctx, transport: Literal["http", "stdio"]) -> None: + from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var from starlette.requests import Request from litellm.proxy._experimental.mcp_server.mcp_debug import MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, MCPAuthDiagnostics from litellm.types.mcp_server.mcp_server_manager import MCPServer diagnostics = MCPAuthDiagnostics() - token = request_ctx.set(RequestContext( - request_id=1, meta=None, session=MagicMock(), lifespan_context=None, - request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), - )) + token = active_mcp_request_ctx_var.set( + _mcp_request_ctx( + request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), + ) + ) try: server = MCPServer( - server_id="signed", name="signed", transport=transport, - url="https://up.example/mcp", auth_type="aws_sigv4", - aws_access_key_id="AKIDEXAMPLE", aws_secret_access_key="test-signing-secret", - aws_region_name="us-east-1", aws_service_name="execute-api", - command="python", args=["-c", "pass"], + server_id="signed", + name="signed", + transport=transport, + url="https://up.example/mcp", + auth_type="aws_sigv4", + aws_access_key_id="AKIDEXAMPLE", + aws_secret_access_key="test-signing-secret", + aws_region_name="us-east-1", + aws_service_name="execute-api", + command="python", + args=["-c", "pass"], ) client = await MCPServerManager()._create_mcp_client(server) if transport == "stdio": @@ -12804,7 +12824,7 @@ async def test_debug_reports_legacy_signing_and_non_http_transport(transport: Li assert request.headers["Authorization"].startswith("AWS4-HMAC-SHA256 ") assert "Credential=AKIDEXAMPLE/" in request.headers["Authorization"] finally: - request_ctx.reset(token) + active_mcp_request_ctx_var.reset(token) @pytest.mark.asyncio @@ -13049,6 +13069,32 @@ class _DiscoveryClock: return self.now +from pydantic import TypeAdapter +from mcp.types import JSONRPCMessage + +_JSONRPC_ADAPTER = TypeAdapter(JSONRPCMessage) + + +@contextlib.contextmanager +def _mcp_upstream(respond): + """Drive the SDK's streamable-HTTP transport off an httpx2 MockTransport; respx only sees httpx.""" + from litellm.experimental_mcp_client.client import MCPClient + + def make_client(self, *args, **kwargs): + return httpx2.AsyncClient( + transport=httpx2.MockTransport(respond), + headers=kwargs.get("headers"), + auth=kwargs.get("auth") or self._resolved_auth or self._aws_auth, + ) + + with ( + patch.object( # test-quality-ok: respx cannot intercept httpx2; inject MockTransport through the client factory + MCPClient, "_create_httpx_client_factory", lambda self: functools.partial(make_client, self) + ) + ): + yield + + class _DiscoveryUpstream: def __init__(self) -> None: self.requests: tuple[tuple[str, str], ...] = () @@ -13057,37 +13103,47 @@ class _DiscoveryUpstream: self.release = asyncio.Event() self.release.set() - async def respond(self, request: httpx.Request) -> httpx.Response: - from mcp.types import JSONRPCMessage, JSONRPCRequest + async def respond(self, request: httpx2.Request) -> httpx2.Response: + from mcp.types import JSONRPCRequest if request.method == "DELETE": - return httpx.Response(200) - payload: Final = JSONRPCMessage.model_validate_json(request.content).root + return httpx2.Response(200) + payload: Final = _JSONRPC_ADAPTER.validate_json(request.content) if not isinstance(payload, JSONRPCRequest): - return httpx.Response(202) + return httpx2.Response(202) self.requests = (*self.requests, (payload.method, request.headers.get("authorization", ""))) if payload.method == "initialize": - return httpx.Response(200, json={ - "jsonrpc": "2.0", "id": payload.id, - "result": {"protocolVersion": "2025-03-26", "serverInfo": {"name": "discovery", "version": "1"}, - "capabilities": {} if self.outcome == "unsupported" else {"prompts": {}, "resources": {}}}, - }) + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "protocolVersion": "2025-03-26", + "serverInfo": {"name": "discovery", "version": "1"}, + "capabilities": {} if self.outcome == "unsupported" else {"prompts": {}, "resources": {}}, + }, + }, + ) self.entered.set() await self.release.wait() if self.outcome == "failure": - return httpx.Response(503) + return httpx2.Response(503) if self.outcome == "cancelled": raise asyncio.CancelledError() if self.outcome == "rejected": - return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, - "error": {"code": -32601, "message": "Unsupported"}}) + return httpx2.Response( + 200, json={"jsonrpc": "2.0", "id": payload.id, "error": {"code": -32601, "message": "Unsupported"}} + ) result: Final = { "prompts/list": {"prompts": [{"name": "example", "description": "original"}]}, "resources/list": {"resources": [{"name": "example", "uri": "test://example", "description": "original"}]}, - "resources/templates/list": {"resourceTemplates": [{"name": "example", "uriTemplate": "test://{name}", "description": "original"}]}, + "resources/templates/list": { + "resourceTemplates": [{"name": "example", "uriTemplate": "test://{name}", "description": "original"}] + }, "tools/list": {"tools": []}, }[payload.method] - return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) @property def initializes(self) -> int: @@ -13106,11 +13162,13 @@ async def test_discovery_cache_reuses_raw_results_and_expires(kind: str) -> None clock: Final = _DiscoveryClock() manager: Final = MCPServerManager(discovery_clock=clock) upstream: Final = _DiscoveryUpstream() - operation: Final = {"prompts": manager.get_prompts_from_server, "resources": manager.get_resources_from_server, - "templates": manager.get_resource_templates_from_server}[kind] + operation: Final = { + "prompts": manager.get_prompts_from_server, + "resources": manager.get_resources_from_server, + "templates": manager.get_resource_templates_from_server, + }[kind] server: Final = _discovery_server() - with respx.mock(base_url="https://discovery.example") as router: - router.route().mock(side_effect=upstream.respond) + with _mcp_upstream(upstream.respond): first: Final = await operation(server, None) assert len(first) == 1 assert first[0].name == "discovery-example" @@ -13136,10 +13194,12 @@ async def test_discovery_cache_empty_results_and_failures(kind: str, outcome: st manager: Final = MCPServerManager() upstream: Final = _DiscoveryUpstream() upstream.outcome = outcome - operation: Final = {"prompts": manager.get_prompts_from_server, "resources": manager.get_resources_from_server, - "templates": manager.get_resource_templates_from_server}[kind] - with respx.mock(base_url="https://discovery.example") as router: - router.route().mock(side_effect=upstream.respond) + operation: Final = { + "prompts": manager.get_prompts_from_server, + "resources": manager.get_resources_from_server, + "templates": manager.get_resource_templates_from_server, + }[kind] + with _mcp_upstream(upstream.respond): assert await operation(_discovery_server(), None) == [] assert await operation(_discovery_server(), None) == [] assert upstream.initializes == (2 if outcome == "failure" else 1) @@ -13158,15 +13218,25 @@ async def test_discovery_cache_isolates_forwarded_credentials_and_shares_static_ server: Final = _discovery_server() first_user: Final = UserAPIKeyAuth(user_id="first") second_user: Final = UserAPIKeyAuth(user_id="second") - with respx.mock(base_url="https://discovery.example") as router: - router.route().mock(side_effect=upstream.respond) + with _mcp_upstream(upstream.respond): for user in (first_user, second_user): assert len(await manager.get_prompts_from_server(server, user)) == 1 assert upstream.initializes == 1 for credential in ("first-secret", "second-secret", "first-secret"): - assert len(await manager.get_prompts_from_server(server, first_user, extra_headers={"Authorization": credential})) == 1 + assert ( + len( + await manager.get_prompts_from_server( + server, first_user, extra_headers={"Authorization": credential} + ) + ) + == 1 + ) assert upstream.initializes == 3 - assert {auth for method, auth in upstream.requests if method == "prompts/list"} == {"", "first-secret", "second-secret"} + assert {auth for method, auth in upstream.requests if method == "prompts/list"} == { + "", + "first-secret", + "second-secret", + } @pytest.mark.asyncio @@ -13176,9 +13246,10 @@ async def test_discovery_cache_coalesces_and_survives_waiter_cancellation() -> N manager: Final = MCPServerManager() upstream: Final = _DiscoveryUpstream() upstream.release.clear() - with respx.mock(base_url="https://discovery.example") as router: - router.route().mock(side_effect=upstream.respond) - tasks: Final = tuple(asyncio.create_task(manager.get_prompts_from_server(_discovery_server(), None)) for _ in range(10)) + with _mcp_upstream(upstream.respond): + tasks: Final = tuple( + asyncio.create_task(manager.get_prompts_from_server(_discovery_server(), None)) for _ in range(10) + ) await asyncio.wait_for(upstream.entered.wait(), timeout=5) tasks[0].cancel() with pytest.raises(asyncio.CancelledError): @@ -13199,8 +13270,7 @@ async def test_discovery_cache_invalidation_during_fetch_does_not_repopulate_old manager: Final = MCPServerManager() upstream: Final = _DiscoveryUpstream() upstream.release.clear() - with respx.mock(base_url="https://discovery.example") as router: - router.route().mock(side_effect=upstream.respond) + with _mcp_upstream(upstream.respond): task: Final = asyncio.create_task(manager.get_prompts_from_server(_discovery_server(), None)) await asyncio.wait_for(upstream.entered.wait(), timeout=5) manager._invalidate_discovery_lists("discovery") @@ -13220,8 +13290,7 @@ async def test_discovery_cache_can_be_disabled(monkeypatch: pytest.MonkeyPatch) monkeypatch.setenv("LITELLM_MCP_DISCOVERY_CACHE_TTL", "0") manager: Final = MCPServerManager() upstream: Final = _DiscoveryUpstream() - with respx.mock(base_url="https://discovery.example") as router: - router.route().mock(side_effect=upstream.respond) + with _mcp_upstream(upstream.respond): assert len(await manager.get_prompts_from_server(_discovery_server(), None)) == 1 assert len(await manager.get_prompts_from_server(_discovery_server(), None)) == 1 assert upstream.initializes == 2 @@ -13345,32 +13414,41 @@ async def test_discovery_cache_tracks_resolved_credentials_across_workers() -> N source: Final = CredentialSource() managers: Final = (MCPServerManager(cred_provider=source), MCPServerManager(cred_provider=source)) server: Final = MCPServer( - server_id="discovery", name="discovery", url="https://discovery.example/mcp", transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id="discovery-client", - authorization_url="https://discovery.example/authorize", token_url="https://discovery.example/token", + server_id="discovery", + name="discovery", + url="https://discovery.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + client_id="discovery-client", + authorization_url="https://discovery.example/authorize", + token_url="https://discovery.example/token", ) user: Final = UserAPIKeyAuth(user_id="same-user", api_key="same-key") upstream: Final = _DiscoveryUpstream() - async def respond(request: httpx.Request) -> httpx.Response: + async def respond(request: httpx2.Request) -> httpx2.Response: response: Final = await upstream.respond(request) if '"prompts/list"' not in request.content.decode(): return response - from mcp.types import JSONRPCMessage, JSONRPCRequest + from mcp.types import JSONRPCRequest - payload: Final = JSONRPCMessage.model_validate_json(request.content).root + payload: Final = _JSONRPC_ADAPTER.validate_json(request.content) assert isinstance(payload, JSONRPCRequest) name: Final = {"Bearer token-a": "account-a", "Bearer token-b": "account-b"}[request.headers["authorization"]] - return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": {"prompts": [{"name": name}]}}) + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": {"prompts": [{"name": name}]}}) - with respx.mock(base_url="https://discovery.example") as router: - router.route().mock(side_effect=respond) + with _mcp_upstream(respond): for manager in managers: - assert [item.name for item in await manager.get_prompts_from_server(server, user)] == ["discovery-account-a"] + assert [item.name for item in await manager.get_prompts_from_server(server, user)] == [ + "discovery-account-a" + ] assert upstream.initializes == 2 source.token = "token-b" for manager in managers: - assert [item.name for item in await manager.get_prompts_from_server(server, user)] == ["discovery-account-b"] + assert [item.name for item in await manager.get_prompts_from_server(server, user)] == [ + "discovery-account-b" + ] assert upstream.initializes == 4 source.token = None for manager in managers: @@ -13397,14 +13475,19 @@ async def test_discovery_resolves_stored_oauth_for_the_requesting_user() -> None store: Final = TokenStore() manager: Final = MCPServerManager(per_user_oauth_token_store=store) server: Final = MCPServer( - server_id="discovery", name="discovery", url="https://discovery.example/mcp", transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id="discovery-client", - authorization_url="https://discovery.example/authorize", token_url="https://discovery.example/token", + server_id="discovery", + name="discovery", + url="https://discovery.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + client_id="discovery-client", + authorization_url="https://discovery.example/authorize", + token_url="https://discovery.example/token", ) user: Final = UserAPIKeyAuth(user_id="requesting-user") upstream: Final = _DiscoveryUpstream() - with respx.mock(base_url="https://discovery.example") as router: - router.route().mock(side_effect=upstream.respond) + with _mcp_upstream(upstream.respond): assert len(await manager.get_prompts_from_server(server, user)) == 1 assert len(await manager.get_prompts_from_server(server, user)) == 1 assert store.calls == (("requesting-user", "discovery"), ("requesting-user", "discovery")) @@ -13506,7 +13589,7 @@ class TestProtectedCredentialPreparation: if dispatch == "managed" else await _handle_local_mcp_tool(add_server_prefix_to_name("echo", get_server_prefix(server)), {}) ) - assert result.isError is True + assert result.is_error is True assert "requires a usable upstream credential" in result.content[0].text assert destination.call_count == 0 @@ -13937,5 +14020,5 @@ async def test_request_selected_during_guardrail_runs_concurrently_with_tool(mon ), timeout=5) assert tool_started.is_set() assert guardrail_started.is_set() is selected - assert result.isError is False + assert result.is_error is False assert result.content[0].text == "executed" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py index e814425c9a2..8cf3bc6fcc7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py @@ -1,7 +1,7 @@ """ Tests for AWS SigV4 authentication in MCP client. -Tests the MCPSigV4Auth httpx.Auth subclass that enables per-request +Tests the MCPSigV4Auth httpx2.Auth subclass that enables per-request SigV4 signing for Bedrock AgentCore MCP servers, plus DB/UI path tests for credential encryption, merge-on-update, and build_from_table. """ @@ -11,7 +11,7 @@ import json import pytest from unittest.mock import patch, MagicMock, AsyncMock -import httpx +import httpx2 from litellm.experimental_mcp_client.client import MCPSigV4Auth, MCPClient from litellm.types.mcp import MCPAuth, MCPTransport @@ -103,7 +103,7 @@ class TestMCPSigV4Auth: aws_service_name="bedrock-agentcore", ) - request = httpx.Request( + request = httpx2.Request( method="POST", url="https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes/test/invocations", headers={"Content-Type": "application/json"}, @@ -128,13 +128,13 @@ class TestMCPSigV4Auth: aws_region_name="us-east-1", ) - request1 = httpx.Request( + request1 = httpx2.Request( method="POST", url="https://example.com/mcp", headers={"Content-Type": "application/json"}, content=b'{"jsonrpc":"2.0","method":"tools/list","id":1}', ) - request2 = httpx.Request( + request2 = httpx2.Request( method="POST", url="https://example.com/mcp", headers={"Content-Type": "application/json"}, @@ -156,7 +156,7 @@ class TestMCPSigV4Auth: aws_region_name="us-east-1", ) - request = httpx.Request( + request = httpx2.Request( method="POST", url="https://example.com/mcp", headers={"Content-Type": "application/json"}, @@ -265,7 +265,7 @@ class TestMCPSigV4AssumeRole: aws_service_name="bedrock-agentcore", ) - request = httpx.Request( + request = httpx2.Request( method="POST", url="https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes/test/invocations", headers={"Content-Type": "application/json"}, @@ -306,7 +306,7 @@ class TestMCPClientSigV4Integration: def test_mcp_client_stores_aws_auth(self): """MCPClient stores the aws_auth parameter.""" - mock_auth = MagicMock(spec=httpx.Auth) + mock_auth = MagicMock(spec=httpx2.Auth) client = MCPClient( server_url="https://example.com/mcp", transport_type=MCPTransport.http, @@ -330,7 +330,7 @@ class TestMCPClientSigV4Integration: factory = client._create_httpx_client_factory() httpx_client = factory( headers={"Content-Type": "application/json"}, - timeout=httpx.Timeout(30.0), + timeout=httpx2.Timeout(30.0), ) # Verify the auth object was actually wired into the httpx client @@ -342,7 +342,7 @@ class TestMCPClientSigV4Integration: aws_access_key_id="AKIAIOSFODNN7EXAMPLE", aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", ) - explicit_auth = MagicMock(spec=httpx.Auth) + explicit_auth = MagicMock(spec=httpx2.Auth) client = MCPClient( server_url="https://example.com/mcp", @@ -353,7 +353,7 @@ class TestMCPClientSigV4Integration: factory = client._create_httpx_client_factory() httpx_client = factory( headers={"Content-Type": "application/json"}, - timeout=httpx.Timeout(30.0), + timeout=httpx2.Timeout(30.0), auth=explicit_auth, ) @@ -370,7 +370,7 @@ class TestMCPClientSigV4Integration: factory = client._create_httpx_client_factory() httpx_client = factory( headers={"Content-Type": "application/json"}, - timeout=httpx.Timeout(30.0), + timeout=httpx2.Timeout(30.0), ) # No auth should be set when aws_auth is not configured assert httpx_client._auth is None @@ -380,7 +380,7 @@ class TestMCPServerManagerSigV4: """Tests for MCPServerManager config loading with SigV4.""" @pytest.mark.asyncio - async def test_load_config_with_aws_sigv4(self): + async def test_load_config_with_aws_sigv4(self, config_only_mcp_manager_factory): """Config loading correctly parses aws_sigv4 auth type and AWS fields.""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, @@ -398,7 +398,7 @@ class TestMCPServerManagerSigV4: } } - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(config) server = next(iter(manager.config_mcp_servers.values())) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index b8935d07774..cb43d2c2592 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -85,6 +85,13 @@ FAKE_VECTORS: dict[str, Vector] = { } + + +def _paged_params(): + from mcp.types import PaginatedRequestParams + + return PaginatedRequestParams() + class RecordingEmbedder: def __init__(self) -> None: self.calls: list[tuple[str, ...]] = [] @@ -113,7 +120,7 @@ class TestSearchMcpTools: assert _names(results) == [FX_TOOL.name, WEATHER_TOOL.name, CALENDAR_TOOL.name] assert not isinstance(results, EmbeddingFailed) assert results[0]["score"] > results[1]["score"] > results[2]["score"] - assert results[0]["inputSchema"] == FX_TOOL.inputSchema + assert results[0]["inputSchema"] == FX_TOOL.input_schema @pytest.mark.asyncio async def test_similarity_threshold_drops_weak_matches(self) -> None: @@ -313,10 +320,10 @@ class TestGetVirtualToolDefinitions: for definition in get_virtual_tool_definitions(): tool = Tool.model_validate(definition) - required_arguments = {name: "x" for name in tool.inputSchema["required"]} - validate(instance=required_arguments, schema=tool.inputSchema) + required_arguments = {name: "x" for name in tool.input_schema["required"]} + validate(instance=required_arguments, schema=tool.input_schema) with pytest.raises(ValidationError): - validate(instance={}, schema=tool.inputSchema) + validate(instance={}, schema=tool.input_schema) def test_all_tools_have_description(self) -> None: for tool in get_virtual_tool_definitions(): @@ -562,7 +569,7 @@ class TestCallToolRestApiVirtualTools: mock_tool = MagicMock() mock_tool.name = "github-create_issue" mock_tool.description = "Create a GitHub issue" - mock_tool.inputSchema = {"type": "object", "properties": {}} + mock_tool.input_schema = {"type": "object", "properties": {}} with patch( "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", @@ -633,7 +640,7 @@ class TestCallToolRestApiVirtualTools: mock_fire_logging.assert_awaited_once() assert mock_execute.await_args.kwargs["name"] == "github-create_issue" - assert result.isError is False + assert result.is_error is False assert result.content[0].text == "Issue created" @pytest.mark.asyncio @@ -730,7 +737,7 @@ class TestCallToolRestApiVirtualTools: ): result = await self._get_call_fn()(request=request, user_api_key_dict=user_api_key_dict) - assert result.isError is False + assert result.is_error is False assert mock_search.await_args.kwargs["user_api_key_dict"] is user_api_key_dict assert json.loads(result.content[0].text) == [ { @@ -766,7 +773,7 @@ class TestCallToolRestApiVirtualTools: ) as mock_search: result = await self._get_call_fn()(request=request, user_api_key_dict=user_api_key_dict) - assert result.isError is False + assert result.is_error is False assert mock_search.await_args.kwargs["top_k"] == DEFAULT_SKILL_SEARCH_TOP_K assert mock_search.await_args.kwargs["query"] == "translate a document" @@ -790,7 +797,7 @@ class TestCallToolRestApiVirtualTools: ): result = await self._get_call_fn()(request=request, user_api_key_dict=user_api_key_dict) - assert result.isError is True + assert result.is_error is True assert result.content[0].text == "set agent_search_embedding_model" def _semantic_request(self, query: str = "FX") -> MagicMock: @@ -835,7 +842,7 @@ class TestCallToolRestApiVirtualTools: assert mock_list.await_args.kwargs["user_api_key_auth"] is user_api_key_dict assert key_limits.pre_call_hook.await_args.kwargs["call_type"] == "aembedding" assert key_limits.pre_call_hook.await_args.kwargs["data"]["model"] == "emb" - assert result.isError is False + assert result.is_error is False assert [t["name"] for t in json.loads(result.content[0].text)] == [FX_TOOL.name] @pytest.mark.asyncio @@ -846,7 +853,7 @@ class TestCallToolRestApiVirtualTools: "litellm.proxy.proxy_server.llm_router", None ): result = await self._get_call_fn()(request=self._semantic_request(), user_api_key_dict=user_api_key_dict) - assert result.isError is True + assert result.is_error is True assert "mcp_tool_search.embedding_model" in result.content[0].text @pytest.mark.asyncio @@ -856,7 +863,7 @@ class TestCallToolRestApiVirtualTools: monkeypatch.setattr(litellm, "mcp_tool_search", {"top_k": 0}) user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True)) result = await self._get_call_fn()(request=self._semantic_request(), user_api_key_dict=user_api_key_dict) - assert result.isError is True + assert result.is_error is True assert "top_k" in result.content[0].text @pytest.mark.asyncio @@ -920,7 +927,7 @@ class TestDispatchVirtualMcpTool: client_ip=None, ) assert result is not None - assert result.isError is True + assert result.is_error is True @pytest.mark.asyncio async def test_routes_search_with_client_ip(self) -> None: @@ -977,7 +984,7 @@ class TestDispatchVirtualMcpTool: name=AGENT_SEARCH_TOOL_NAME, arguments={"query": "x"}, user_api_key_auth=uak, client_ip=None ) assert result is not None - assert result.isError is True + assert result.is_error is True @pytest.mark.asyncio async def test_routes_call_with_client_ip(self) -> None: @@ -1144,78 +1151,28 @@ class TestDispatchVirtualMcpTool: class TestCaptureHostProgressCallback: - """Covers the host progress-forwarding helper extracted from the tool call path.""" + @pytest.mark.parametrize("meta", [None, {}, {"traceparent": "trace"}]) + def test_returns_none_without_progress(self, _mcp_request_ctx, meta) -> None: + from litellm.proxy._experimental.mcp_server.server import _capture_host_progress_callback - def test_returns_none_when_request_context_unavailable(self) -> None: - from litellm.proxy._experimental.mcp_server.server import ( - _capture_host_progress_callback, - ) - - class _NoCtx: - @property - def request_context(self): # type: ignore[no-untyped-def] - raise RuntimeError("no context") - - assert _capture_host_progress_callback(_NoCtx()) is None - - def test_returns_none_when_no_progress_token(self) -> None: - from litellm.proxy._experimental.mcp_server.server import ( - _capture_host_progress_callback, - ) - - host = MagicMock() - host.request_context.meta.progressToken = None - assert _capture_host_progress_callback(host) is None - - def test_returns_callable_when_token_present(self) -> None: - from litellm.proxy._experimental.mcp_server.server import ( - _capture_host_progress_callback, - ) - - host = MagicMock() - host.request_context.meta.progressToken = "tok12345" - host.request_context.session = MagicMock() - assert callable(_capture_host_progress_callback(host)) - - def test_returns_callable_when_token_is_integer(self) -> None: - from litellm.proxy._experimental.mcp_server.server import ( - _capture_host_progress_callback, - ) - - host = MagicMock() - host.request_context.meta.progressToken = 12345 - host.request_context.session = MagicMock() - assert callable(_capture_host_progress_callback(host)) - - def test_returns_callable_when_token_is_zero(self) -> None: - from litellm.proxy._experimental.mcp_server.server import ( - _capture_host_progress_callback, - ) - - host = MagicMock() - host.request_context.meta.progressToken = 0 - host.request_context.session = MagicMock() - assert callable(_capture_host_progress_callback(host)) + assert _capture_host_progress_callback(_mcp_request_ctx(meta=meta)) is None @pytest.mark.asyncio - async def test_forwarded_progress_token_preserves_integer_value(self) -> None: - from litellm.proxy._experimental.mcp_server.server import ( - _capture_host_progress_callback, + @pytest.mark.parametrize("token", ["tok12345", 12345, 0]) + async def test_forwards_wire_progress_token(self, _mcp_request_ctx, token) -> None: + from mcp.types import CallToolRequestParams + + from litellm.proxy._experimental.mcp_server.server import _capture_host_progress_callback + + params = CallToolRequestParams.model_validate( + {"name": "tool", "_meta": {"progressToken": token}}, by_name=False ) - - host = MagicMock() - host.request_context.meta.progressToken = 12345 session = AsyncMock() - host.request_context.session = session - - callback = _capture_host_progress_callback(host) + callback = _capture_host_progress_callback(_mcp_request_ctx(meta=params.meta, session=session)) assert callback is not None await callback(0.5, 1.0) - session.send_progress_notification.assert_awaited_once_with( - progress_token=12345, - progress=0.5, - total=1.0, + progress_token=token, progress=0.5, total=1.0 ) @@ -1223,7 +1180,7 @@ class TestHandleListToolsVirtual: """Covers the protocol list_tools early-return when the flag is enabled.""" @pytest.mark.asyncio - async def test_returns_virtual_tools_when_flag_enabled(self) -> None: + async def test_returns_virtual_tools_when_flag_enabled(self, _mcp_request_ctx) -> None: from litellm.proxy._experimental.mcp_server import server as srv uak = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True)) @@ -1232,9 +1189,9 @@ class TestHandleListToolsVirtual: new_callable=AsyncMock, return_value=(uak, None, None, None, None, None, None), ): - tools = await srv.handle_list_tools() + result = await srv.handle_list_tools(_mcp_request_ctx(), _paged_params()) - assert {t.name for t in tools} == { + assert {t.name for t in result.tools} == { MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME, @@ -1247,7 +1204,7 @@ class TestMcpServerToolCallErrorHandling: isError CallToolResult instead of letting them raise out of the handler.""" @pytest.mark.asyncio - async def test_virtual_tool_error_returns_iserror_not_raised(self) -> None: + async def test_virtual_tool_error_returns_iserror_not_raised(self, _mcp_request_ctx) -> None: from fastapi import HTTPException from litellm.proxy._experimental.mcp_server import server as srv @@ -1265,12 +1222,17 @@ class TestMcpServerToolCallErrorHandling: side_effect=HTTPException(status_code=403, detail="User not allowed to call this tool"), ), ): + from mcp.types import CallToolRequestParams + result = await srv.mcp_server_tool_call( - name=MCP_TOOL_CALL_TOOL_NAME, - arguments={"tool_name": "other-server-tool", "arguments": {}}, + _mcp_request_ctx(), + CallToolRequestParams( + name=MCP_TOOL_CALL_TOOL_NAME, + arguments={"tool_name": "other-server-tool", "arguments": {}}, + ), ) - assert result.isError is True + assert result.is_error is True assert "User not allowed to call this tool" in result.content[0].text diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py index 334bee9800c..ac716bace3c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py @@ -459,7 +459,7 @@ async def test_legacy_local_tool_fallback_still_dispatches_entitled_caller( user_api_key_auth=user, ) - assert result.isError is False + assert result.is_error is False assert executed == [{}] assert "legacy local tool ran" in result.content[0].text @@ -663,12 +663,12 @@ async def test_local_dispatch_reports_the_outcome_instead_of_success(failure: st failure may propagate. `_handle_local_mcp_tool` used to catch every exception and return it as TextContent, and both of - its callers then stamped `isError=False`, so an upstream rejection was served as tool output and + its callers then stamped `is_error=False`, so an upstream rejection was served as tool output and `extract_mcp_tool_result_error_message` logged the request as a success. The two kinds are split by consequence. `MCPUpstreamAuthError` propagates because both renderers know it: the streamable path names the status and the REST path relays a real 401 with the - upstream's WWW-Authenticate. Anything else is reported as `isError=True` right here, because + upstream's WWW-Authenticate. Anything else is reported as `is_error=True` right here, because `call_tool_rest_api` turns an unrecognized exception into HTTP 500 and an upstream 403 or 429 is not a gateway crash. """ @@ -729,7 +729,7 @@ async def test_local_dispatch_reports_the_outcome_instead_of_success(failure: st result = await call # A non-auth upstream failure stays a 200 with isError, so REST does not report it as a gateway 500 - assert result.isError is True + assert result.is_error is True assert "upstream returned HTTP 429" in result.content[0].text diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index ac3ad9ed89e..13af58c15c0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -3177,7 +3177,7 @@ async def test_request_selected_tool_specific_guardrail_applies_to_virtual_execu upstream.assert_not_awaited() else: result: Final = await rest_endpoints.call_tool_rest_api(request, user_api_key_dict=caller) - assert result.isError is False + assert result.is_error is False upstream.assert_awaited_once() assert upstream.await_args.kwargs == {"q": "redacted" if selected else "confidential"} @@ -3198,7 +3198,7 @@ class TestGetToolsForSingleServer: def __init__(self, name, description): self.name = name self.description = description - self.inputSchema = {} + self.input_schema = {} mock_tools = [ MockTool("tool1", "First tool"), @@ -3259,7 +3259,7 @@ class TestGetToolsForSingleServer: def __init__(self, name, description): self.name = name self.description = description - self.inputSchema = {} + self.input_schema = {} mock_tools = [ MockTool("tool1", "First tool"), @@ -3307,7 +3307,7 @@ class TestGetToolsForSingleServer: def __init__(self, name, description): self.name = name self.description = description - self.inputSchema = {} + self.input_schema = {} mock_tools = [ MockTool("tool1", "First tool"), @@ -3360,7 +3360,7 @@ class TestGetToolsForSingleServer: def __init__(self, name, description): self.name = name self.description = description - self.inputSchema = {} + self.input_schema = {} mock_tools = [ MockTool("tool1", "First tool"), @@ -3413,7 +3413,7 @@ class TestGetToolsForSingleServer: def __init__(self, name, description): self.name = name self.description = description - self.inputSchema = {} + self.input_schema = {} mock_tools = [ MockTool("tool1", "First tool"), @@ -3475,7 +3475,7 @@ class TestGetToolsForSingleServer: def __init__(self, name): self.name = name self.description = name - self.inputSchema = {} + self.input_schema = {} mock_tools = [MockTool("tool1"), MockTool("tool2"), MockTool("tool3")] @@ -3903,11 +3903,11 @@ class TestConnectionErrorMessage: assert "secret" not in message def test_closed_connection_explains_incomplete_request(self) -> None: - from mcp import McpError + from mcp import MCPError from mcp.types import ErrorData message: Final = rest_endpoints._connection_error_message( - McpError(ErrorData(code=-32000, message="Connection closed", data="secret-data")), None, 30 + MCPError(code=-32000, message="Connection closed", data="secret-data"), None, 30 ) assert "connection was closed before the request completed" in message assert "secret" not in message @@ -3920,8 +3920,8 @@ class TestConnectionErrorMessage: @pytest.mark.parametrize("sdk_timeout", [True, False]) @pytest.mark.parametrize("read_timeout", [0, 1]) async def test_timeout_message_uses_the_deadline_that_expired(self, sdk_timeout: bool, read_timeout: int) -> None: - from mcp import McpError - from mcp.types import ErrorData + from mcp import MCPError + from mcp.types import REQUEST_TIMEOUT, ErrorData async def operation(client: rest_endpoints.MCPClient) -> dict[str, object]: try: @@ -3930,8 +3930,8 @@ class TestConnectionErrorMessage: if not sdk_timeout: raise try: - raise McpError(ErrorData(code=408, message="secret-sdk-timeout")) from elapsed - except McpError as sdk_error: + raise MCPError(code=REQUEST_TIMEOUT, message="secret-sdk-timeout") from elapsed + except MCPError as sdk_error: raise TimeoutError() from sdk_error payload: Final = NewMCPServerRequest( @@ -3947,11 +3947,11 @@ class TestConnectionErrorMessage: assert "reference" in message.lower() def test_sdk_session_terminated_explains_endpoint_and_retry(self) -> None: - from mcp.shared.exceptions import McpError + from mcp.shared.exceptions import MCPError from mcp.types import ErrorData message: Final = rest_endpoints._connection_error_message( - McpError(ErrorData(code=32600, message="Session terminated")), "https://example.com/mcp", 30.0 + MCPError(code=32600, message="Session terminated"), "https://example.com/mcp", 30.0 ) assert "session was terminated" in message @@ -3962,11 +3962,11 @@ class TestConnectionErrorMessage: @pytest.mark.parametrize("code", [-32700, -32601, -32602, -32603, -32000, 32600, 408]) def test_rpc_errors_include_code_without_echoing_upstream_data(self, code: int) -> None: - from mcp.shared.exceptions import McpError + from mcp.shared.exceptions import MCPError from mcp.types import ErrorData message: Final = rest_endpoints._connection_error_message( - McpError(ErrorData(code=code, message="secret-message", data={"token": "secret-data"})), + MCPError(code=code, message="secret-message", data={"token": "secret-data"}), "https://example.com/secret-path?token=secret-query", 30.0, ) @@ -4150,6 +4150,12 @@ class TestToolResponseMcpInfoEnrichment: "alias": "atlassian", } + from fastapi.encoders import jsonable_encoder + + wire = jsonable_encoder(result[0]) + assert wire["inputSchema"] == {"type": "object"} + assert wire["mcp_info"] == result[0].mcp_info + def test_alias_none_is_explicit_in_mcp_info(self): from mcp.types import Tool as MCPTool diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py index 0252fb9843d..842859e5a1e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py @@ -269,3 +269,17 @@ class TestBuildSyntheticMcpRequest: ) assert request.headers.get("x-user-email") == "alice@corp.example" + + +@pytest.mark.parametrize("field", ["structuredContent", "structured_content"]) +def test_structured_content_redaction_updates_shared_dictionary(field): + from litellm.proxy._experimental.mcp_server.utils import ( + mcp_tool_result_structured_content, + set_mcp_tool_result_structured_content, + ) + + result = {field: {"secret": "sensitive"}, "content": []} + logging_reference = result + assert set_mcp_tool_result_structured_content(result, {"secret": "[REDACTED]"}) is True + assert mcp_tool_result_structured_content(logging_reference) == {"secret": "[REDACTED]"} + assert set(result) == {field, "content"} diff --git a/tests/test_litellm/proxy/config_resolvers/test_settings_store.py b/tests/test_litellm/proxy/config_resolvers/test_settings_store.py index 1182bcdce3c..daf6609325e 100644 --- a/tests/test_litellm/proxy/config_resolvers/test_settings_store.py +++ b/tests/test_litellm/proxy/config_resolvers/test_settings_store.py @@ -315,3 +315,47 @@ def test_settings_store_still_accepts_a_write_to_a_key_the_config_does_not_own() store["max_parallel_requests"] = 7 assert store["max_parallel_requests"] == 7 + + +def test_settings_store_reports_a_config_owned_key_whose_stored_value_differs() -> None: + store: Final = SettingsStore("general_settings") + store.load_yaml({"allowed_ips": ["1.2.3.4"]}) + store.apply_db_row("general_settings", {"allowed_ips": ["1.2.3.4", "5.6.7.8"], "max_parallel_requests": 7}) + + assert store.shadowed_db_keys() == ("allowed_ips",) + assert store.shadows_db_value("allowed_ips") is True + assert store.shadows_db_value("max_parallel_requests") is False + assert store["max_parallel_requests"] == 7 + + +def test_settings_store_reports_no_shadowing_when_the_stored_value_agrees() -> None: + store: Final = SettingsStore("general_settings") + store.load_yaml({"allowed_ips": ["1.2.3.4"]}) + store.apply_db_row("general_settings", {"allowed_ips": ["1.2.3.4"]}) + + assert store.shadowed_db_keys() == () + assert store.shadows_db_value("allowed_ips") is False + + +def test_settings_store_says_the_stored_value_is_ignored_when_it_refuses_a_write() -> None: + store: Final = SettingsStore("general_settings") + store.load_yaml({"allowed_ips": ["1.2.3.4"]}) + store.apply_db_row("general_settings", {"allowed_ips": ["1.2.3.4", "5.6.7.8"]}) + + with pytest.raises(ConfigOwnedKeyError) as refused: + store["allowed_ips"] = ["9.9.9.9"] + + assert refused.value.shadows_db_value is True + assert "stored in the database" in str(refused.value) + + +def test_settings_store_refusal_stays_quiet_about_the_database_when_nothing_is_stored() -> None: + store: Final = SettingsStore("general_settings") + store.load_yaml({"allowed_ips": ["1.2.3.4"]}) + + with pytest.raises(ConfigOwnedKeyError) as refused: + store["allowed_ips"] = ["9.9.9.9"] + + assert refused.value.shadows_db_value is False + assert "stored in the database" not in str(refused.value) + assert "config file" in str(refused.value) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_mcp.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_mcp.py index 137b7d24023..826edab694d 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_mcp.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_mcp.py @@ -376,9 +376,10 @@ class TestCiscoAIDefenseMCPMode: assert sent_payload["result"]["content"][0]["text"] == text_content assert result is None + @pytest.mark.parametrize("use_wrapper", [True, False]) @pytest.mark.asyncio - async def test_mcp_response_hook_through_real_logging_wrapper(self): - from mcp.types import CallToolResult, TextContent + async def test_mcp_response_hook_through_real_logging_wrapper(self, use_wrapper): + from mcp.types import AudioContent, CallToolResult, EmbeddedResource, ImageContent, TextContent, TextResourceContents from litellm.types.mcp import MCPPostCallResponseObject @@ -387,7 +388,14 @@ class TestCiscoAIDefenseMCPMode: ) real_result = CallToolResult( - content=[TextContent(type="text", text="leak 9045629876")], + content=[ + TextContent(type="text", text="leak 9045629876"), + ImageContent(type="image", data="aGVsbG8=", mimeType="image/png"), + AudioContent(type="audio", data="aGVsbG8=", mimeType="audio/wav"), + EmbeddedResource(type="resource", resource=TextResourceContents( + uri="memo://status", mimeType="text/plain", text="resource text" + )), + ], structuredContent={"patient": {"ssn": "123-45-6789"}}, isError=False, ) @@ -396,15 +404,6 @@ class TestCiscoAIDefenseMCPMode: hidden_params={}, ) - assert isinstance(wrapped.mcp_tool_call_response, list) - assert all( - isinstance(item, tuple) and len(item) == 2 - for item in wrapped.mcp_tool_call_response - ), ( - "Pydantic coercion shape changed — update the normalizer to " - "match the new wire format." - ) - post_mock = AsyncMock(return_value=_safe_response(url=MCP_URL)) with _patch_inspection_post(g, post_mock): result = await g.async_post_mcp_tool_call_hook( @@ -414,7 +413,7 @@ class TestCiscoAIDefenseMCPMode: "mcp_server_name": "vault", "litellm_call_id": "real-wire-call", }, - response_obj=wrapped, + response_obj=wrapped if use_wrapper else real_result, start_time=datetime.now(), end_time=datetime.now(), ) @@ -428,8 +427,8 @@ class TestCiscoAIDefenseMCPMode: sent_payload = post_mock.call_args.kwargs["json"] content_items = sent_payload["result"]["content"] - assert len(content_items) == 1, ( - f"expected exactly 1 content item from the real " + assert len(content_items) == 4, ( + f"expected exactly 4 content items from the real " f"CallToolResult.content list, got {len(content_items)}: " f"{content_items!r}" ) @@ -441,6 +440,11 @@ class TestCiscoAIDefenseMCPMode: f"``content`` field." ) assert content_items[0].get("type") == "text" + assert content_items[1:] == [ + {"type": "image", "data": "aGVsbG8=", "mimeType": "image/png"}, + {"type": "audio", "data": "aGVsbG8=", "mimeType": "audio/wav"}, + {"type": "resource", "resource": {"uri": "memo://status", "mimeType": "text/plain", "text": "resource text"}}, + ] assert sent_payload["result"]["structuredContent"] == { "patient": {"ssn": "123-45-6789"} } @@ -482,7 +486,6 @@ class TestCiscoAIDefenseMCPMode: class TestCiscoAIDefenseRedactListShape: - @staticmethod def _violation_with_redact_response(text: str = "[REDACTED tool output]"): return _mock_inspect_response( @@ -512,8 +515,8 @@ class TestCiscoAIDefenseRedactListShape: tuples_list = [ ("meta", None), ("content", inner_content), - ("structuredContent", {"patient": {"ssn": "123-45-6789"}}), - ("isError", False), + ("structured_content", {"patient": {"ssn": "123-45-6789"}}), + ("is_error", False), ] return tuples_list, lambda: inner_content[0].text @@ -526,16 +529,12 @@ class TestCiscoAIDefenseRedactListShape: from litellm.types.mcp import MCPPostCallResponseObject - g = _make_guardrail( - inspection_type="mcp", event_hook=["pre_mcp_call", "during_mcp_call"] - ) + g = _make_guardrail(inspection_type="mcp", event_hook=["pre_mcp_call", "during_mcp_call"]) content, get_text = getattr(self, factory_name)() response_obj = _mcp_response(content) - with _patch_inspection_post( - g, AsyncMock(return_value=self._violation_with_redact_response()) - ): + with _patch_inspection_post(g, AsyncMock(return_value=self._violation_with_redact_response())): result = await g.async_post_mcp_tool_call_hook( kwargs={"name": "leak", "arguments": {}}, response_obj=response_obj, @@ -544,15 +543,13 @@ class TestCiscoAIDefenseRedactListShape: ) assert result is None or not isinstance(result, MCPPostCallResponseObject), ( - f"Redact silently fell through to block for {factory_name}. " - f"result={result!r}" + f"Redact silently fell through to block for {factory_name}. result={result!r}" ) assert get_text() == "[REDACTED tool output]", ( - f"Redact silently failed for {factory_name}; original text " - f"not rewritten." + f"Redact silently failed for {factory_name}; original text not rewritten." ) if factory_name == "_pydantic_tuple_list_factory": - structured_content = dict(content)["structuredContent"] + structured_content = dict(content)["structured_content"] assert structured_content == {"result": "[REDACTED tool output]"} assert "123-45-6789" not in json.dumps(structured_content) @@ -591,12 +588,12 @@ class TestCiscoAIDefenseRedactListShape: ) assert original_response.content[0].text == "[REDACTED tool output]" - assert "123-45-6789" not in json.dumps(original_response.structuredContent), ( + assert "123-45-6789" not in json.dumps(original_response.structured_content), ( "Redact verdict left the client-visible MCP tool output unchanged. " "The post-call hook receives a wrapped MCPPostCallResponseObject but " "the endpoint returns kwargs['original_response'], so the redaction " "must rewrite that object too. structuredContent still leaks: " - f"{original_response.structuredContent!r}" + f"{original_response.structured_content!r}" ) @@ -712,11 +709,11 @@ class TestCiscoAIDefenseMCPBlockingContract: "Hook must keep returning a MCPPostCallResponseObject for " "dispatcher paths that do honor returned replacements." ) - assert raw_response.isError is True + assert raw_response.is_error is True assert "Blocked by Cisco AI Defense" in raw_response.content[0].text - assert raw_response.structuredContent is not None - assert "Blocked by Cisco AI Defense" in raw_response.structuredContent["result"] - assert "exfiltrated" not in raw_response.structuredContent["result"] + assert raw_response.structured_content is not None + assert "Blocked by Cisco AI Defense" in raw_response.structured_content["result"] + assert "exfiltrated" not in raw_response.structured_content["result"] logging_stub = Logging.__new__(Logging) logging_stub.model_call_details = {} parsed = logging_stub._parse_post_mcp_call_hook_response(response=result) diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index fc3ede88aa9..baaf3f4ba2f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -643,6 +643,24 @@ def test_key_metadata_includes_recovered_user_email(): assert meta.user_email == "alice@example.com" +def test_key_metadata_includes_user_id_without_user_email(): + from litellm.proxy.management_endpoints.common_daily_activity import _key_metadata + + meta = _key_metadata( + { + "dirty-key": { + "key_alias": "batch-worker", + "team_id": "team-1", + "user_id": "user-123", + } + }, + "dirty-key", + ) + + assert meta.user_id == "user-123" + assert meta.user_email is None + + def test_update_breakdown_metrics_includes_user_email(): from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics @@ -913,6 +931,67 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys(): assert key_data.metrics.spend == 10.0 +@pytest.mark.asyncio +async def test_aggregated_activity_flags_only_keys_that_key_info_can_still_resolve(): + """/key/info reads the active key table only, so deleted and never-stored (session) keys must not claim to exist.""" + mock_prisma = MagicMock() + base = { + "date": "2024-01-01", + "endpoint": "/v1/chat/completions", + "model": None, + "model_group": None, + "custom_llm_provider": None, + "mcp_namespaced_tool_name": None, + "group_level": 30, + "distinct_api_keys": 1, + "spend": 1.0, + "prompt_tokens": 10, + "completion_tokens": 5, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + "compression_saved_tokens": 0, + "compression_savings_spend": 0.0, + "prompt_caching_savings_spend": 0.0, + "gateway_injected_caching_savings_spend": 0.0, + "autorouter_savings_spend": 0.0, + "total_response_time_ms": 0, + "timed_requests": 0, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + mock_prisma.db.query_raw = AsyncMock( + return_value=[{**base, "api_key": key} for key in ("active-key", "deleted-key", "session-key")] + ) + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[SimpleNamespace(token="active-key", key_alias="active", team_id=None, user_id="owner")] + ) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock( + return_value=[SimpleNamespace(token="deleted-key", key_alias="deleted", team_id=None, user_id="owner")] + ) + mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + + result = await get_daily_activity_aggregated( + prisma_client=mock_prisma, + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id=None, + entity_metadata_field=None, + start_date="2024-01-01", + end_date="2024-01-01", + model=None, + api_key=None, + ) + + key_breakdown = result.results[0].breakdown.endpoints["/v1/chat/completions"].api_key_breakdown + assert {key: data.metadata.key_exists for key, data in key_breakdown.items()} == { + "active-key": True, + "deleted-key": False, + "session-key": False, + } + assert key_breakdown["deleted-key"].metadata.key_alias == "deleted" + + def _daily_user_spend_record(*, user_id, api_key, spend, model="gpt-4", model_group="gpt-4"): """A LiteLLM_DailyUserSpend row as the per-user breakdown reads it.""" return SimpleNamespace( diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index f53fbde6b51..f7abb209015 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -121,6 +121,137 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( assert "error" in events[-1] +@pytest.mark.asyncio +async def test_responses_api_background_polling_rejects_missing_input(): + from fastapi import Response as FastAPIResponse + from starlette.requests import Request + + from litellm.proxy._types import ProxyException, UserAPIKeyAuth + from litellm.proxy.response_api_endpoints.endpoints import responses_api + + processor = MagicMock() + + async def return_exception(*, e: Exception, **kwargs: object) -> Exception: + return e + + processor._handle_llm_api_exception = AsyncMock(side_effect=return_exception) + processor.common_processing_pre_call_logic = AsyncMock(return_value=({"model": "gpt-4o"}, MagicMock())) + + async def receive(): + return { + "type": "http.request", + "body": b'{"model":"gpt-4o","background":true}', + "more_body": False, + } + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/responses", + "headers": [(b"content-type", b"application/json")], + }, + receive, + ) + + with ( + patch( # test-quality-ok: endpoint constructs the processor directly + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + return_value=processor, + ), + patch( # test-quality-ok: polling decision is imported inside the endpoint + "litellm.proxy.response_polling.polling_handler.should_use_polling_for_request", + return_value=True, + ), + patch( # test-quality-ok: background task is imported inside the endpoint + "litellm.proxy.response_polling.background_streaming.background_streaming_task", + new_callable=AsyncMock, + ) as mock_background_streaming_task, + patch( # test-quality-ok: polling handler is imported inside the endpoint + "litellm.proxy.response_polling.polling_handler.ResponsePollingHandler.create_initial_state", + new_callable=AsyncMock, + ) as mock_create_initial_state, + ): + with pytest.raises(ProxyException) as exc_info: + await responses_api( + request=request, + fastapi_response=FastAPIResponse(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert exc_info.value.code == "400" + assert exc_info.value.param == "input" + processor.common_processing_pre_call_logic.assert_awaited_once() + mock_background_streaming_task.assert_not_called() + mock_create_initial_state.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_responses_api_background_polling_accepts_input_from_prompt_template(): + from fastapi import Response as FastAPIResponse + from starlette.requests import Request + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.response_api_endpoints.endpoints import responses_api + + processor = MagicMock() + processor.common_processing_pre_call_logic = AsyncMock( + return_value=({"model": "gpt-4o", "input": "hello from prompt"}, MagicMock()) + ) + initial_state = MagicMock() + + async def receive(): + return { + "type": "http.request", + "body": b'{"model":"gpt-4o","prompt_id":"greeting","background":true}', + "more_body": False, + } + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/responses", + "headers": [(b"content-type", b"application/json")], + }, + receive, + ) + + with ( + patch( # test-quality-ok: endpoint constructs the processor directly + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + return_value=processor, + ), + patch( # test-quality-ok: polling decision is imported inside the endpoint + "litellm.proxy.response_polling.polling_handler.should_use_polling_for_request", + return_value=True, + ), + patch( # test-quality-ok: background task is imported inside the endpoint + "litellm.proxy.response_polling.background_streaming.background_streaming_task", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: avoid scheduling a background task in this unit test + "litellm.proxy.response_api_endpoints.endpoints.asyncio.create_task", + ), + patch( # test-quality-ok: polling handler is imported inside the endpoint + "litellm.proxy.response_polling.polling_handler.ResponsePollingHandler.create_initial_state", + new_callable=AsyncMock, + ) as mock_create_initial_state, + ): + mock_create_initial_state.return_value = initial_state + result = await responses_api( + request=request, + fastapi_response=FastAPIResponse(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert result is initial_state + processor.common_processing_pre_call_logic.assert_awaited_once() + mock_create_initial_state.assert_awaited_once() + request_data = mock_create_initial_state.await_args.kwargs["request_data"] + assert request_data["input"] == "hello from prompt" + + class TestResponsesAPIEndpoints(unittest.TestCase): @pytest.mark.asyncio @patch("litellm.proxy.proxy_server.llm_router") diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 811d11bf0bf..935cc6ad8b7 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5516,6 +5516,37 @@ async def test_router_settings_reload_keeps_db_values_writable(tmp_path, monkeyp assert proxy_config.router_settings.rejected_writes({"disable_cooldowns": False}) == ("disable_cooldowns",) +@pytest.mark.asyncio +async def test_boot_warns_that_a_shadowed_database_value_will_never_apply(tmp_path, monkeypatch, caplog): + from litellm.proxy.proxy_server import ProxyConfig + + config_path: Final = tmp_path / "config.yaml" + config_path.write_text( + yaml.safe_dump({"model_list": [], "general_settings": {"allowed_ips": ["1.2.3.4"], "max_file_size_mb": 5}}) + ) + db_row: Final = types.SimpleNamespace(param_value={"allowed_ips": ["1.2.3.4", "5.6.7.8"], "max_parallel_requests": 7}) + + async def read_config_row(_prisma_client, param_name): + return db_row if param_name == "general_settings" else None + + monkeypatch.setattr(proxy_server_module, "get_config_param", read_config_row) + monkeypatch.setattr(proxy_server_module, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server_module, "store_model_in_db", True) + monkeypatch.setattr(proxy_server_module, "user_config_file_path", None) + proxy_config: Final = ProxyConfig() + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await proxy_config.get_config(config_file_path=str(config_path)) + + warnings: Final = " ".join(record.getMessage() for record in caplog.records) + assert "allowed_ips" in warnings + assert "ignored" in warnings + assert "max_parallel_requests" not in warnings + assert "max_file_size_mb" not in warnings + assert proxy_config.settings["allowed_ips"] == ["1.2.3.4"] + assert proxy_config.settings["max_parallel_requests"] == 7 + + @pytest.mark.asyncio async def test_model_info_v1_oci_secrets_not_leaked(): """ diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index f7021763a4d..6cbbc279748 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -1043,6 +1043,7 @@ async def test_route_request_override_enable_tag_filtering_beats_body_value(): [ ("acompletion", "messages", "/chat/completions"), ("aembedding", "input", "/embeddings"), + ("aresponses", "input", "/responses"), ("acreate_batch", "input_file_id", "/batches"), ], ) @@ -1090,6 +1091,8 @@ def test_raise_if_required_body_param_missing_names_first_missing_batch_param(da ("acompletion", {"model": "gpt-4o", "messages": []}), ("atext_completion", {"model": "gpt-4o"}), ("aembedding", {"model": "text-embedding-3-small", "input": "hi"}), + ("aresponses", {"model": "gpt-4o", "input": "hi"}), + ("aresponses", {"model": "gpt-4o", "input": []}), ("arerank", {"model": "rerank-model"}), ("aimage_generation", {"model": "dall-e-3"}), ( @@ -1120,6 +1123,20 @@ async def test_route_request_rejects_chat_completion_without_messages(): llm_router.acompletion.assert_not_called() +@pytest.mark.asyncio +async def test_route_request_rejects_responses_without_input(): + from litellm.proxy.route_llm_request import ProxyMissingRequiredParamError + + llm_router = MagicMock() + + with pytest.raises(ProxyMissingRequiredParamError) as exc_info: + await route_request({"model": "gpt-4o"}, llm_router, None, "aresponses") + + assert exc_info.value.code == "400" + assert exc_info.value.param == "input" + llm_router.aresponses.assert_not_called() + + class FakeProxyModelTable: def __init__(self, rows): self.rows = rows diff --git a/tests/test_litellm/rust_bridge/test_catalog.py b/tests/test_litellm/rust_bridge/test_catalog.py index e9fdbf859f4..147e863baf5 100644 --- a/tests/test_litellm/rust_bridge/test_catalog.py +++ b/tests/test_litellm/rust_bridge/test_catalog.py @@ -85,3 +85,15 @@ def test_first_matching_rule_respects_every_constraint(context: Context, expecte ) assert catalog.decision(context, rules) is expected + + +@pytest.mark.parametrize("process", (None, False, True)) +@pytest.mark.parametrize("environment", (None, "0", "1")) +def test_textract_ocr_has_no_python_path_to_opt_out_to( + monkeypatch: pytest.MonkeyPatch, process: bool | None, environment: str | None +) -> None: + configuration.rust(process) + if environment is not None: + monkeypatch.setenv("LITELLM_RUST", environment) + + assert catalog.decision(Context(Route.OCR, provider="aws_textract", model="m")) is Decision.RUST_REQUIRED diff --git a/tests/test_litellm/test_auto_merge_price_sync.py b/tests/test_litellm/test_auto_merge_price_sync.py deleted file mode 100644 index 3e8c0dc024c..00000000000 --- a/tests/test_litellm/test_auto_merge_price_sync.py +++ /dev/null @@ -1,219 +0,0 @@ -"""Tests for .github/scripts/auto_merge_price_sync.py. - -`evaluate` is pure: it takes the pull request plus the fetched facts and -returns a Verdict, so each gate is exercised by building inputs where exactly -one condition fails and asserting the matching hold reason. A merge verdict -is the thing that spends an unreviewed merge, so the defaults below are the -happy path that every case perturbs one part of. -""" - -import importlib.util -import sys -from datetime import datetime, timezone -from pathlib import Path -from typing import Final - -import pytest - -_REPO_ROOT = Path(__file__).resolve().parents[2] -_MODULE_PATH = _REPO_ROOT / ".github" / "scripts" / "auto_merge_price_sync.py" -_spec = importlib.util.spec_from_file_location("auto_merge_price_sync", _MODULE_PATH) -merger = importlib.util.module_from_spec(_spec) -sys.modules[_spec.name] = merger -_spec.loader.exec_module(merger) - -HEAD_SHA: Final = "deadbeef" * 5 -ALLOWLIST: Final = frozenset({"berriai-litellm-provider-info-sync[bot]"}) -COST_MAP_FILES: Final = ("model_prices_and_context_window.json",) - - -def _pr(**overrides: object) -> merger.PullRequest: - base: Final = { - "number": 1, - "title": "sync prices", - "author_login": "berriai-litellm-provider-info-sync[bot]", - "state": "open", - "draft": False, - "mergeable": True, - "mergeable_state": "clean", - "head_sha": HEAD_SHA, - } - return merger.PullRequest(**{**base, **overrides}) - - -def _inputs(**overrides: object) -> merger.EvaluationInputs: - base: Final = { - "pr": _pr(), - "changed_files": COST_MAP_FILES, - "required_contexts": frozenset({"build"}), - "check_runs": (merger.CheckRun(name="build", status="completed", conclusion="success"),), - "statuses": (), - "reviews": (), - "self_check_name": "auto-merge-price-sync", - "author_allowlist": ALLOWLIST, - } - return merger.EvaluationInputs(**{**base, **overrides}) - - -def _evaluate(inputs: merger.EvaluationInputs) -> merger.Verdict: - return merger.evaluate(inputs, classify=lambda files: "run") - - -def _holds(inputs: merger.EvaluationInputs, fragment: str) -> merger.Verdict: - verdict: Final = _evaluate(inputs) - assert not verdict.merge - assert any(fragment in reason for reason in verdict.reasons), verdict.reasons - return verdict - - -def test_happy_path_merges() -> None: - verdict: Final = _evaluate(_inputs()) - assert verdict.merge - assert verdict.reasons == () - - -def test_non_allowlisted_author_holds() -> None: - _holds(_inputs(pr=_pr(author_login="octocat")), "not in allowlist") - - -def test_closed_pr_holds() -> None: - _holds(_inputs(pr=_pr(state="closed")), "pr not open") - - -def test_draft_pr_holds() -> None: - _holds(_inputs(pr=_pr(draft=True)), "draft") - - -def test_unmergeable_pr_holds() -> None: - _holds(_inputs(pr=_pr(mergeable=False)), "not mergeable") - - -def test_dirty_pr_holds() -> None: - _holds(_inputs(pr=_pr(mergeable_state="dirty")), "merge conflicts") - - -def test_non_cost_map_files_hold() -> None: - verdict: Final = merger.evaluate(_inputs(changed_files=("litellm/utils.py",)), classify=lambda files: "skip") - assert not verdict.merge - assert any("cost-map-only" in reason for reason in verdict.reasons) - - -def test_required_context_missing_holds() -> None: - _holds(_inputs(check_runs=()), "required check 'build' not green") - - -def test_required_context_via_commit_status_passes() -> None: - verdict: Final = _evaluate( - _inputs( - check_runs=(), - statuses=(merger.CommitStatus(context="build", state="success"),), - ) - ) - assert verdict.merge - - -def test_failing_check_run_holds() -> None: - _holds( - _inputs( - check_runs=( - merger.CheckRun(name="build", status="completed", conclusion="success"), - merger.CheckRun(name="lint", status="completed", conclusion="failure"), - ) - ), - "check run 'lint' is completed/failure", - ) - - -def test_in_progress_check_run_holds() -> None: - _holds( - _inputs( - check_runs=( - merger.CheckRun(name="build", status="completed", conclusion="success"), - merger.CheckRun(name="ui", status="in_progress", conclusion=None), - ) - ), - "check run 'ui'", - ) - - -def test_own_check_run_is_ignored() -> None: - verdict: Final = _evaluate( - _inputs( - check_runs=( - merger.CheckRun(name="build", status="completed", conclusion="success"), - merger.CheckRun(name="auto-merge-price-sync", status="in_progress", conclusion=None), - ) - ) - ) - assert verdict.merge - - -def test_pending_commit_status_holds() -> None: - _holds( - _inputs(statuses=(merger.CommitStatus(context="codecov", state="pending"),)), - "commit status 'codecov' is pending", - ) - - -def test_changes_requested_holds() -> None: - _holds( - _inputs( - reviews=( - merger.Review( - author_login="human-reviewer", - state="CHANGES_REQUESTED", - body="", - commit_id=HEAD_SHA, - submitted_at=datetime(2026, 1, 12, tzinfo=timezone.utc), - ), - ) - ), - "changes requested by human-reviewer", - ) - - -def test_superseded_changes_requested_merges() -> None: - verdict: Final = _evaluate( - _inputs( - reviews=( - merger.Review( - author_login="human-reviewer", - state="CHANGES_REQUESTED", - body="", - commit_id=HEAD_SHA, - submitted_at=datetime(2026, 1, 11, tzinfo=timezone.utc), - ), - merger.Review( - author_login="human-reviewer", - state="APPROVED", - body="", - commit_id=HEAD_SHA, - submitted_at=datetime(2026, 1, 13, tzinfo=timezone.utc), - ), - ) - ) - ) - assert verdict.merge - - -def test_merge_request_pins_evaluated_head_sha() -> None: - body: Final = merger.merge_request_body(_pr(number=7, title="sync prices")) - assert body["sha"] == HEAD_SHA - assert body["merge_method"] == "merge" - assert body["commit_title"] == "sync prices (#7)" - - -def test_classifier_cost_map_set_runs() -> None: - assert merger._classify(["model_prices_and_context_window.json", "tests/test_litellm/test_x.py"]) == "run" - - -def test_classifier_backend_file_skips() -> None: - assert merger._classify(["model_prices_and_context_window.json", "litellm/main.py"]) == "skip" - - -def test_main_without_token_logs_and_exits_zero( - monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] -) -> None: - monkeypatch.delenv("GH_TOKEN", raising=False) - assert merger.main() == 0 - assert "app credentials not configured" in capsys.readouterr().out diff --git a/tests/test_litellm/test_circleci_path_filter.py b/tests/test_litellm/test_circleci_path_filter.py index 07fab42bad6..84e2327057d 100644 --- a/tests/test_litellm/test_circleci_path_filter.py +++ b/tests/test_litellm/test_circleci_path_filter.py @@ -49,6 +49,16 @@ CI = [".github/workflows/test-litellm-ui-unit.yml"] @pytest.mark.parametrize( "category,changed,expected", [ + ("mcp-dependencies", ["pyproject.toml"], "run"), + ("mcp-dependencies", ["uv.lock"], "run"), + ("mcp-dependencies", ["litellm/experimental_mcp_client/client.py"], "run"), + ("mcp-dependencies", ["tests/e2e/mcp/oauth_chat_client.py"], "run"), + ("mcp-dependencies", ["litellm-proxy-extras/pyproject.toml"], "run"), + ("mcp-dependencies", ["scripts/check_mcp_sdk_install.py"], "run"), + ("mcp-dependencies", [".github/workflows/test-mcp-dependency-resolution.yml"], "run"), + ("mcp-dependencies", [".circleci/scripts/classify_changes.sh"], "run"), + ("mcp-dependencies", ["litellm/llms/openai/chat/gpt_transformation.py"], "skip"), + ("mcp-dependencies", DOCS + CLIENT, "skip"), ("provider-harness", ["tests/e2e/provider_cache.py"], "run"), ("provider-harness", ["tests/e2e/conftest.py"], "run"), ("provider-harness", ["tests/e2e/e2e_http.py"], "run"), diff --git a/tests/test_litellm/test_router/test_enforce_model_rate_limits.py b/tests/test_litellm/test_router/test_enforce_model_rate_limits.py index 1def253ac93..7577064b7f9 100644 --- a/tests/test_litellm/test_router/test_enforce_model_rate_limits.py +++ b/tests/test_litellm/test_router/test_enforce_model_rate_limits.py @@ -6,6 +6,7 @@ regardless of the routing strategy being used. """ import asyncio +from datetime import timedelta from unittest.mock import AsyncMock, MagicMock import pytest @@ -13,10 +14,28 @@ import pytest import litellm from litellm import Router from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) +TPM_DEPLOYMENT = { + "tpm": 1000, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "replica-test-id"}, + "model_name": "test-model", +} + + +def _dual_cache_with_local_tpm(local_tpm: int, redis_cache: MagicMock | None) -> DualCache: + dual_cache = DualCache(redis_cache=redis_cache) + check = ModelRateLimitingCheck(dual_cache=dual_cache) + now = litellm.utils.get_utc_datetime() + for minute in (now, now + timedelta(minutes=1)): + tpm_key, _ = check._get_cache_keys(TPM_DEPLOYMENT, minute.strftime("%H-%M")) + dual_cache.set_cache(key=tpm_key, value=local_tpm, local_only=True) + return dual_cache + class TestModelRateLimitingCheck: """Test the ModelRateLimitingCheck class directly.""" @@ -144,6 +163,50 @@ class TestModelRateLimitingCheck: assert "TPM limit=1000" in str(exc_info.value) assert "current usage=1000" in str(exc_info.value) + def test_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): + redis_cache = MagicMock() + redis_cache.get_cache.return_value = 1000 + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(TPM_DEPLOYMENT) + + assert "current usage=1000" in str(exc_info.value) + + @pytest.mark.parametrize( + "redis_get", [MagicMock(return_value=None), MagicMock(side_effect=RedisCircuitBreakerOpenError())] + ) + def test_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): + redis_cache = MagicMock() + redis_cache.get_cache = redis_get + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(TPM_DEPLOYMENT) + + assert "current usage=1000" in str(exc_info.value) + + def test_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open(self): + redis_cache = MagicMock() + redis_cache.get_cache.side_effect = RedisCircuitBreakerOpenError() + redis_cache.increment_cache.return_value = 2 + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) + + assert "RPM limit=1" in str(exc_info.value) + + def test_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) + deployment = {**TPM_DEPLOYMENT, "rpm": 1} + + assert check.pre_call_check(deployment) == deployment + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(deployment) + + assert "RPM limit=1" in str(exc_info.value) + def test_log_success_event_increments_cache(self): """Test that log_success_event correctly increments the cache.""" mock_cache = MagicMock() @@ -245,6 +308,56 @@ class TestModelRateLimitingCheckAsync: assert "TPM limit=1000" in str(exc_info.value) + @pytest.mark.asyncio + async def test_async_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): + redis_cache = MagicMock() + redis_cache.async_get_cache = AsyncMock(return_value=1000) + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(TPM_DEPLOYMENT) + + assert "current usage=1000" in str(exc_info.value) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "redis_get", [AsyncMock(return_value=None), AsyncMock(side_effect=RedisCircuitBreakerOpenError())] + ) + async def test_async_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): + redis_cache = MagicMock() + redis_cache.async_get_cache = redis_get + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(TPM_DEPLOYMENT) + + assert "current usage=1000" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open( + self, + ): + redis_cache = MagicMock() + redis_cache.async_get_cache = AsyncMock(side_effect=RedisCircuitBreakerOpenError()) + redis_cache.async_increment = AsyncMock(return_value=2) + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) + + assert "RPM limit=1" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) + deployment = {**TPM_DEPLOYMENT, "rpm": 1} + + assert await check.async_pre_call_check(deployment) == deployment + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "RPM limit=1" in str(exc_info.value) + @pytest.mark.asyncio async def test_async_log_success_event_increments_cache(self): """Test that async_log_success_event correctly increments the cache.""" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx index 5846a63bc70..6bd16095351 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx @@ -5,7 +5,39 @@ import type { ReactNode } from "react"; import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; import useTeams from "@/app/(dashboard)/hooks/useTeams"; import * as networking from "@/components/networking"; +import type { DailyData, KeyMetadata, KeyMetricWithMetadata, SpendMetrics } from "@/components/UsagePage/types"; import EntityUsage from "./EntityUsage"; +import { getGlobalTopKeys, getTopAPIKeys } from "./entityUsageAggregations"; + +const emptySpendMetrics: SpendMetrics = { + spend: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + api_requests: 0, + successful_requests: 0, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, +}; + +const createKeyMetrics = (spend: number, metadata: KeyMetadata): KeyMetricWithMetadata => ({ + metrics: { ...emptySpendMetrics, spend }, + metadata, +}); + +const createDailyData = (date: string, apiKeys: Record): DailyData => ({ + date, + metrics: { ...emptySpendMetrics }, + breakdown: { + models: {}, + model_groups: {}, + mcp_servers: {}, + providers: {}, + api_keys: apiKeys, + entities: {}, + }, +}); beforeAll(() => { if (typeof window !== "undefined" && !window.ResizeObserver) { @@ -44,10 +76,10 @@ vi.mock("../EndpointUsage/EndpointUsage", () => ({ })); vi.mock("@/components/UsagePage/components/EntityUsage/TopKeyView", () => ({ - default: ({ topKeys }: { topKeys: { api_key: string; spend: number }[] }) => ( + default: ({ topKeys }: { topKeys: { api_key: string; user?: string | null; spend: number }[] }) => (
Top Keys - {`top-keys:${topKeys.map((row) => `${row.api_key}=${row.spend}`).join("|")}`} + {`top-keys:${topKeys.map((row) => `${row.api_key}=${row.spend}=${row.user ?? "-"}`).join("|")}`}
), })); @@ -431,6 +463,55 @@ describe("EntityUsage", () => { ); }); + describe("top key aggregations", () => { + it("sums, sorts, limits, and carries email attribution for global top keys", () => { + const results = [ + createDailyData("2025-01-01", { + "key-low": createKeyMetrics(10, { key_alias: "Low", team_id: null, user_email: "low@example.com" }), + "key-high": createKeyMetrics(25, { key_alias: "High", team_id: null, user_email: "high@example.com" }), + }), + createDailyData("2025-01-02", { + "key-low": createKeyMetrics(30, { key_alias: "Low", team_id: null, user_email: "low@example.com" }), + }), + ]; + + expect(getGlobalTopKeys(results, 1)).toEqual([ + { + api_key: "key-low", + key_alias: "Low", + user: "low@example.com", + tags: [], + spend: 40, + }, + ]); + }); + + it("falls back to user ID attribution for global and entity top keys", () => { + const results = [ + createDailyData("2025-01-01", { + "key-123": createKeyMetrics(12.5, { key_alias: "User ID key", team_id: null, user_id: "user-123" }), + }), + ]; + + expect(getGlobalTopKeys(results, 5)[0]?.user).toBe("user-123"); + expect(getTopAPIKeys(results, 5)[0]?.user).toBe("user-123"); + }); + + it("carries whether each key still exists for global and entity top keys", () => { + const results = [ + createDailyData("2025-01-01", { + "stored-key": createKeyMetrics(20, { key_alias: "Stored", team_id: null, key_exists: true }), + "session-key": createKeyMetrics(10, { key_alias: null, team_id: null, key_exists: false }), + }), + ]; + const existsByKey = (rows: { api_key: string; key_exists?: boolean | null }[]) => + Object.fromEntries(rows.map((row) => [row.api_key, row.key_exists])); + + expect(existsByKey(getGlobalTopKeys(results, 5))).toEqual({ "stored-key": true, "session-key": false }); + expect(existsByKey(getTopAPIKeys(results, 5))).toEqual({ "stored-key": true, "session-key": false }); + }); + }); + it("should render with tag entity type and display spend metrics", async () => { render(); @@ -1099,7 +1180,12 @@ describe("EntityUsage", () => { breakdown: { ...mockSpendData.results[0].breakdown, model_groups: { "gpt-4o": { metrics: { ...usageMetrics, spend: 70.25 }, metadata: {} } }, - api_keys: { "sk-abc": { metrics: usageMetrics, metadata: { key_alias: "prod-key", team_id: null } } }, + api_keys: { + "sk-abc": { + metrics: usageMetrics, + metadata: { key_alias: "prod-key", team_id: null, user_email: "alice@example.com" }, + }, + }, }, }, ], @@ -1108,7 +1194,7 @@ describe("EntityUsage", () => { render(); await waitFor(() => { - expect(screen.getByText("top-keys:sk-abc=30.75")).toBeInTheDocument(); + expect(screen.getByText("top-keys:sk-abc=30.75=alice@example.com")).toBeInTheDocument(); }); expect(screen.getByText("top-models:gpt-4o=70.25")).toBeInTheDocument(); expect(screen.getByText(/^top-models:Code Review Agent=/)).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts index d482a5576ae..54569608b0e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts @@ -1,4 +1,5 @@ import { keyActivityLabel } from "@/components/UsagePage/keyActivityLabel"; +import type { TopKeyItem } from "@/components/UsagePage/components/EntityUsage/TopKeyView"; import { BreakdownMetrics, DailyData, KeyMetricWithMetadata, TagUsage } from "@/components/UsagePage/types"; export type ExtendedDailyData = DailyData & { @@ -85,7 +86,59 @@ export const getTopAgents = (results: ExtendedDailyData[], topAgentsLimit: numbe .slice(0, topAgentsLimit); }; -export const getTopAPIKeys = (results: ExtendedDailyData[], topKeysLimit: number) => { +export const getGlobalTopKeys = (results: DailyData[], topKeysLimit: number): TopKeyItem[] => { + const keySpend: { [key: string]: KeyMetricWithMetadata } = {}; + results.forEach((day) => { + Object.entries(day.breakdown.api_keys || {}).forEach(([key, metrics]) => { + if (!keySpend[key]) { + keySpend[key] = { + metrics: { + spend: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + api_requests: 0, + successful_requests: 0, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + metadata: { + key_alias: metrics.metadata.key_alias, + team_id: null, + user_id: metrics.metadata.user_id, + user_email: metrics.metadata.user_email, + key_exists: metrics.metadata.key_exists, + tags: metrics.metadata.tags || [], + }, + }; + } + keySpend[key].metrics.spend += metrics.metrics.spend; + keySpend[key].metrics.prompt_tokens += metrics.metrics.prompt_tokens; + keySpend[key].metrics.completion_tokens += metrics.metrics.completion_tokens; + keySpend[key].metrics.total_tokens += metrics.metrics.total_tokens; + keySpend[key].metrics.api_requests += metrics.metrics.api_requests; + keySpend[key].metrics.successful_requests += metrics.metrics.successful_requests; + keySpend[key].metrics.failed_requests += metrics.metrics.failed_requests; + keySpend[key].metrics.cache_read_input_tokens += metrics.metrics.cache_read_input_tokens || 0; + keySpend[key].metrics.cache_creation_input_tokens += metrics.metrics.cache_creation_input_tokens || 0; + }); + }); + + return Object.entries(keySpend) + .map(([api_key, metrics]) => ({ + api_key, + key_alias: keyActivityLabel(metrics.metadata), + user: metrics.metadata.user_email ?? metrics.metadata.user_id ?? null, + key_exists: metrics.metadata.key_exists, + tags: metrics.metadata.tags || [], + spend: metrics.metrics.spend, + })) + .sort((a, b) => b.spend - a.spend) + .slice(0, topKeysLimit); +}; + +export const getTopAPIKeys = (results: ExtendedDailyData[], topKeysLimit: number): TopKeyItem[] => { const keySpend: { [key: string]: KeyMetricWithMetadata } = {}; results.forEach((day) => { const { breakdown } = day; @@ -119,7 +172,9 @@ export const getTopAPIKeys = (results: ExtendedDailyData[], topKeysLimit: number metadata: { key_alias: metrics.metadata.key_alias, team_id: metrics.metadata.team_id || null, + user_id: metrics.metadata.user_id, user_email: metrics.metadata.user_email, + key_exists: metrics.metadata.key_exists, tags: tagDictionary[key] || [], }, }; @@ -140,7 +195,9 @@ export const getTopAPIKeys = (results: ExtendedDailyData[], topKeysLimit: number .map(([api_key, metrics]) => ({ api_key, key_alias: keyActivityLabel(metrics.metadata), - tags: metrics.metadata.tags || "-", + user: metrics.metadata.user_email ?? metrics.metadata.user_id ?? null, + key_exists: metrics.metadata.key_exists, + tags: metrics.metadata.tags || [], spend: metrics.metrics.spend, })) .sort((a, b) => b.spend - a.spend) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index a9ab0f17f40..228d8acf146 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -46,8 +46,7 @@ import { Tag } from "@/components/tag_management/types"; import UserAgentActivity from "@/components/user_agent_activity"; import ViewUserSpend from "@/components/view_user_spend"; import { usePaginatedDailyActivity } from "../hooks/usePaginatedDailyActivity"; -import { keyActivityLabel } from "@/components/UsagePage/keyActivityLabel"; -import { DailyData, KeyMetricWithMetadata, MetricWithMetadata } from "@/components/UsagePage/types"; +import { DailyData, MetricWithMetadata } from "@/components/UsagePage/types"; import { valueFormatterSpend } from "@/components/UsagePage/utils/value_formatters"; import { fetchedRangeKey, @@ -63,7 +62,8 @@ import EntityUsage, { EntityList } from "./EntityUsage/EntityUsage"; import ModelViewToggle, { ModelViewType } from "./ModelViewToggle"; import SpendByProvider from "./EntityUsage/SpendByProvider"; import { TOP_MODEL_LIMITS } from "./EntityUsage/TopModelView"; -import TopKeyView from "@/components/UsagePage/components/EntityUsage/TopKeyView"; +import TopKeyView, { type TopKeyItem } from "@/components/UsagePage/components/EntityUsage/TopKeyView"; +import { getGlobalTopKeys } from "./EntityUsage/entityUsageAggregations"; import UsageAIChatPanel from "./UsageAIChatPanel"; import { UsageOption, UsageViewSelect } from "./UsageViewSelect/UsageViewSelect"; @@ -422,53 +422,10 @@ const UsagePage: React.FC = ({ teams, organizations }) => { }, [userSpendData.results]); // Calculate top API keys from the breakdown data - const topKeys = useMemo(() => { - const keySpend: { [key: string]: KeyMetricWithMetadata } = {}; - userSpendData.results.forEach((day) => { - Object.entries(day.breakdown.api_keys || {}).forEach(([key, metrics]) => { - if (!keySpend[key]) { - keySpend[key] = { - metrics: { - spend: 0, - prompt_tokens: 0, - completion_tokens: 0, - total_tokens: 0, - api_requests: 0, - successful_requests: 0, - failed_requests: 0, - cache_read_input_tokens: 0, - cache_creation_input_tokens: 0, - }, - metadata: { - key_alias: metrics.metadata.key_alias, - team_id: null, - user_email: metrics.metadata.user_email, - tags: metrics.metadata.tags || [], - }, - }; - } - keySpend[key].metrics.spend += metrics.metrics.spend; - keySpend[key].metrics.prompt_tokens += metrics.metrics.prompt_tokens; - keySpend[key].metrics.completion_tokens += metrics.metrics.completion_tokens; - keySpend[key].metrics.total_tokens += metrics.metrics.total_tokens; - keySpend[key].metrics.api_requests += metrics.metrics.api_requests; - keySpend[key].metrics.successful_requests += metrics.metrics.successful_requests; - keySpend[key].metrics.failed_requests += metrics.metrics.failed_requests; - keySpend[key].metrics.cache_read_input_tokens += metrics.metrics.cache_read_input_tokens || 0; - keySpend[key].metrics.cache_creation_input_tokens += metrics.metrics.cache_creation_input_tokens || 0; - }); - }); - - return Object.entries(keySpend) - .map(([api_key, metrics]) => ({ - api_key, - key_alias: keyActivityLabel(metrics.metadata), - tags: metrics.metadata.tags || [], - spend: metrics.metrics.spend, - })) - .sort((a, b) => b.spend - a.spend) - .slice(0, topKeysLimit); - }, [userSpendData.results, topKeysLimit]); + const topKeys = useMemo( + () => getGlobalTopKeys(userSpendData.results, topKeysLimit), + [userSpendData.results, topKeysLimit], + ); const sortedDailyResults = useMemo( () => [...userSpendData.results].sort((a, b) => new Date(a.date).getTime() - new Date(b.date).getTime()), diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx index c2837cf412e..e65094c5228 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx @@ -29,6 +29,8 @@ vi.mock("../../../templates/key_info_view", () => ({ ), })); +const chartBars = (container: HTMLElement) => Array.from(container.querySelectorAll("path.recharts-rectangle")); + describe("TopKeyView", () => { const mockUseAuthorized = vi.mocked(useAuthorized); const mockKeyInfoV1Call = vi.mocked(networking.keyInfoV1Call); @@ -102,6 +104,7 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "Test Key", + user: null, spend: 100, tags: [ { tag: "tag-1", usage: 50 }, @@ -118,6 +121,60 @@ describe("TopKeyView", () => { expect(screen.getByText("$100.00")).toBeInTheDocument(); }); + it("should render User column only when a row has user attribution", () => { + const { rerender } = render( + , + ); + + expect(screen.queryByText("User")).not.toBeInTheDocument(); + + rerender( + , + ); + + expect(screen.getByText("User")).toBeInTheDocument(); + expect(screen.getByText("alice@example.com")).toBeInTheDocument(); + }); + + it("should render a user ID in the User column", () => { + render( + , + ); + + expect(screen.getByText("User")).toBeInTheDocument(); + expect(screen.getByText("user-123")).toBeInTheDocument(); + }); + it("should switch to chart view when chart view button is clicked", async () => { const user = userEvent.setup(); render(); @@ -142,6 +199,7 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "A Very Long Key Alias", + user: null, spend: 100, }, ]} @@ -150,7 +208,7 @@ describe("TopKeyView", () => { await user.click(screen.getByRole("button", { name: "Chart View" })); - const bars = container.querySelectorAll("path.recharts-rectangle"); + const bars = chartBars(container); expect(bars).toHaveLength(1); expect(bars[0]).toHaveAttribute("fill", "var(--color-cyan-500, #06b6d4)"); expect(screen.getAllByText("A Very Lon...").length).toBeGreaterThan(0); @@ -197,6 +255,7 @@ describe("TopKeyView", () => { { api_key: "sk-1234567890abcdef", key_alias: "Test Key", + user: null, spend: 100, }, ]} @@ -215,12 +274,13 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "", + user: null, spend: 100, }, ]} />, ); - expect(screen.getByText("-")).toBeInTheDocument(); + expect(screen.getAllByText("-")).toHaveLength(1); }); it("should format spend values with two decimal places", () => { @@ -231,6 +291,7 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "Test Key", + user: null, spend: 123.456, }, ]} @@ -247,6 +308,7 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "Test Key", + user: null, spend: 0.004, }, ]} @@ -263,12 +325,13 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "Test Key", + user: null, spend: 0, }, ]} />, ); - expect(screen.getByText("-")).toBeInTheDocument(); + expect(screen.getAllByText("-")).toHaveLength(1); expect(screen.queryByText("$0.00")).not.toBeInTheDocument(); }); @@ -280,6 +343,7 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "Test Key", + user: null, spend: 100, tags: [], }, @@ -298,6 +362,7 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "Test Key", + user: null, spend: 100, }, ]} @@ -315,6 +380,7 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "Test Key", + user: null, spend: 100, tags: [ { tag: "tag-1", usage: 50 }, @@ -340,6 +406,7 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "Test Key", + user: null, spend: 100, tags: [ { tag: "tag-1", usage: 50 }, @@ -367,6 +434,7 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "Test Key", + user: null, spend: 100, tags: [ { tag: "tag-1", usage: 50 }, @@ -404,6 +472,7 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "Test Key", + user: null, spend: 100, }, ]} @@ -424,6 +493,35 @@ describe("TopKeyView", () => { }); }); + it("should only look up keys that still exist in the database, from both the table and the chart", async () => { + mockKeyInfoV1Call.mockResolvedValue({ key: "info" }); + mockTransformKeyInfo.mockReturnValue({ transformed: "data" } as unknown as KeyResponse); + + const user = userEvent.setup(); + const { container } = render( + , + ); + + expect(screen.getByRole("button", { name: "stored-key" })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "session-key" })).not.toBeInTheDocument(); + await user.click(screen.getByText("session-key")); + + await user.click(screen.getByRole("button", { name: "Chart View" })); + const bars = chartBars(container); + expect(bars).toHaveLength(2); + bars.forEach((bar) => fireEvent.click(bar)); + + expect(await screen.findByText("Key Info View for stored-key")).toBeInTheDocument(); + expect(mockKeyInfoV1Call).toHaveBeenCalledTimes(1); + expect(mockKeyInfoV1Call).toHaveBeenCalledWith("test-token", "stored-key"); + }); + it("should close modal when close button is clicked", async () => { const mockKeyInfo = { key: "info" }; const mockTransformedData = { transformed: "data" } as unknown as KeyResponse; @@ -438,6 +536,7 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "Test Key", + user: null, spend: 100, }, ]} @@ -475,6 +574,7 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "Test Key", + user: null, spend: 100, }, ]} @@ -511,6 +611,7 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "Test Key", + user: null, spend: 100, }, ]} @@ -550,6 +651,7 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "Test Key", + user: null, spend: 100, }, ]} @@ -580,6 +682,7 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "Test Key", + user: null, spend: 100, }, ]} @@ -610,6 +713,7 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "Test Key", + user: null, spend: 100, tags: [ { tag: "tag-low", usage: 10 }, @@ -643,6 +747,7 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: "This is a very long key alias", + user: null, spend: 100, }, ]} @@ -659,11 +764,12 @@ describe("TopKeyView", () => { { api_key: "key-123", key_alias: null, + user: null, spend: 100, }, ]} />, ); - expect(screen.getByText("-")).toBeInTheDocument(); + expect(screen.getAllByText("-")).toHaveLength(1); }); }); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx index df1a51d8e38..178a5b7ec37 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx @@ -15,8 +15,22 @@ import { TagUsage } from "../../types"; const TOP_KEYS_LIMITS = [5, 10, 25, 50] as const; +export interface TopKeyItem { + api_key: string; + key_alias: string | null; + user?: string | null; + key_exists?: boolean | null; + tags?: TagUsage[] | null; + spend: number; +} + +const KEY_NOT_IN_DATABASE_TOOLTIP = + "This key is no longer in the database (deleted, or a CLI/SSO session key), so its details can't be opened"; + +const canOpenKeyInfo = (item: TopKeyItem) => item.key_exists !== false; + interface TopKeyViewProps { - topKeys: any[]; + topKeys: TopKeyItem[]; teams: any[] | null; showTags?: boolean; topKeysLimit: number; @@ -43,8 +57,8 @@ const TopKeyView: React.FC = ({ topKeys, teams, showTags = fals }); }; - const handleKeyClick = async (item: any) => { - if (!accessToken) return; + const handleKeyClick = async (item: TopKeyItem) => { + if (!accessToken || !canOpenKeyInfo(item)) return; try { const keyInfo = await keyInfoV1Call(accessToken, item.api_key); @@ -88,13 +102,27 @@ const TopKeyView: React.FC = ({ topKeys, teams, showTags = fals { header: "Key ID", accessorKey: "api_key", - cell: (info: any) => handleKeyClick(info.row.original)} />, + cell: (info: any) => + canOpenKeyInfo(info.row.original) ? ( + handleKeyClick(info.row.original)} /> + ) : ( + + ), }, { header: "Key Alias", accessorKey: "key_alias", cell: (info: any) => info.getValue() || "-", }, + ...(topKeys.some((k) => k.user) + ? [ + { + header: "User", + accessorKey: "user", + cell: (info: any) => info.getValue() || "-", + }, + ] + : []), ]; const tagsColumn = { diff --git a/ui/litellm-dashboard/src/components/UsagePage/types.ts b/ui/litellm-dashboard/src/components/UsagePage/types.ts index e8bd3cb3a87..d53db68bb9f 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/types.ts +++ b/ui/litellm-dashboard/src/components/UsagePage/types.ts @@ -50,6 +50,7 @@ export interface KeyMetadata { team_id: string | null; user_id?: string | null; user_email?: string | null; + key_exists?: boolean | null; tags?: { tag: string; usage: number }[]; } diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index d43adfe1ae4..7aa34c5752c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7805,7 +7805,7 @@ export interface paths { * - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules. * - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. * - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely. - * - enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Anthropic and Bedrock Claude models only. + * - enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Supported Claude models on Anthropic, Bedrock, Vertex AI, and Azure AI only. * - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} * - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. @@ -8286,7 +8286,7 @@ export interface paths { * - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules. * - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. * - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely. - * - enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Anthropic and Bedrock Claude models only. + * - enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Supported Claude models on Anthropic, Bedrock, Vertex AI, and Azure AI only. * - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. * - blocked: Optional[bool] - Whether the key is blocked * - aliases: Optional[dict] - Model aliases for the key - [Docs](https://litellm.vercel.app/docs/proxy/virtual_keys#model-aliases) @@ -29459,6 +29459,8 @@ export interface components { KeyMetadata: { /** Key Alias */ key_alias?: string | null; + /** Key Exists */ + key_exists?: boolean | null; /** Team Id */ team_id?: string | null; /** User Email */ diff --git a/ui/litellm-dashboard/tests/top_key_view.test.tsx b/ui/litellm-dashboard/tests/top_key_view.test.tsx index 51662b8f453..6751b639339 100644 --- a/ui/litellm-dashboard/tests/top_key_view.test.tsx +++ b/ui/litellm-dashboard/tests/top_key_view.test.tsx @@ -28,12 +28,15 @@ describe("TopKeyView", () => { teams: null, premiumUser: true, showTags: false, + topKeysLimit: 5, + setTopKeysLimit: vi.fn(), }; const mockKeysWithTags = [ { api_key: "key-1", key_alias: "Production Key", + user: null, tags: [ { tag: "production", usage: 0.005 } as TagUsage, // <$0.01 { tag: "high-volume", usage: 125.5 } as TagUsage, // High spend @@ -44,6 +47,7 @@ describe("TopKeyView", () => { { api_key: "key-2", key_alias: "Staging Key", + user: null, tags: [ { tag: "staging", usage: 45.75 } as TagUsage, // Medium spend { tag: "testing", usage: 0.008 } as TagUsage, // <$0.01 @@ -54,6 +58,7 @@ describe("TopKeyView", () => { { api_key: "key-3", key_alias: "Development Key", + user: null, tags: [ { tag: "dev", usage: 0.002 } as TagUsage, // <$0.01 { tag: "experimental", usage: 0.001 } as TagUsage, // <$0.01 @@ -65,11 +70,15 @@ describe("TopKeyView", () => { beforeEach(() => { vi.clearAllMocks(); mockUseAuthorized.mockReturnValue({ + isLoading: false, + isAuthorized: true, token: "mock-token", accessToken: mockProps.accessToken, userId: mockProps.userID, userEmail: "test@example.com", userRole: mockProps.userRole, + userRoleLabel: mockProps.userRole, + isViewOnly: false, premiumUser: mockProps.premiumUser, disabledPersonalKeyCreation: false, showSSOBanner: false, @@ -181,13 +190,14 @@ describe("TopKeyView", () => { { api_key: "key-no-tags", key_alias: "No Tags Key", + user: null, tags: [], spend: 10.0, }, ]; renderWithProviders(); - expect(screen.getByText("-")).toBeInTheDocument(); + expect(screen.getAllByText("-")).toHaveLength(1); }); it("should handle keys with undefined tags", () => { @@ -195,13 +205,14 @@ describe("TopKeyView", () => { { api_key: "key-undefined-tags", key_alias: "Undefined Tags Key", + user: null, tags: undefined, spend: 5.0, }, ]; renderWithProviders(); - expect(screen.getByText("-")).toBeInTheDocument(); + expect(screen.getAllByText("-")).toHaveLength(1); }); it("should handle keys with null tags", () => { @@ -209,13 +220,14 @@ describe("TopKeyView", () => { { api_key: "key-null-tags", key_alias: "Null Tags Key", + user: null, tags: null, spend: 3.0, }, ]; renderWithProviders(); - expect(screen.getByText("-")).toBeInTheDocument(); + expect(screen.getAllByText("-")).toHaveLength(1); }); }); @@ -225,6 +237,7 @@ describe("TopKeyView", () => { { api_key: "key-long-tags", key_alias: "Long Tags Key", + user: null, tags: [{ tag: "very-long-tag-name", usage: 10.0 } as TagUsage, { tag: "short", usage: 5.0 } as TagUsage], spend: 15.0, }, @@ -245,12 +258,14 @@ describe("TopKeyView", () => { { api_key: "key-mixed-1", key_alias: "Mixed Key 1", + user: null, tags: [{ tag: "expensive", usage: 999.99 } as TagUsage, { tag: "cheap", usage: 0.001 } as TagUsage], spend: 1000.0, }, { api_key: "key-mixed-2", key_alias: "Mixed Key 2", + user: null, tags: [{ tag: "moderate", usage: 50.0 } as TagUsage, { tag: "tiny", usage: 0.005 } as TagUsage], spend: 50.01, }, @@ -292,6 +307,7 @@ describe("TopKeyView", () => { { api_key: "test-key-123", key_alias: "Test Key", + user: null, tags: [], spend: 25.5, }, diff --git a/uv.lock b/uv.lock index ab9582ed134..db2fb11c6e3 100644 --- a/uv.lock +++ b/uv.lock @@ -3290,6 +3290,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" }, ] +[[package]] +name = "httpcore2" +version = "2.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "h11" }, + { name = "truststore" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/15/8c/e925b1c92018abb3a1863ce1549d76d2381e334d21d65d4ac8f65dabd78a/httpcore2-2.13.0.tar.gz", hash = "sha256:2adc8be4fb285fbcd6d894298db3b52c177e74b6674eda3a76bd36be3292a3db", size = 67740, upload-time = "2026-09-14T14:18:04.717Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/0d/117a771a2bb91df334b66bf4da14cd02f21aefbcfe53180f336ce55e8f90/httpcore2-2.13.0-py3-none-any.whl", hash = "sha256:35ae5be347aa40467b4a5dc032ac67ebb6d27189fc97e8cebcf99616f6a1bb9e", size = 83162, upload-time = "2026-09-14T14:18:02.529Z" }, +] + [[package]] name = "httplib2" version = "0.32.0" @@ -3331,6 +3344,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d2/fd/6668e5aec43ab844de6fc74927e155a3b37bf40d7c3790e49fc0406b6578/httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc", size = 8960, upload-time = "2025-10-10T21:48:21.158Z" }, ] +[[package]] +name = "httpx2" +version = "2.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio", marker = "sys_platform != 'emscripten'" }, + { name = "httpcore2", marker = "sys_platform != 'emscripten'" }, + { name = "httpx2-jsfetch", marker = "python_full_version >= '3.12' and sys_platform == 'emscripten'" }, + { name = "idna" }, + { name = "truststore", marker = "sys_platform != 'emscripten'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b9/a0/e9deef4654132857b5a5dbe4eddd0ac59c2814500e11f2f5044cd81103ee/httpx2-2.13.0.tar.gz", hash = "sha256:81bd07dc67a3701729ef1f777a3c00c915d4539604fdb5afd327f8682f6b7b44", size = 100290, upload-time = "2026-09-14T14:18:05.486Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/d1/a0c72b0e006df654709fbc366cc5bcb53e5aee13e1e3395152c6dd293376/httpx2-2.13.0-py3-none-any.whl", hash = "sha256:fc12720cedf72faa26cca6b4ca394e05c894e7d7933fc45cafe767960804e49a", size = 95565, upload-time = "2026-09-14T14:18:03.553Z" }, +] + +[[package]] +name = "httpx2-jsfetch" +version = "1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cd/c4/0e5636363151a2a1795e0a77617168b9ca438e1748ec05fc9b5687f93d64/httpx2_jsfetch-1.0.tar.gz", hash = "sha256:70a0e3eabfef7cce5ad9c629f7d01ca05e418f586646f4ddf14782e4c1454c60", size = 6872, upload-time = "2026-08-07T00:13:07.492Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9b/43/832f631d32e4f1211caa2ba368317739fe71f0b8530e4c9d15dc454bac2a/httpx2_jsfetch-1.0-py3-none-any.whl", hash = "sha256:cb916b707601e69a07721aabc8f3f6659be3a6893bc1ff5c6f9e02241df2da32", size = 6382, upload-time = "2026-08-07T00:13:06.567Z" }, +] + [[package]] name = "huey" version = "2.6.0" @@ -3494,11 +3533,11 @@ wheels = [ [[package]] name = "idna" -version = "3.15" +version = "3.19" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/82/77/7b3966d0b9d1d31a36ddf1746926a11dface89a83409bf1483f0237aa758/idna-3.15.tar.gz", hash = "sha256:ca962446ea538f7092a95e057da437618e886f4d349216d2b1e294abfdb65fdc", size = 199245, upload-time = "2026-05-12T22:45:57.011Z" } +sdist = { url = "https://files.pythonhosted.org/packages/5f/f7/abb373e5757eaec4b922b92f97ec8d6d7e057cf06778247604fbc4e7c3f3/idna-3.19.tar.gz", hash = "sha256:5e0811a4383b21dc5838069f801c4fb62113b7447663d2530d2bd6e77b49bf15", size = 215237, upload-time = "2026-08-18T05:14:24.27Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d2/23/408243171aa9aaba178d3e2559159c24c1171a641aa83b67bdd3394ead8e/idna-3.15-py3-none-any.whl", hash = "sha256:048adeaf8c2d788c40fee287673ccaa74c24ffd8dcf09ffa555a2fbb59f10ac8", size = 72340, upload-time = "2026-05-12T22:45:55.733Z" }, + { url = "https://files.pythonhosted.org/packages/57/b0/0e52c878c53f245edd3a11020f20979b3f490f245af532c7cae3027754b5/idna-3.19-py3-none-any.whl", hash = "sha256:815e7be7a7806d54abb586dc943addc79e8b2ee16915059658cbeff4b1b43bf4", size = 68550, upload-time = "2026-08-18T05:14:22.343Z" }, ] [[package]] @@ -4187,20 +4226,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/13/d6/bdf6f0481cc57ef300d6b1eb48cf1400c0409be715d6eb3cabadd1142a09/langchain_core-1.4.8-py3-none-any.whl", hash = "sha256:d84c28b05e3ba8d4271d0827aad5b592ccdaaf986e76768c23503f0a2045e8aa", size = 557416, upload-time = "2026-06-18T19:39:21.902Z" }, ] -[[package]] -name = "langchain-mcp-adapters" -version = "0.2.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "langchain-core" }, - { name = "mcp" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/d9/52/cebf0ef5b1acef6cbc63d671171d43af70f12d19f55577909c7afa79fb6e/langchain_mcp_adapters-0.2.1.tar.gz", hash = "sha256:58e64c44e8df29ca7eb3b656cf8c9931ef64386534d7ca261982e3bdc63f3176", size = 36394, upload-time = "2025-12-09T16:28:38.98Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/03/81/b2479eb26861ab36be851026d004b2d391d789b7856e44c272b12828ece0/langchain_mcp_adapters-0.2.1-py3-none-any.whl", hash = "sha256:9f96ad4c64230f6757297fec06fde19d772c99dbdfbca987f7b7cfd51ff77240", size = 22708, upload-time = "2025-12-09T16:28:37.877Z" }, -] - [[package]] name = "langchain-openai" version = "1.1.14" @@ -4537,7 +4562,9 @@ grpc = [ { name = "grpcio" }, ] mcp = [ + { name = "httpx2" }, { name = "mcp" }, + { name = "pydantic" }, ] mlflow = [ { name = "mlflow" }, @@ -4555,12 +4582,14 @@ proxy = [ { name = "granian" }, { name = "gunicorn" }, { name = "hiredis" }, + { name = "httpx2" }, { name = "inquirerpy" }, { name = "litellm-enterprise" }, { name = "litellm-proxy-extras" }, { name = "mcp" }, { name = "orjson" }, { name = "polars" }, + { name = "pydantic" }, { name = "pyjwt" }, { name = "pynacl" }, { name = "pyroscope-io", marker = "sys_platform != 'win32'" }, @@ -4632,7 +4661,6 @@ ci = [ { name = "google-generativeai" }, { name = "jsonlines" }, { name = "langchain" }, - { name = "langchain-mcp-adapters" }, { name = "langchain-openai" }, { name = "langgraph" }, { name = "langgraph-prebuilt" }, @@ -4752,6 +4780,8 @@ requires-dist = [ { name = "gunicorn", marker = "extra == 'proxy'", specifier = ">=23.0.0,<24.0" }, { name = "hiredis", marker = "extra == 'proxy'", specifier = ">=3.0.0,<4.0" }, { name = "httpx", extras = ["http2"], specifier = ">=0.28.0,<1.0" }, + { name = "httpx2", marker = "extra == 'mcp'", specifier = ">=2.5.0,<3" }, + { name = "httpx2", marker = "extra == 'proxy'", specifier = ">=2.5.0,<3" }, { name = "importlib-metadata", specifier = ">=8.0.0,<9.0" }, { name = "inquirerpy", marker = "extra == 'cli'", specifier = ">=0.3.4,<1.0" }, { name = "inquirerpy", marker = "extra == 'proxy'", specifier = ">=0.3.4,<1.0" }, @@ -4763,8 +4793,8 @@ requires-dist = [ { name = "litellm-proxy-extras", marker = "extra == 'proxy'", editable = "litellm-proxy-extras" }, { name = "llm-sandbox", marker = "extra == 'proxy-runtime'", specifier = ">=0.3.39,<1.0" }, { name = "mangum", marker = "extra == 'proxy-runtime'", specifier = ">=0.17.0,<1.0" }, - { name = "mcp", marker = "extra == 'mcp'", specifier = ">=1.28.1,<2.0" }, - { name = "mcp", marker = "extra == 'proxy'", specifier = ">=1.28.1,<2.0" }, + { name = "mcp", marker = "extra == 'mcp'", specifier = ">=2.2.0,<3" }, + { name = "mcp", marker = "extra == 'proxy'", specifier = ">=2.2.0,<3" }, { name = "mlflow", marker = "extra == 'mlflow'", specifier = ">=3.11.1,<4.0" }, { name = "numpy", marker = "extra == 'stt-nvidia-riva'", specifier = ">=1.26.0" }, { name = "numpydoc", marker = "extra == 'utils'", specifier = ">=1.8.0,<2.0" }, @@ -4780,7 +4810,10 @@ requires-dist = [ { name = "prometheus-client", marker = "extra == 'proxy-runtime'", specifier = ">=0.20.0,<1.0" }, { name = "psycopg", marker = "extra == 'extra-proxy'", specifier = ">=3.2,<4.0" }, { name = "psycopg-binary", marker = "extra == 'extra-proxy'", specifier = ">=3.2,<4.0" }, - { name = "pydantic", specifier = ">=2.10.0,<3.0.0" }, + { name = "pydantic", marker = "python_full_version < '3.14'", specifier = ">=2.11.0,<3.0.0" }, + { name = "pydantic", marker = "python_full_version >= '3.14'", specifier = ">=2.12.0,<3.0.0" }, + { name = "pydantic", marker = "extra == 'mcp'", specifier = ">=2.12.0,<3" }, + { name = "pydantic", marker = "extra == 'proxy'", specifier = ">=2.12.0,<3" }, { name = "pydantic-settings", specifier = ">=2.14.1,<3.0" }, { name = "pyjwt", marker = "extra == 'proxy'", specifier = ">=2.13.0,<3.0" }, { name = "pynacl", marker = "extra == 'proxy'", specifier = ">=1.6.2,<2.0" }, @@ -4803,7 +4836,8 @@ requires-dist = [ { name = "soundfile", marker = "extra == 'proxy'", specifier = ">=0.12.1,<1.0" }, { name = "soundfile", marker = "extra == 'stt-nvidia-riva'", specifier = ">=0.12.1" }, { name = "starlette", marker = "extra == 'proxy'", specifier = ">=1.0.1,<2.0" }, - { name = "tiktoken", specifier = ">=0.8.0,<1.0" }, + { name = "tiktoken", marker = "python_full_version < '3.14'", specifier = ">=0.8.0,<1.0" }, + { name = "tiktoken", marker = "python_full_version >= '3.14'", specifier = ">=0.12.0,<1.0" }, { name = "tokenizers", specifier = ">=0.21.0,<1.0" }, { name = "tomlkit", marker = "extra == 'cli'", specifier = ">=0.13.3,<1.0" }, { name = "tomlkit", marker = "extra == 'proxy'", specifier = ">=0.13.3,<1.0" }, @@ -4827,7 +4861,6 @@ ci = [ { name = "google-generativeai", specifier = "==0.8.6" }, { name = "jsonlines", specifier = "==4.0.0" }, { name = "langchain", specifier = "==1.3.9" }, - { name = "langchain-mcp-adapters", specifier = "==0.2.1" }, { name = "langchain-openai", specifier = "==1.1.14" }, { name = "langgraph", specifier = ">=1.2.4,<1.3.0" }, { name = "langgraph-prebuilt", specifier = ">=1.1.0,<1.3.0" }, @@ -4886,7 +4919,7 @@ dev = [ ] e2e-dev = [ { name = "locust", specifier = "==2.45.0" }, - { name = "mcp", specifier = ">=1.28.1,<2.0" }, + { name = "mcp", specifier = ">=2.2.0,<3" }, { name = "playwright", specifier = "==1.61.0" }, { name = "psutil", specifier = "==7.2.2" }, { name = "websockets", specifier = ">=15.0.1,<16.0" }, @@ -5364,15 +5397,15 @@ wheels = [ [[package]] name = "mcp" -version = "1.28.1" +version = "2.2.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, - { name = "httpx" }, - { name = "httpx-sse" }, + { name = "httpx2" }, { name = "jsonschema" }, + { name = "mcp-types" }, + { name = "opentelemetry-api" }, { name = "pydantic" }, - { name = "pydantic-settings" }, { name = "pyjwt", extra = ["crypto"] }, { name = "python-multipart" }, { name = "pywin32", marker = "sys_platform == 'win32'" }, @@ -5382,9 +5415,22 @@ dependencies = [ { name = "typing-inspection" }, { name = "uvicorn", marker = "sys_platform != 'emscripten'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/6e/77/9450b8f251a13affb6281997d0523c4615f8a8b35d0b21ff30db3a5aac9d/mcp-1.28.1.tar.gz", hash = "sha256:d51e36a5f5644faea4f85ea649bfffa6bc6c26770d42798ad6a3de3d2ba69683", size = 638501, upload-time = "2026-06-26T12:57:29.093Z" } +sdist = { url = "https://files.pythonhosted.org/packages/76/31/ac54fb0fdd5b37de704486e288bba4fbbb463f24cfcfedbede407b854513/mcp-2.2.0.tar.gz", hash = "sha256:2dc37ecb1974becdcebdbf7561e7c15a07dbbf20ba21ba16c3593b3038b3afbd", size = 4084129, upload-time = "2026-09-07T16:06:23.439Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e2/5e/d118fce19f87a2e7d8101c35c8ae0ec289098a4df0ff244cec23e415aca0/mcp-1.28.1-py3-none-any.whl", hash = "sha256:2726bca5e7193f61c5dde8b12500a6de2d9acf6d1a1c0be9e8c2e706437991df", size = 222620, upload-time = "2026-06-26T12:57:27.218Z" }, + { url = "https://files.pythonhosted.org/packages/1b/ff/8e7eade68b8a28f7da0ed1085544341b51f9c935dbf6b95c76b7edfea6a0/mcp-2.2.0-py3-none-any.whl", hash = "sha256:bde982589473a060ae145e3406e9a5333fe538c97229ba841f5a7f92be004f81", size = 365656, upload-time = "2026-09-07T16:06:19.711Z" }, +] + +[[package]] +name = "mcp-types" +version = "2.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydantic" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ae/91/762d7755d971aff8a28d75f7961656148edf27875c8026e6385aaab08ae7/mcp_types-2.2.0.tar.gz", hash = "sha256:d3ed53703ddd10d9c6399f29d322bb66f3f67ab41348ac8556ba23e07fedefad", size = 65892, upload-time = "2026-09-07T16:06:25.187Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8f/d7/6ffba5d8cd5dd9b8a19478875c50e04945314ba5074e84d749283f27f62d/mcp_types-2.2.0-py3-none-any.whl", hash = "sha256:ea476b73ee86709ab5abc9452385ed36cc05907e582355622e294595c9a04f13", size = 69106, upload-time = "2026-09-07T16:06:21.461Z" }, ] [[package]] @@ -9845,6 +9891,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ce/13/53c2ab6ac27804769314554a062e0651a44db2360be47e21cf0a29d202ee/traceloop_sdk-0.33.12-py3-none-any.whl", hash = "sha256:d47a474afbf4a68ff38a702dbaca7b17d2d4f0b0e14dc2f1560b6bdd3859ac75", size = 25932, upload-time = "2024-11-13T20:29:25.174Z" }, ] +[[package]] +name = "truststore" +version = "0.10.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/53/a3/1585216310e344e8102c22482f6060c7a6ea0322b63e026372e6dcefcfd6/truststore-0.10.4.tar.gz", hash = "sha256:9d91bd436463ad5e4ee4aba766628dd6cd7010cf3e2461756b3303710eebc301", size = 26169, upload-time = "2025-08-12T18:49:02.73Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/19/97/56608b2249fe206a67cd573bc93cd9896e1efb9e98bce9c163bcdc704b88/truststore-0.10.4-py3-none-any.whl", hash = "sha256:adaeaecf1cbb5f4de3b1959b42d41f6fab57b2b1666adb59e89cb0b53361d981", size = 18660, upload-time = "2025-08-12T18:49:01.46Z" }, +] + [[package]] name = "typer" version = "0.25.1"