From da558936569b2bd13ce533e3115d3284cb775cb8 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 14 Nov 2026 15:41:30 -0800 Subject: [PATCH] ci(unit): cut unit check time with cheaper per-test resets and fewer non-unit tests (#51864) * test(unit): reset module state only when a test changed it The proxy conftest deep-copied every snapshotted litellm attribute before each test, and the copy of litellm.model_cost alone was a median 0.12 s for each of 56,031 proxy tests on the runner. The reset now copies an attribute back only when it no longer equals the snapshot. The enterprise and MCP server conftests reloaded litellm before every test, 3.27 s and 0.21 s per test. They now use the same reset. * test(unit): remove six flaky tests They failed on main without a code change behind them between 2026-11-05 and 2026-11-14: the vector store hook SDK fallback, the concurrent user config TPM count, three Prometheus metrics server port tests and the OpenCode list cost billing test. * ci: run four heavy unit shards on 16-core runners and keep two smoke jobs off pull requests proxy-feature-endpoints, proxy-endpoints, root and llms-providers ran 1.9 to 2.2 times faster with 8 workers on a 16-core runner in a benchmark run. windows-sdk and proxy-build-from-pip still run on pushes to main and manual runs. proxy-build-from-pip installs a published release and never runs the pull request's code. * test: move 109 tests that start a process or open a socket into tests/integration tests/unit is in-process by its own rules, and these tests spawn a Python interpreter, a shell or a local TCP or HTTP server. They took 921 s of unit shard time per run, 396 s of it in the MCP client transport tests. Each one now lives in a *_real_io.py file under the integration directory for its domain and runs with that directory's group. The tests themselves are unchanged apart from one repository-root path in the CLI test. * test(unit): copy back only the dict entries a test changed A test module that registers models at import leaves litellm.model_cost four entries larger than the snapshot before every test in its directory, so the reset still deep-copied the whole price map for each of them. It now rebuilds a changed dict from the live entries that still match and copies only the rest. * ci(smoke): keep the two jobs that skip on pull requests out of the smoke gate A skipped job passing the gate is what test_unit_passed_gate guards against, so the gate goes back to requiring success from every job it needs and windows-sdk and proxy-build-from-pip stop being part of it. They still run and report on pushes to main and manual runs. --- .github/workflows/test-merge-smoke.yml | 4 +- .github/workflows/test-unit.yml | 14 +- tests/_module_state.py | 107 ++ .../test_unit_passed_gate.py | 2 +- tests/integration/AGENTS.md | 4 + .../test_litellm_http_client_sso_real_io.py} | 0 .../test_master_key_boot_check_real_io.py | 124 ++ ..._user_api_key_auth_request_flow_real_io.py | 155 ++ .../test_evicted_client_closer_real_io.py | 119 ++ .../test_valkey_semantic_cache_real_io.py | 32 + .../test_claude_settings_real_io.py | 272 +++ .../configuration/test_main_real_io.py | 20 + .../configuration/test_runtime_real_io.py | 83 + .../test_litellm_proxy_mcp_handler_real_io.py | 88 + .../mcp/test_mcp_client_real_io.py | 1527 +++++++++++++++++ .../test_openapi_to_mcp_generator_real_io.py | 44 + ...st_anthropic_cache_control_hook_real_io.py | 65 + .../test_anthropic_common_utils_real_io.py | 41 + .../sdk/test_azure_common_utils_real_io.py | 204 +++ .../sdk/test_configuration_real_io.py | 39 + tests/integration/sdk/test_context_real_io.py | 146 ++ tests/integration/sdk/test_docker_real_io.py | 50 + .../sdk/test_drop_params_env_var_real_io.py} | 0 .../sdk/test_eager_tiktoken_load_real_io.py} | 0 .../sdk/test_http_handler_real_io.py | 296 ++++ .../sdk/test_jwt_signing_real_io.py | 37 + .../sdk/test_lazy_imports_real_io.py | 23 + .../sdk/test_litellm_cli_real_io.py | 384 +++++ .../sdk/test_long_decode_real_io.py | 323 ++++ ...uter_fuzz_regressions_rust_only_real_io.py | 193 +++ .../sdk/test_router_host_real_io.py | 28 + .../sdk/test_streaming_handler_real_io.py | 367 ++++ .../sdk/test_token_counter_real_io.py | 207 +++ .../integration/sdk/test_tokenizer_real_io.py | 119 ++ .../caching/test_evicted_client_closer.py | 87 - .../caching/test_valkey_semantic_cache.py | 29 - tests/unit/enterprise/conftest.py | 40 +- tests/unit/harness/sandbox/test_docker.py | 45 - .../test_anthropic_cache_control_hook.py | 52 +- .../test_vector_store_pre_call_hook.py | 30 - .../test_streaming_handler.py | 66 - .../litellm_core_utils/test_token_counter.py | 175 +- .../unit/litellm_core_utils/test_tokenizer.py | 111 -- .../anthropic/test_anthropic_common_utils.py | 38 - .../llms/azure/test_azure_common_utils.py | 23 - .../llms/base_llm/auth/test_jwt_signing.py | 29 - .../llms/custom_httpx/test_http_handler.py | 243 --- .../llms/opencode/chat/test_transformation.py | 20 - .../vertex_ai/count_tokens/test_context.py | 132 -- .../llms/vertex_ai/gemini/test_long_decode.py | 220 +-- tests/unit/mcp_client/test_mcp_client.py | 1379 +-------------- .../proxy/auth/test_master_key_boot_check.py | 63 - .../test_user_api_key_auth_request_flow.py | 127 -- .../proxy/client/cli/test_claude_settings.py | 62 - tests/unit/proxy/conftest.py | 111 +- .../proxy/hooks/litellm_skills/test_main.py | 20 - tests/unit/proxy/mcp_server/conftest.py | 13 +- .../test_openapi_to_mcp_generator.py | 18 - .../rust_forward/supervisor/test_runtime.py | 60 +- .../proxy/test_prometheus_metrics_server.py | 47 - tests/unit/proxy/test_route_llm_request.py | 27 - .../mcp/test_litellm_proxy_mcp_handler.py | 48 - tests/unit/rust_bridge/test_configuration.py | 22 - tests/unit/test_lazy_imports.py | 19 - tests/unit/test_litellm_cli.py | 369 ---- .../test_router_fuzz_regressions_rust_only.py | 106 -- tests/unit/test_router/test_router_host.py | 14 - tests/windows_tests/AGENTS.md | 2 +- 68 files changed, 5143 insertions(+), 3821 deletions(-) create mode 100644 tests/_module_state.py rename tests/{unit/proxy/management_endpoints/sso/test_litellm_http_client_sso.py => integration/authorization/test_litellm_http_client_sso_real_io.py} (100%) create mode 100644 tests/integration/authorization/test_master_key_boot_check_real_io.py create mode 100644 tests/integration/authorization/test_user_api_key_auth_request_flow_real_io.py create mode 100644 tests/integration/caching/test_evicted_client_closer_real_io.py create mode 100644 tests/integration/caching/test_valkey_semantic_cache_real_io.py create mode 100644 tests/integration/configuration/test_claude_settings_real_io.py create mode 100644 tests/integration/configuration/test_main_real_io.py create mode 100644 tests/integration/configuration/test_runtime_real_io.py create mode 100644 tests/integration/mcp/test_litellm_proxy_mcp_handler_real_io.py create mode 100644 tests/integration/mcp/test_mcp_client_real_io.py create mode 100644 tests/integration/mcp/test_openapi_to_mcp_generator_real_io.py create mode 100644 tests/integration/observability/test_anthropic_cache_control_hook_real_io.py create mode 100644 tests/integration/sdk/test_anthropic_common_utils_real_io.py create mode 100644 tests/integration/sdk/test_azure_common_utils_real_io.py create mode 100644 tests/integration/sdk/test_configuration_real_io.py create mode 100644 tests/integration/sdk/test_context_real_io.py create mode 100644 tests/integration/sdk/test_docker_real_io.py rename tests/{unit/test_drop_params_env_var.py => integration/sdk/test_drop_params_env_var_real_io.py} (100%) rename tests/{unit/test_eager_tiktoken_load.py => integration/sdk/test_eager_tiktoken_load_real_io.py} (100%) create mode 100644 tests/integration/sdk/test_http_handler_real_io.py create mode 100644 tests/integration/sdk/test_jwt_signing_real_io.py create mode 100644 tests/integration/sdk/test_lazy_imports_real_io.py create mode 100644 tests/integration/sdk/test_litellm_cli_real_io.py create mode 100644 tests/integration/sdk/test_long_decode_real_io.py create mode 100644 tests/integration/sdk/test_router_fuzz_regressions_rust_only_real_io.py create mode 100644 tests/integration/sdk/test_router_host_real_io.py create mode 100644 tests/integration/sdk/test_streaming_handler_real_io.py create mode 100644 tests/integration/sdk/test_token_counter_real_io.py create mode 100644 tests/integration/sdk/test_tokenizer_real_io.py diff --git a/.github/workflows/test-merge-smoke.yml b/.github/workflows/test-merge-smoke.yml index 8ed1876a081..88acf50837f 100644 --- a/.github/workflows/test-merge-smoke.yml +++ b/.github/workflows/test-merge-smoke.yml @@ -178,6 +178,7 @@ jobs: windows-sdk: name: windows-sdk + if: github.event_name != 'pull_request' runs-on: windows-2025 timeout-minutes: 30 env: @@ -213,6 +214,7 @@ jobs: proxy-build-from-pip: name: proxy-build-from-pip + if: github.event_name != 'pull_request' runs-on: ubuntu-24.04 timeout-minutes: 10 services: @@ -352,7 +354,7 @@ jobs: smoke-passed: name: smoke passed - needs: [dashboard-build, core-checks, all-extras, windows-sdk, proxy-build-from-pip] + needs: [dashboard-build, core-checks, all-extras] if: always() runs-on: ubuntu-latest timeout-minutes: 2 diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index f27b58149ab..50bcc6c2b42 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -126,7 +126,7 @@ jobs: name: ${{ matrix.shard }} needs: [rust-bridge, assert-shard-coverage] if: ${{ !cancelled() }} - runs-on: ubuntu-latest + runs-on: ${{ matrix.runner || 'ubuntu-latest' }} timeout-minutes: ${{ matrix.job-timeout-minutes }} permissions: contents: read @@ -220,6 +220,7 @@ jobs: dist: loadscope - shard: llms-providers + runner: ubuntu-latest-16-cores test-path: |- tests/unit/llms/test_api_base_trailing_slash_routes.py tests/unit/llms/test_cache_control_and_reasoning.py @@ -380,7 +381,7 @@ jobs: tests/unit/llms/zai --ignore=tests/unit/llms/base_llm/batches/base_batches_config_test.py python-version: "3.12" - workers: 4 + workers: 8 reruns: 0 timeout-minutes: 20 job-timeout-minutes: 60 @@ -398,10 +399,11 @@ jobs: dist: loadscope - shard: root + runner: ubuntu-latest-16-cores test-path: >- tests/unit/test_*.py python-version: "3.12" - workers: 4 + workers: 8 reruns: 2 timeout-minutes: 20 job-timeout-minutes: 60 @@ -522,6 +524,7 @@ jobs: dist: loadscope - shard: proxy-endpoints + runner: ubuntu-latest-16-cores test-path: |- tests/unit/proxy/management_endpoints tests/unit/proxy/management_helpers @@ -539,13 +542,14 @@ jobs: --ignore=tests/unit/proxy/management_endpoints/test_roi_calculator_endpoints.py --ignore=tests/unit/proxy/management_helpers/test_audit_logs_proxy.py python-version: "3.12" - workers: 4 + workers: 8 reruns: 0 timeout-minutes: 20 job-timeout-minutes: 60 dist: loadscope - shard: proxy-feature-endpoints + runner: ubuntu-latest-16-cores test-path: |- tests/unit/proxy/guardrails --ignore=tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py @@ -578,7 +582,7 @@ jobs: tests/unit/proxy/config_resolvers tests/unit/proxy/utils python-version: "3.12" - workers: 4 + workers: 8 reruns: 0 timeout-minutes: 20 job-timeout-minutes: 60 diff --git a/tests/_module_state.py b/tests/_module_state.py new file mode 100644 index 00000000000..b168616d4c9 --- /dev/null +++ b/tests/_module_state.py @@ -0,0 +1,107 @@ +import copy +import inspect +import itertools +import warnings +from collections.abc import Mapping +from types import ModuleType +from typing import Final + +_SNAPSHOT_TYPES: Final = (list, dict, set, tuple, str, int, float, bool, bytes) +_MUTABLE_CONTAINERS: Final = (list, dict, set, bytearray) +_ABSENT: Final = object() + + +def _holds_mutable_container(value: object) -> bool: + if isinstance(value, _MUTABLE_CONTAINERS): + return True + if isinstance(value, tuple): + return any(_holds_mutable_container(element) for element in value) + return False + + +def _restored_value(value: object) -> object: + return copy.deepcopy(value) if _holds_mutable_container(value) else value + + +def _readable_attributes(module: ModuleType) -> tuple[tuple[str, object], ...]: + return tuple((attr, value) for attr, value in tuple(vars(module).items()) if not attr.startswith("_")) + + +def _snapshot_of(module: ModuleType, attr: str, value: object) -> object: + try: + return _restored_value(value) + except Exception as exc: + warnings.warn(f"conftest: could not snapshot {module.__name__}.{attr}: {exc}", stacklevel=2) + return _ABSENT + + +def snapshot_mutable_state(module: ModuleType) -> Mapping[str, object]: + candidates: Final = ( + (attr, _snapshot_of(module, attr, value)) + for attr, value in _readable_attributes(module) + if value is None or isinstance(value, _SNAPSHOT_TYPES) + ) + return {attr: value for attr, value in candidates if value is not _ABSENT} + + +def _still_matches(current: object, default: object) -> bool: + if current is default: + return True + if type(current) is not type(default): + return False + try: + return bool(current == default) + except Exception: + return False + + +def _restored_entries(current: dict[object, object], default: dict[object, object]) -> dict[object, object]: + return { + key: current[key] + if key in current and type(current[key]) is type(value) and current[key] == value + else _restored_value(value) + for key, value in default.items() + } + + +def _restored_from(current: object, default: object) -> object: + if type(current) is dict and type(default) is dict: + return _restored_entries(current, default) + return _restored_value(default) + + +def restore_mutable_state(module: ModuleType, snapshot: Mapping[str, object]) -> None: + live: Final = vars(module) + for attr, default in snapshot.items(): + current = live.get(attr, _ABSENT) + if _still_matches(current, default): + continue + try: + setattr(module, attr, _restored_from(current, default)) + except Exception as exc: + warnings.warn(f"conftest: could not restore {module.__name__}.{attr}: {exc}", stacklevel=2) + + +def _is_flushable(value: object) -> bool: + if inspect.isclass(value) or inspect.ismodule(value): + return False + return callable(getattr(value, "flush_cache", None)) + + +def _flushable_attributes(module: ModuleType) -> tuple[tuple[ModuleType, str], ...]: + return tuple((module, attr) for attr, value in _readable_attributes(module) if _is_flushable(value)) + + +def collect_flushable_caches(modules: tuple[ModuleType, ...]) -> tuple[tuple[ModuleType, str], ...]: + return tuple(itertools.chain.from_iterable(_flushable_attributes(module) for module in modules)) + + +def flush_caches(targets: tuple[tuple[ModuleType, str], ...]) -> None: + for module, attr in targets: + flush = getattr(vars(module).get(attr), "flush_cache", None) + if not callable(flush): + continue + try: + flush() + except Exception as exc: + warnings.warn(f"conftest: flush_cache failed on {module.__name__}.{attr}: {exc}", stacklevel=2) diff --git a/tests/code_coverage_tests/test_unit_passed_gate.py b/tests/code_coverage_tests/test_unit_passed_gate.py index 5a0d04681d4..2bb480d67c1 100644 --- a/tests/code_coverage_tests/test_unit_passed_gate.py +++ b/tests/code_coverage_tests/test_unit_passed_gate.py @@ -11,7 +11,7 @@ _WORKFLOWS_DIR: Final = Path(__file__).resolve().parents[2] / ".github" / "workf _TIERS: Final = ( ("test-linting.yml", "lint-passed", frozenset()), ("test-unit.yml", "unit-passed", frozenset({"coverage"})), - ("test-merge-smoke.yml", "smoke-passed", frozenset()), + ("test-merge-smoke.yml", "smoke-passed", frozenset({"windows-sdk", "proxy-build-from-pip"})), ) _Tier = tuple[str, str, frozenset[str]] diff --git a/tests/integration/AGENTS.md b/tests/integration/AGENTS.md index b430ab553a5..96ff87994a1 100644 --- a/tests/integration/AGENTS.md +++ b/tests/integration/AGENTS.md @@ -26,6 +26,10 @@ directory that a `GROUPS` entry in `run.py` selects. There is no manifest and no tests. A product bug the test exposes is `pytest.skip("BUG: ")` at the top of the body, not a fix in the test and not a deletion. Needs no network, proxy, DB or Redis: `tests/unit` +A test that starts its own process or opens its own socket but needs no proxy also lives here, in a +`*_real_io.py` file next to the domain it covers, because `tests/unit` is in-process. It asks for none of +the gateway fixtures, so it runs as plain pytest inside its group + GitHub Actions integration groups run every file here except the OCI and sandbox tests, which no CI runs, and the files in `GITHUB_FILES` in `run.py`. `test-integration.yml` runs the same groups through the same `run_integration.sh` on GitHub Actions, on main every three hours from 8am to 8pm Pacific Monday to Saturday, diff --git a/tests/unit/proxy/management_endpoints/sso/test_litellm_http_client_sso.py b/tests/integration/authorization/test_litellm_http_client_sso_real_io.py similarity index 100% rename from tests/unit/proxy/management_endpoints/sso/test_litellm_http_client_sso.py rename to tests/integration/authorization/test_litellm_http_client_sso_real_io.py diff --git a/tests/integration/authorization/test_master_key_boot_check_real_io.py b/tests/integration/authorization/test_master_key_boot_check_real_io.py new file mode 100644 index 00000000000..68d62ed8316 --- /dev/null +++ b/tests/integration/authorization/test_master_key_boot_check_real_io.py @@ -0,0 +1,124 @@ +import re +import shutil +import subprocess +import sys +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +import pytest + +from litellm.proxy.auth.master_key_boot_check import ( + GENERATE_MASTER_KEY_COMMAND, + MASTER_KEY_ENV_VAR, + MIGRATE_FROM_MASTER_KEY_ENV_VAR, + PRINT_NEW_MASTER_KEY_COMMAND, + ConfigFileSource, + EnvironmentSource, + MasterKeyBootVerdict, + SafeMasterKey, + StoredSecretsMigration, + UnsafeMasterKeyReason, + UnsafeMasterKeyRefused, + announce_on_stderr_at_exit, + master_key_boot_verdict, + render_refusal, +) + +PUBLICLY_KNOWN_KEY: Final = "sk-" + "1234" + + +def _verdict( + master_key: str | None, + general_settings: Mapping[str, object] | None = None, + *, + environment_master_key: str | None = None, + config_file_path: str | None = None, + override_env_is_on: bool = False, + salt_key_is_set: bool = False, + database_is_configured: bool = False, +) -> MasterKeyBootVerdict: + return master_key_boot_verdict( + master_key=master_key, + environment_master_key=environment_master_key, + general_settings=general_settings or {}, + config_file_path=config_file_path, + override_env_is_on=override_env_is_on, + salt_key_is_set=salt_key_is_set, + database_is_configured=database_is_configured, + ) + + +def _refusal( + reason: UnsafeMasterKeyReason = UnsafeMasterKeyReason.PUBLICLY_KNOWN, + source: ConfigFileSource | EnvironmentSource = EnvironmentSource(), + environment_variable_is_set: bool = False, + migration: StoredSecretsMigration | None = None, +) -> UnsafeMasterKeyRefused: + return UnsafeMasterKeyRefused( + reason=reason, + source=source, + environment_variable_is_set=environment_variable_is_set, + migration=migration, + ) + + +_MIGRATION = StoredSecretsMigration(from_master_key=PUBLICLY_KNOWN_KEY, encrypted_value_count=3) + + +@pytest.mark.skipif(shutil.which("openssl") is None, reason="the printed command shells out to openssl") +@pytest.mark.parametrize("from_master_key", [PUBLICLY_KNOWN_KEY, "", " " + PUBLICLY_KNOWN_KEY]) +def test_printed_migration_commands_save_both_keys_to_the_env_file(tmp_path: Path, from_master_key: str): + from dotenv import dotenv_values + + text = render_refusal( + _refusal(migration=StoredSecretsMigration(from_master_key=from_master_key, encrypted_value_count=1)) + ) + commands = [line.strip() for line in text.splitlines() if line.strip().startswith("echo ")] + + subprocess.run(["bash", "-c", "\n".join(commands)], cwd=tmp_path, capture_output=True, text=True, check=True) + + saved = dotenv_values(tmp_path / ".env") + assert saved[MIGRATE_FROM_MASTER_KEY_ENV_VAR] == from_master_key + assert _verdict(saved[MASTER_KEY_ENV_VAR]) == SafeMasterKey() + assert len(commands) == 2 + + +@pytest.mark.skipif(shutil.which("openssl") is None, reason="the printed command shells out to openssl") +def test_printed_command_saves_a_key_the_boot_check_accepts(tmp_path: Path): + completed = subprocess.run( + ["bash", "-c", GENERATE_MASTER_KEY_COMMAND], cwd=tmp_path, capture_output=True, text=True, check=True + ) + + saved = (tmp_path / ".env").read_text() + match = re.fullmatch(rf"{MASTER_KEY_ENV_VAR}=(sk-[0-9a-f]{{64}})\n", saved) + assert match is not None + assert completed.stdout == saved + assert _verdict(match.group(1)) == SafeMasterKey() + + +@pytest.mark.skipif(shutil.which("openssl") is None, reason="the printed command shells out to openssl") +def test_rotation_command_prints_a_key_the_boot_check_accepts_and_saves_nothing(tmp_path: Path): + completed = subprocess.run( + ["bash", "-c", PRINT_NEW_MASTER_KEY_COMMAND], cwd=tmp_path, capture_output=True, text=True, check=True + ) + + assert re.fullmatch(r"sk-[0-9a-f]{64}\n", completed.stdout) is not None + assert _verdict(completed.stdout.strip()) == SafeMasterKey() + assert list(tmp_path.iterdir()) == [] + + +def test_announced_fix_is_the_last_thing_a_crashing_process_prints(): + crash_after_announcing = ( + "from litellm.proxy.auth.master_key_boot_check import announce_on_stderr_at_exit\n" + "announce_on_stderr_at_exit('THE FIX')\n" + "print('buffered stdout')\n" + "raise RuntimeError('lifespan failed')\n" + ) + + completed = subprocess.run([sys.executable, "-c", crash_after_announcing], capture_output=True, text=True) + + assert completed.returncode != 0 + assert "RuntimeError: lifespan failed" in completed.stderr + assert completed.stderr.endswith("THE FIX") + assert completed.stdout == "buffered stdout\n" diff --git a/tests/integration/authorization/test_user_api_key_auth_request_flow_real_io.py b/tests/integration/authorization/test_user_api_key_auth_request_flow_real_io.py new file mode 100644 index 00000000000..42e77e5fa07 --- /dev/null +++ b/tests/integration/authorization/test_user_api_key_auth_request_flow_real_io.py @@ -0,0 +1,155 @@ +import json +import os +import subprocess +import sys +from datetime import datetime, timezone +from pathlib import Path +from textwrap import dedent +from types import SimpleNamespace +from typing import Final +from unittest.mock import ANY + +import litellm +import litellm.proxy.proxy_server +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.types.videos.utils import encode_video_id_with_provider + + +class _RoutingRequest: + def __init__(self, headers=None, query_params=None): + self.headers = headers or {} + self.query_params = query_params or {} + self.state = SimpleNamespace() + + +# --------------------------------------------------------------------------- + +# _run_centralized_common_checks — centralized authz gate +# --------------------------------------------------------------------------- + + +_RECORDING_DDTRACE = dedent( + """ + import functools + import inspect + + + class _Span: + def __enter__(self): + return self + + def __exit__(self, *exc): + return None + + + class _Tracer: + def __init__(self): + self.spans = [] + + def wrap(self, name=None, **kwargs): + def decorator(f): + span_name = name or f"{f.__module__}.{f.__name__}" + if inspect.iscoroutinefunction(f): + + @functools.wraps(f) + async def async_wrapped(*args, **kw): + self.spans.append(span_name) + return await f(*args, **kw) + + return async_wrapped + + @functools.wraps(f) + def wrapped(*args, **kw): + self.spans.append(span_name) + return f(*args, **kw) + + return wrapped + + return decorator + + def trace(self, name, **kwargs): + return _Span() + + def current_span(self): + return None + + def current_root_span(self): + return None + + + tracer = _Tracer() + """ +) + +_DDTRACE_AUTH_PROBE = dedent( + """ + import asyncio + import json + + import ddtrace + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + proxy_server.master_key = "sk-probe" + + + async def auth(api_key): + request = Request(scope={"type": "http", "headers": [], "method": "POST", "path": "/chat/completions"}) + request._url = URL(url="/chat/completions") + try: + await user_api_key_auth( + request=request, + api_key=api_key, + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + custom_litellm_key_header=None, + ) + return "accepted" + except ProxyException: + return "rejected" + + + async def main(): + outcomes = [await auth("Bearer sk-probe"), await auth("Bearer sk-wrong")] + print(json.dumps({"outcomes": outcomes, "spans": ddtrace.tracer.spans})) + + + asyncio.run(main()) + """ +) + + +def test_user_api_key_auth_opens_a_datadog_span_for_accepted_and_rejected_keys(tmp_path: Path): + stub_root = tmp_path / "site" + (stub_root / "ddtrace").mkdir(parents=True) + (stub_root / "ddtrace" / "__init__.py").write_text(_RECORDING_DDTRACE) + probe = tmp_path / "probe.py" + probe.write_text(_DDTRACE_AUTH_PROBE) + repo_root = Path(litellm.__file__).resolve().parent.parent + env = { + **os.environ, + "USE_DDTRACE": "true", + "PYTHONPATH": os.pathsep.join( + [str(stub_root), str(repo_root)] + [p for p in (os.environ.get("PYTHONPATH"),) if p] + ), + } + + result = subprocess.run( + [sys.executable, str(probe)], env=env, cwd=repo_root, capture_output=True, text=True, check=False + ) + + assert result.returncode == 0, result.stderr[-4000:] + report = json.loads(result.stdout.strip().splitlines()[-1]) + assert report["outcomes"] == ["accepted", "rejected"] + auth_span = "litellm.proxy.auth.user_api_key_auth.user_api_key_auth" + assert [span for span in report["spans"] if span == auth_span] == [auth_span, auth_span] + + +_VIDEO_ID: Final = encode_video_id_with_provider(video_id="video_raw", provider="openai", model_id="sora-2") +_EXPIRED_BEFORE_ANY_TEST_RUN: Final = datetime(2000, 1, 1, tzinfo=timezone.utc) diff --git a/tests/integration/caching/test_evicted_client_closer_real_io.py b/tests/integration/caching/test_evicted_client_closer_real_io.py new file mode 100644 index 00000000000..c747d4bd798 --- /dev/null +++ b/tests/integration/caching/test_evicted_client_closer_real_io.py @@ -0,0 +1,119 @@ +""" +Tests for EvictedClientCloser. + +An evicted client must stay open long enough for a request that already holds it +to finish, and must then actually be closed, otherwise its connection pool is +retained until a generational collection runs. A client the caller supplied is +never closed, because litellm does not own its lifecycle. +""" + +import asyncio + +import httpx +import pytest + +from litellm.caching.evicted_client_closer import EvictedClientCloser +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + +class FakeClock: + """Hand-advanced monotonic clock, so grace windows need no real waiting.""" + + def __init__(self) -> None: + self.now = 1000.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +def make_closer(clock: FakeClock, grace_seconds: float = 60.0) -> EvictedClientCloser: + return EvictedClientCloser(grace_seconds=grace_seconds, clock=clock) + + +async def _trickling_upstream(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + """Serves a chunked body slowly, so a request stays on the wire long enough to observe.""" + await reader.read(4096) + writer.write(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n") + await writer.drain() + for _ in range(6): + writer.write(b"5\r\nhello\r\n") + await writer.drain() + await asyncio.sleep(0.1) + writer.write(b"0\r\n\r\n") + await writer.drain() + + +@pytest.mark.asyncio +async def test_a_client_serving_a_request_is_not_closed_when_its_grace_window_ends(): + """The grace window on its own cannot promise that a request has finished. + + ``litellm.request_timeout`` defaults to 6000 seconds and a streaming response + is bounded only by how long the upstream keeps sending, so a client past its + deadline is closed only once its own pool reports nothing in flight. + """ + server = await asyncio.start_server(_trickling_upstream, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + clock = FakeClock() + closer = make_closer(clock) + client = httpx.AsyncClient() + + closer.mark_owned(client) + closer.schedule(client) + + async def read_the_stream() -> int: + received = 0 + async with client.stream("GET", f"http://127.0.0.1:{port}/") as response: + async for chunk in response.aiter_bytes(): + received += len(chunk) + return received + + streaming = asyncio.create_task(read_the_stream()) + await asyncio.sleep(0.25) # the request is on the wire + clock.advance(3600.0) # and its grace window is long gone + closer.reap() + await asyncio.sleep(0.05) + + assert client.is_closed is False, "closed a client that was serving a request" + assert await streaming > 0, "the in-flight request did not survive the reap" + + clock.advance(3600.0) + closer.reap() + await asyncio.sleep(0.05) + + assert client.is_closed is True, "an idle client past its grace window must be closed" + assert closer.pending_count == 0 + server.close() + + +@pytest.mark.asyncio +async def test_the_aiohttp_backed_handler_is_not_closed_mid_request(): + """The default async path is aiohttp-backed, whose pool accounts for its own leases.""" + server = await asyncio.start_server(_trickling_upstream, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + clock = FakeClock() + closer = make_closer(clock) + handler = AsyncHTTPHandler() + held_client = handler.client + + closer.mark_owned(handler) + closer.schedule(handler) + + request = asyncio.create_task(handler.get(f"http://127.0.0.1:{port}/")) + await asyncio.sleep(0.25) + clock.advance(3600.0) + closer.reap() + await asyncio.sleep(0.05) + + assert held_client.is_closed is False, "closed a handler that was serving a request" + assert (await request).status_code == 200 + + clock.advance(3600.0) + closer.reap() + await asyncio.sleep(0.05) + + assert held_client.is_closed is True + assert handler.client.is_closed is False, "a held handler must self-heal after its evicted client is closed" + server.close() diff --git a/tests/integration/caching/test_valkey_semantic_cache_real_io.py b/tests/integration/caching/test_valkey_semantic_cache_real_io.py new file mode 100644 index 00000000000..0bd3f8b47e2 --- /dev/null +++ b/tests/integration/caching/test_valkey_semantic_cache_real_io.py @@ -0,0 +1,32 @@ +import os +import subprocess +import sys +import textwrap + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")) + + +def test_importing_caching_does_not_require_redis(): + # redis is an optional dependency (extra_proxy), so the base SDK can be + # installed without it. Selecting valkey-semantic needs redis, but merely + # importing litellm.caching.caching must not, or `import litellm` breaks for + # every base-SDK user. This runs in a subprocess with redis blocked so the + # check is not polluted by redis already being imported in this session. + code = textwrap.dedent(""" + import sys + for name in ("redis", "redis.asyncio", "redis.commands", + "redis.commands.search"): + sys.modules[name] = None + import litellm.caching.caching # must not import redis at module top + from litellm.types.caching import LiteLLMCacheType + assert LiteLLMCacheType.VALKEY_SEMANTIC == "valkey-semantic" + print("ok") + """) + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + env={**os.environ, "PYTHONPATH": _REPO_ROOT}, + ) + assert result.returncode == 0, result.stderr + assert "ok" in result.stdout diff --git a/tests/integration/configuration/test_claude_settings_real_io.py b/tests/integration/configuration/test_claude_settings_real_io.py new file mode 100644 index 00000000000..5f2ed05021a --- /dev/null +++ b/tests/integration/configuration/test_claude_settings_real_io.py @@ -0,0 +1,272 @@ +import os +import pathlib +import shlex +import stat +import subprocess +import sys +import time + +import pytest +from click.testing import CliRunner +from litellm.proxy.client.cli.commands.claude_settings import ( + ApiKeyHelper, + KeepModel, + StartOn, + StaticToken, + UnpinModel, + install_statusline_script, + installed_executable, + lite_api_key_helper, + statusline_command, +) + + +class TestLiteApiKeyHelper: + @pytest.mark.skipif(sys.platform.startswith("win"), reason="runs the command through sh, as Claude Code does on POSIX") + def test_sh_runs_the_lite_on_path_against_this_proxy(self, tmp_path): + lite = tmp_path / "my tools" / "lite" + lite.parent.mkdir() + lite.write_text('#!/bin/sh\nprintf "%s\\n" "$0" "$@"\n') + lite.chmod(0o700) + + helper = lite_api_key_helper("https://proxy.example.com/", platform="linux", which=lambda name: str(lite)) + + assert helper is not None + ran = subprocess.run(["sh", "-c", helper.command], capture_output=True, text=True, check=True) + assert ran.stdout.splitlines() == [str(lite), "--base-url", "https://proxy.example.com", "auth", "print-token"] + + @pytest.mark.skipif(sys.platform.startswith("win"), reason="runs the command through sh, as Claude Code does on POSIX") + def test_a_lite_found_through_a_relative_path_entry_keeps_running_from_any_directory(self, tmp_path, monkeypatch): + for repo, name in ((tmp_path / "login-dir", "mine"), (tmp_path / "other-repo", "planted")): + lite = repo / "bin" / "lite" + lite.parent.mkdir(parents=True) + lite.write_text(f"#!/bin/sh\necho {name}\n") + lite.chmod(0o700) + monkeypatch.chdir(tmp_path / "login-dir") + + helper = lite_api_key_helper("https://proxy.example.com", platform="linux", which=lambda name: "bin/lite") + + assert helper is not None + ran = subprocess.run( + ["sh", "-c", helper.command], cwd=tmp_path / "other-repo", capture_output=True, text=True, check=True + ) + assert ran.stdout == "mine\n" + + +@pytest.mark.skipif(sys.platform.startswith("win"), reason="runs the command through sh, as Claude Code does on POSIX") +class TestInstalledExecutable: + @staticmethod + def _script(path, body): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(f"#!/bin/sh\n{body}\n") + path.chmod(0o700) + + def test_the_helper_runs_the_lite_installed_beside_this_python_over_a_shim_earlier_on_path( + self, tmp_path, monkeypatch + ): + self._script(tmp_path / "shims" / "lite", "echo 'pyenv: lite: command not found' >&2; exit 127") + self._script(tmp_path / "venv" / "bin" / "lite", "echo real") + monkeypatch.setenv("PATH", f"{tmp_path / 'shims'}{os.pathsep}{os.environ['PATH']}") + + helper = lite_api_key_helper( + "https://proxy.example.com", + platform="linux", + which=lambda name: installed_executable(name, scripts_dir=str(tmp_path / "venv" / "bin")), + ) + + assert helper is not None + ran = subprocess.run(["sh", "-c", helper.command], capture_output=True, text=True) + assert (ran.returncode, ran.stdout) == (0, "real\n") + + +PROXY = "http://127.0.0.1:4000" +ANTHROPIC = "https://api.anthropic.com" +RELOGIN = StaticToken("sk-fresh-login") +HELPER = ApiKeyHelper("/usr/local/bin/lite --base-url http://127.0.0.1:4000 auth print-token") +ORIGINAL = { + "theme": "dark", + "permissions": {"allow": ["Bash"]}, + "env": {"KEEP_ME": "1", "ANTHROPIC_API_KEY": "sk-ant-mine", "ANTHROPIC_BASE_URL": ANTHROPIC}, + "apiKeyHelper": "/usr/local/bin/lite auth print-token", + "model": "claude-opus-5", +} + + +def _set(path, value): + """A user edit: set (or with `_ABSENT`, remove) the key at a dotted path in the settings file.""" + + def edit(settings): + section, _, key = path.rpartition(".") + container = settings.setdefault(section, {}) if section else settings + if value is _ABSENT: + container.pop(key, None) + else: + container[key] = value + return settings + + return edit + + +_ABSENT = object() + + +# Each row: initial file, steps (configure kwargs dicts or edit callables) between the first configure +# and unconfigure, the expected file afterwards, and the expected outcome fields. Sequences that used +# to be one test each; the receipt's rules are what make them all come out right. +UNDO_SCENARIOS = { + "plain round trip": (ORIGINAL, [], ORIGINAL, {"kept": ()}), + "helper round trip": (ORIGINAL, [], ORIGINAL, {"kept": ()}, {"credential": HELPER, "model": KeepModel()}), + "helper with no file before": (None, [], None, {"file_removed": True}, {"credential": HELPER}), + "re-logins switching between the token and the helper keep the first snapshot": ( + ORIGINAL, + [{"credential": HELPER, "model": KeepModel()}, {"credential": RELOGIN, "model": KeepModel()}], + ORIGINAL, + {"kept": ()}, + {"credential": HELPER, "model": KeepModel()}, + ), + "no file before": (None, [], None, {"file_removed": True}), + "no env before": ({"theme": "dark"}, [], {"theme": "dark"}, {}), + "null env before": ({"theme": "dark", "env": None}, [], {"theme": "dark", "env": None}, {}), + "empty env before": ({"theme": "dark", "env": {}}, [], {"theme": "dark", "env": {}}, {}), + "user edits stay and are named": ( + ORIGINAL, + [_set("env.ENABLE_TOOL_SEARCH", "false"), _set("model", "claude-sonnet-4-6")], + {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "claude-sonnet-4-6"}, + {"kept": {"env.ENABLE_TOOL_SEARCH", "model"}, "withheld": ()}, + ), + "user filled an env configure created": (None, [_set("env.MY_VAR", "mine")], {"env": {"MY_VAR": "mine"}}, {}), + "user deleted the file": (None, [lambda s: None], None, {"file_removed": True, "restored": (), "kept": ()}), + "user removed our key: neither restored nor kept": ( + ORIGINAL, + [_set("env.ANTHROPIC_AUTH_TOKEN", _ABSENT)], + ORIGINAL, + {"not_restored": {"env.ANTHROPIC_AUTH_TOKEN"}, "kept": ()}, + ), + "restored names only what changed": ( + {"model": "claude-opus-5"}, + [], + {"model": "claude-opus-5"}, + { + "restored": { + "env.ANTHROPIC_BASE_URL", + "env.ANTHROPIC_AUTH_TOKEN", + "env.ENABLE_TOOL_SEARCH", + "env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", + "statusLine", + }, + "kept": (), + }, + {"credential": RELOGIN, "model": KeepModel()}, + ), + "repeat across credential kinds keeps the first snapshot": ( + ORIGINAL, + [ + {"credential": RELOGIN, "model": UnpinModel()}, + {"credential": StaticToken("sk-rotated"), "model": StartOn("claude-sonnet-4-6")}, + ], + ORIGINAL, + {}, + ), + "repeat without a model lets go of our pin, user had none": ({}, [{"model": UnpinModel()}], {}, {}), + "repeat without a model lets go of our pin, user had one": ( + {"model": "claude-opus-5"}, + [{"model": UnpinModel()}], + {"model": "claude-opus-5"}, + {}, + ), + "re-login keeps our pin": (None, [{"credential": RELOGIN, "model": KeepModel()}], None, {"file_removed": True}), + "edit between configures survives an unpin repeat": ( + ORIGINAL, + [ + _set("model", "my-favourite"), + _set("env.ENABLE_TOOL_SEARCH", "false"), + {"credential": RELOGIN, "model": UnpinModel()}, + ], + {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"}, + {"kept": {"env.ENABLE_TOOL_SEARCH", "model"}}, + ), + "edit between configures survives a re-login": ( + ORIGINAL, + [ + _set("model", "my-favourite"), + _set("env.ENABLE_TOOL_SEARCH", "false"), + {"credential": RELOGIN, "model": KeepModel()}, + ], + {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"}, + {"kept": {"env.ENABLE_TOOL_SEARCH", "model"}}, + ), + "edit between configures: a same-model repeat displaces it, so it is what comes back": ( + ORIGINAL, + [_set("model", "my-favourite"), _set("env.ENABLE_TOOL_SEARCH", "false"), {"credential": RELOGIN}], + {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"}, + {"kept": {"env.ENABLE_TOOL_SEARCH"}, "restored_includes": {"model"}}, + ), + "base URL changed since: credentials withheld, receipt kept": ( + ORIGINAL, + [_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000")], + {**ORIGINAL, "env": {"KEEP_ME": "1", "ANTHROPIC_BASE_URL": "http://other-proxy:4000"}, "apiKeyHelper": _ABSENT}, + { + "withheld": {("env.ANTHROPIC_API_KEY", ANTHROPIC), ("apiKeyHelper", ANTHROPIC)}, + "kept": {"env.ANTHROPIC_BASE_URL"}, + "receipt_kept": True, + }, + ), + "base URL changed and back: judged against the URL the restored file holds": ( + ORIGINAL, + [_set("env.ANTHROPIC_BASE_URL", ANTHROPIC)], + ORIGINAL, + {"withheld": ()}, + ), + "credential captured beside no URL goes back only beside no URL": ( + {"env": {"ANTHROPIC_API_KEY": "sk-default-endpoint"}}, + [_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000")], + {"env": {"ANTHROPIC_BASE_URL": "http://other-proxy:4000"}}, + { + "withheld": {("env.ANTHROPIC_API_KEY", "no ANTHROPIC_BASE_URL (Anthropic's default endpoint)")}, + "receipt_kept": True, + }, + ), + "restored document empty while a credential is withheld: file goes, receipt stays": ( + None, + [ + _set("env.ANTHROPIC_API_KEY", "sk-user"), + {"credential": RELOGIN, "model": KeepModel()}, + _set("env.ANTHROPIC_BASE_URL", _ABSENT), + ], + None, + {"withheld": {("env.ANTHROPIC_API_KEY", PROXY)}, "file_removed": True, "receipt_kept": True}, + {"credential": RELOGIN, "model": KeepModel()}, + ), + "a credential the user changed is kept, never also withheld": ( + ORIGINAL, + [_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000"), _set("apiKeyHelper", "/opt/mine/helper")], + { + **ORIGINAL, + "env": {"KEEP_ME": "1", "ANTHROPIC_BASE_URL": "http://other-proxy:4000"}, + "apiKeyHelper": "/opt/mine/helper", + }, + { + "withheld": {("env.ANTHROPIC_API_KEY", ANTHROPIC)}, + "kept": {"env.ANTHROPIC_BASE_URL", "apiKeyHelper"}, + "receipt_kept": True, + }, + ), +} + + +class TestStatusLine: + """Every configure registers the status line, and only while the slot is empty or already ours.""" + + COMMAND = "/opt/lite/bin/python /Users/me/.litellm/statusline.py" + + def test_the_installed_script_is_the_bundled_one_and_the_command_runs_this_interpreter(self, tmp_path): + from litellm.proxy.client.cli.commands import statusline_script + + script = tmp_path / "lite" / "statusline.py" + command = install_statusline_script(script) + assert script.read_bytes().split(b"\n", 1)[1] == pathlib.Path(statusline_script.__file__).read_bytes() + assert shlex.split(command) == [sys.executable, str(script)] + assert command == statusline_command(script) + assert stat.S_IMODE(script.stat().st_mode) == 0o600 + assert stat.S_IMODE(script.parent.stat().st_mode) == 0o700 + assert install_statusline_script(script) == command diff --git a/tests/integration/configuration/test_main_real_io.py b/tests/integration/configuration/test_main_real_io.py new file mode 100644 index 00000000000..1181ee1eec3 --- /dev/null +++ b/tests/integration/configuration/test_main_real_io.py @@ -0,0 +1,20 @@ +import os +import subprocess +import sys +from typing import Final + + +def test_importing_proxy_hooks_does_not_mutate_litellm_callbacks(): + script: Final = ( + "import litellm; " + "litellm.callbacks = []; " + "import litellm.proxy.hooks; " + "assert len(litellm.callbacks) == 0, f'mutated: {litellm.callbacks}'" + ) + result: Final = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, + ) + assert result.returncode == 0, f"stdout: {result.stdout}\nstderr: {result.stderr}" diff --git a/tests/integration/configuration/test_runtime_real_io.py b/tests/integration/configuration/test_runtime_real_io.py new file mode 100644 index 00000000000..23ca074ef7b --- /dev/null +++ b/tests/integration/configuration/test_runtime_real_io.py @@ -0,0 +1,83 @@ +import signal +import subprocess +import sys +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import anyio +import pytest + +from litellm.proxy.rust_forward.platform import Platform, current_platform +from litellm.proxy.rust_forward.platform.windows import WindowsPlatform +from litellm.proxy.rust_forward.supervisor import supervise +from litellm.proxy.rust_forward.supervisor.children import ChildSpec +from litellm.proxy.rust_forward.supervisor.lifecycle import Child, ChildExited +from litellm.proxy.rust_forward.supervisor.runtime import ( + SpawnFailed, + exit_like, + open_children, + spawn_subprocess, +) +from tests.unit.proxy.rust_forward.supervisor._support import python_server + + +PYTHON_CHILD: Final = """ +import json, os, pathlib, socket, time +time.sleep(0.5) +pathlib.Path(os.environ["MARKERS"], "listening").touch() +private = json.loads(os.environ["_LITELLM_GATEWAY_SUPERVISED"])["listener"] +listener = socket.socket(fileno=private["fd"]) if "fd" in private else socket.create_server(("127.0.0.1", private["port"])) +listener.listen() +time.sleep(10) +""" +RUST_CHILD: Final = """ +import os, pathlib +markers = pathlib.Path(os.environ["MARKERS"]) +(markers / "rust_saw_python_listening").write_text(str((markers / "listening").exists())) +""" + + +def _fake_gateway(scripts: Path) -> Platform: + if sys.platform != "win32": + gateway: Final = scripts / "litellm-gateway" + gateway.write_text(f"#!{sys.executable} -I\n{RUST_CHILD}") + gateway.chmod(0o755) + return current_platform() + script: Final = scripts / "rust_child.py" + script.write_text(RUST_CHILD) + (scripts / "litellm-gateway.cmd").write_text(f'@"{sys.executable}" -I "{script}"\r\n') + return WindowsPlatform(gateway_name="litellm-gateway.cmd") + + +def test_rust_opens_the_public_port_only_once_python_listens(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("MARKERS", str(tmp_path)) + command: Final = (sys.executable, "-I", "-c", PYTHON_CHILD) + supervise(python_server(), command, scripts_dir=str(tmp_path), platform=_fake_gateway(tmp_path)) + assert (tmp_path / "rust_saw_python_listening").read_text() == "True" +POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="the child traps SIGTERM") +RUST_MISSING: Final = SpawnFailed(FileNotFoundError("litellm-gateway")) + + +def _children(argv: tuple[str, ...]) -> MappingProxyType[Child, ChildSpec]: + spec: Final = ChildSpec(argv, MappingProxyType({}), current_platform().spawn_options(())) + return MappingProxyType({Child.PYTHON: spec, Child.RUST: spec}) + + +async def test_exiting_children_are_reported_as_events() -> None: + async with open_children( + current_platform(), spawn_subprocess, _children((sys.executable, "-c", "exit(7)")), 1 + ) as children: + assert await children.spawn(Child.RUST) is None + with anyio.fail_after(5): + assert await children.next_event() == ChildExited(Child.RUST, 7) + + +@POSIX_ONLY +@pytest.mark.parametrize("signum", [signal.SIGTERM, signal.SIGKILL]) +def test_the_supervisor_dies_by_the_signal_that_killed_python(signum: signal.Signals) -> None: + script: Final = ( + "from litellm.proxy.rust_forward.supervisor.runtime import exit_like\n" + f"raise SystemExit(exit_like({-signum}, frozenset()))" + ) + assert subprocess.run((sys.executable, "-I", "-c", script), check=False, timeout=60).returncode == -signum diff --git a/tests/integration/mcp/test_litellm_proxy_mcp_handler_real_io.py b/tests/integration/mcp/test_litellm_proxy_mcp_handler_real_io.py new file mode 100644 index 00000000000..9fda7608e26 --- /dev/null +++ b/tests/integration/mcp/test_litellm_proxy_mcp_handler_real_io.py @@ -0,0 +1,88 @@ +import subprocess +import sys +import textwrap +from types import MappingProxyType +from typing import Final +from mcp.types import Tool as MCPTool +from litellm.types.mcp_server.mcp_toolset import MCPToolset +from litellm.types.mcp import MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +def test_completion_with_function_tools_works_without_fastapi_installed(): + script = textwrap.dedent( + """ + import sys + + class _FastapiBlocker: + def find_spec(self, fullname, path=None, target=None): + if fullname == "fastapi" or fullname.startswith("fastapi."): + raise ModuleNotFoundError("No module named 'fastapi'") + return None + + sys.meta_path.insert(0, _FastapiBlocker()) + + import litellm + + response = litellm.completion( + model="openai/gpt-5.5", + messages=[{"role": "user", "content": "What is the weather in SF?"}], + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a location", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + }, + } + ], + mock_response="sunny", + ) + assert response.choices[0].message.content == "sunny" + """ + ) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stderr + + +TOOLSET_ROWS: Final = ( + MCPToolset( + toolset_id="ts-granted", toolset_name="team-toolset", tools=[{"server_id": "srv-1", "tool_name": "add"}] + ), + MCPToolset( + toolset_id="ts-other", toolset_name="other-toolset", tools=[{"server_id": "srv-2", "tool_name": "lookup"}] + ), + MCPToolset(toolset_id="ts-shouty", toolset_name="CALC", tools=[{"server_id": "srv-2", "tool_name": "lookup"}]), + MCPToolset(toolset_id="ts-group", toolset_name="reference", tools=[{"server_id": "srv-1", "tool_name": "add"}]), + MCPToolset( + toolset_id="ts-prefixed", + toolset_name="prefixed-toolset", + tools=[{"server_id": "srv-1", "tool_name": "add"}, {"server_id": "srv-1", "tool_name": "calc-subtract"}], + ), + MCPToolset( + toolset_id="ts-subtract", toolset_name="subtract-toolset", tools=[{"server_id": "srv-1", "tool_name": "subtract"}] + ), +) +GATEWAY_SERVERS: Final = ( + MCPServer(server_id="srv-1", name="calc", alias="calc", transport=MCPTransport.http, access_groups=["wiki"]), + MCPServer(server_id="srv-2", name="docs", alias="docs", transport=MCPTransport.http, access_groups=["reference"]), +) +UPSTREAM_TOOLS: Final = MappingProxyType( + { + "calc": ( + MCPTool(name="add", inputSchema={"type": "object"}), + MCPTool(name="subtract", inputSchema={"type": "object"}), + ), + "docs": (MCPTool(name="lookup", inputSchema={"type": "object"}),), + } +) diff --git a/tests/integration/mcp/test_mcp_client_real_io.py b/tests/integration/mcp/test_mcp_client_real_io.py new file mode 100644 index 00000000000..f957c6d71e5 --- /dev/null +++ b/tests/integration/mcp/test_mcp_client_real_io.py @@ -0,0 +1,1527 @@ +import asyncio +import json +import sys +from collections.abc import AsyncIterator, Callable +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from pathlib import Path +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import anyio +import httpx2 +import pytest +from mcp import ClientSession, MCPError +from mcp.client.streamable_http import streamable_http_client +from mcp.server.mcpserver import Context +from mcp.server.mcpserver import MCPServer as UpstreamServer +from mcp.server.transport_security import TransportSecuritySettings +from mcp.types import ( + CONNECTION_CLOSED, + CallToolRequestParams, + CallToolResult, + DiscoverResult, + JSONRPCMessage, + JSONRPCRequest, + LoggingMessageNotificationParams, + ServerCapabilities, +) +from pydantic import TypeAdapter, ValidationError +from starlette.applications import Starlette +from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint +from starlette.requests import Request +from starlette.responses import Response +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +# Add the parent directory to the path so we can import litellm +import litellm.mcp_client.client as mcp_client_module +from litellm.mcp_client.client import ( + MCPClient, + PersistentMCPSession, + UpstreamSessionClosedError, + _TransportContext, + as_mcp_read_timeout, +) +from litellm.types.mcp import MCPTransport +from tests._support.clock_loops import AutojumpClockLoop, ManualClockLoop + +_JSONRPC_MESSAGE_ADAPTER: Final = TypeAdapter(JSONRPCMessage) + + +class _MockTransportClient(MCPClient): + """An MCPClient whose streamable-HTTP transport runs on an httpx2 MockTransport.""" + + def __init__( + self, + respond, + *, + http_transport: httpx2.AsyncBaseTransport | None = None, + transport_context: Callable[[str, httpx2.AsyncClient], _TransportContext] | None = None, + **kwargs, + ): + super().__init__(**kwargs) + self._respond = respond + self._http_transport = http_transport + self._transport_context = transport_context + + def _create_transport_context(self) -> tuple[_TransportContext, httpx2.AsyncClient]: + transport: Final = self._http_transport or httpx2.MockTransport(self._respond) + http_client: Final = self._create_httpx_client_factory(transport=transport)( + headers=self._get_auth_headers(), timeout=httpx2.Timeout(self.timeout) + ) + if self._transport_context is not None: + return self._transport_context(self.server_url, http_client), http_client + return streamable_http_client(self.server_url, http_client=http_client), http_client + + +# --------------------------------------------------------------------------- +# _last_initialize_instructions capture +# --------------------------------------------------------------------------- + + +# --------------------------------------------------------------------------- +# Transport error surfacing +# --------------------------------------------------------------------------- + + +if __name__ == "__main__": + pytest.main([__file__]) + + +@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 + + +class _DiagnosticSSEStream(httpx2.AsyncByteStream): + def __init__(self, messages: asyncio.Queue[bytes | Exception | None]) -> None: + self.messages = messages + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b"event: endpoint\ndata: /messages\n\n" + while True: + message: Final = await self.messages.get() + if message is None: + return + if isinstance(message, Exception): + raise message + yield b"event: message\ndata: " + message + b"\n\n" + + +_DIAGNOSTIC_STDIO_SERVER: Final = """ +import json, sys +mode, failure_method = sys.argv[1:] +for line in sys.stdin: + request = json.loads(line) + if "method" not in request or "id" not in request: + continue + if request["method"] == failure_method: + if mode == "bad-json": + print("secret-invalid-json", flush=True) + continue + if mode == "closed": + sys.exit(0) + if mode == "silent": + print(json.dumps({"jsonrpc": "2.0", "method": "notifications/message", "params": {"level": "info", "data": "Waiting"}}), flush=True) + continue + if request["method"] == "initialize": + result = {"protocolVersion": request["params"]["protocolVersion"], "capabilities": {"tools": {}, "logging": {}}, "serverInfo": {"name": "diagnostic", "version": "1"}} + elif request["method"] == "tools/list": + print(json.dumps({"jsonrpc": "2.0", "method": "notifications/message", "params": {"level": "info", "data": "Listing tools"}}), flush=True) + print(json.dumps({"jsonrpc": "2.0", "id": "unmatched", "result": {}}), flush=True) + print(json.dumps({"jsonrpc": "2.0", "id": "server-ping", "method": "ping"}), flush=True) + result = {"tools": [{"name": "ping", "inputSchema": {"type": "object"}}]} + else: + result = {"content": [{"type": "text", "text": "pong"}], "isError": False} + print(json.dumps({"jsonrpc": "2.0", "id": request["id"], "result": result}), flush=True) +""" + + +def _diagnostic_transport(transport: MCPTransport, mode: str, failure_method: str) -> _TransportContext: + from mcp import StdioServerParameters + from mcp.client.sse import sse_client + from mcp.client.stdio import stdio_client + + if transport == MCPTransport.stdio: + return stdio_client( + StdioServerParameters( + command=sys.executable, args=["-u", "-c", _DIAGNOSTIC_STDIO_SERVER, mode, failure_method] + ) + ) + messages: Final[asyncio.Queue[bytes | Exception | None]] = asyncio.Queue() + + async def respond(request: httpx2.Request) -> httpx2.Response: + if request.method == "GET": + 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 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(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 httpx2.Response(202) + if payload["method"] == "tools/list": + for message in ( + { + "jsonrpc": "2.0", + "method": "notifications/message", + "params": {"level": "info", "data": "Listing tools"}, + }, + {"jsonrpc": "2.0", "id": "unmatched", "result": {}}, + {"jsonrpc": "2.0", "id": "server-ping", "method": "ping"}, + ): + await messages.put(json.dumps(message).encode()) + result: Final = ( + { + "protocolVersion": payload["params"]["protocolVersion"], + "capabilities": {"tools": {}, "logging": {}}, + "serverInfo": {"name": "diagnostic", "version": "1"}, + } + if payload["method"] == "initialize" + else {"tools": [{"name": "ping", "inputSchema": {"type": "object"}}]} + if payload["method"] == "tools/list" + else {"content": [{"type": "text", "text": "pong"}], "isError": False} + ) + await messages.put(json.dumps({"jsonrpc": "2.0", "id": payload["id"], "result": result}).encode()) + return httpx2.Response(202) + + def factory( + headers: dict[str, str] | None = None, + 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) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) +@pytest.mark.parametrize("failure_method", ["initialize", "tools/list"]) +async def test_transport_parsing_failure_is_preserved(transport: MCPTransport, failure_method: str) -> None: + client: Final = MCPClient(server_url="https://example.com/sse", transport_type=transport, timeout=0.2) + with pytest.raises(ValidationError): + await asyncio.wait_for( + client._execute_session_operation( + _diagnostic_transport(transport, "bad-json", failure_method), lambda session: session.list_tools() + ), + timeout=3, + ) + + +def test_sse_read_failure_is_preserved() -> None: + loop: Final = AutojumpClockLoop() + + async def run() -> None: + client: Final = MCPClient(server_url="https://example.com/sse", transport_type=MCPTransport.sse, timeout=0.2) + 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(), + ), + timeout=3, + ) + + try: + loop.run_until_complete(run()) + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.run_until_complete(loop.shutdown_default_executor()) + loop.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("protocol_version", ["auto", "2025-06-18"]) +@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) +@pytest.mark.parametrize("mode", ["ok", "closed", "silent"]) +async def test_transport_completion_and_normal_messages( + transport: MCPTransport, mode: str, protocol_version: str +) -> None: + from mcp import ClientSession + + from litellm.proxy.mcp_server.rest_endpoints import _connection_error_message + + logging_callback: Final = AsyncMock() + read_timeout: Final = 0.2 if mode == "silent" else 30 + client: Final = MCPClient( + server_url="https://example.com/sse", + transport_type=transport, + timeout=read_timeout, + logging_callback=logging_callback, + protocol_version=protocol_version, + ) + + async def operation(session: ClientSession) -> CallToolResult: + tools: Final = await session.list_tools() + assert [tool.name for tool in tools.tools] == ["ping"] + return await session.call_tool("ping", {}) + + 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.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: + await asyncio.wait_for(pending, timeout=3) + if mode == "closed": + assert "connection was closed" in _connection_error_message(caught.value, client.server_url, read_timeout) + else: + assert isinstance(as_mcp_read_timeout(caught.value), TimeoutError) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) +async def test_transport_cancellation_cleans_up_a_pending_request(transport: MCPTransport) -> None: + ready: Final = asyncio.Event() + + async def on_log(message: LoggingMessageNotificationParams) -> None: + if message.data == "Waiting": + ready.set() + + client: Final = MCPClient( + server_url="https://example.com/sse", transport_type=transport, timeout=30, logging_callback=on_log + ) + task: Final = asyncio.create_task( + client._execute_session_operation( + _diagnostic_transport(transport, "silent", "tools/list"), lambda session: session.list_tools() + ) + ) + try: + await asyncio.wait_for(ready.wait(), timeout=3) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=3) + + +@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list")) +@pytest.mark.parametrize( + "failure", ("repeat", "cycle", "cap", "method_not_found", "internal_error", "unauthorized", "deadline") +) +@pytest.mark.parametrize("strict", (False, True)) +def test_optional_discovery_rejects_incomplete_walks( + method: str, failure: str, strict: bool, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.setattr( + mcp_client_module, + "MCP_TOOL_LISTING_MAX_PAGES", + 3 if failure in ("cycle", "repeat") else 2, + raising=False, + ) + monkeypatch.setattr(mcp_client_module, "MCP_TOOL_LISTING_TIMEOUT", 0.05) + field: Final = { + "prompts/list": "prompts", + "resources/list": "resources", + "resources/templates/list": "resourceTemplates", + }[method] + entry: Final = { + "prompts/list": {"name": "first"}, + "resources/list": {"name": "first", "uri": "test://first"}, + "resources/templates/list": {"name": "first", "uriTemplate": "test://{name}"}, + }[method] + loop: Final = ManualClockLoop() + + async def run() -> None: + cancelled: Final = asyncio.Event() + + async def respond(request: httpx2.Request) -> httpx2.Response: + payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) + if not isinstance(payload, JSONRPCRequest): + return httpx2.Response(202) + if payload.method == "initialize": + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "protocolVersion": (payload.params or {})["protocolVersion"], + "capabilities": {"prompts": {}, "resources": {}}, + "serverInfo": {"name": "interrupted", "version": "1"}, + }, + }, + ) + assert payload.method == method + cursor: Final = (payload.params or {}).get("cursor") + if cursor is not None: + if failure == "deadline": + loop.advance(0.15) + try: + for _ in range(1_000): + await asyncio.sleep(0) + except asyncio.CancelledError: + cancelled.set() + raise + return httpx2.Response(500) + if failure == "unauthorized": + return httpx2.Response(401) + if failure in ("method_not_found", "internal_error"): + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "error": { + "code": -32601 if failure == "method_not_found" else -32603, + "message": "Later page unavailable", + }, + }, + ) + if failure == "deadline" and cursor is None: + loop.advance(0.1) + next_cursor: Final = ( + "private-cursor-2" if cursor == "private-cursor-1" and failure != "repeat" else "private-cursor-1" + ) + return httpx2.Response( + 200, json={"jsonrpc": "2.0", "id": payload.id, "result": {field: [entry], "nextCursor": next_cursor}} + ) + + responder: Final = AsyncMock(side_effect=respond) + client: Final = _MockTransportClient(responder, server_url="https://example.com/mcp", timeout=0.2) + operation: Final = { + "prompts/list": client.list_prompts, + "resources/list": client.list_resources, + "resources/templates/list": client.list_resource_templates, + }[method] + if strict: + error_type: Final = { + "internal_error": MCPError, + "unauthorized": httpx2.HTTPStatusError, + "deadline": TimeoutError, + }.get(failure, RuntimeError) + with pytest.raises(error_type): + await operation(raise_on_error=True) + else: + assert await operation() == [] + assert len( + tuple( + payload + for call in responder.call_args_list + if isinstance(payload := _JSONRPC_MESSAGE_ADAPTER.validate_json(call.args[0].content), JSONRPCRequest) + and payload.method == method + ) + ) == (3 if failure == "cycle" else 2) + assert "private-cursor" not in caplog.text + if failure == "deadline": + assert cancelled.is_set() + + try: + loop.run_until_complete(run()) + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.run_until_complete(loop.shutdown_default_executor()) + loop.close() + + +def test_client_import_before_proxy_credentials_succeeds_in_fresh_process(): + import subprocess + + result = subprocess.run( + [ + sys.executable, + "-c", + "import litellm.mcp_client.client; from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager; print(MCPServerManager.__name__)", + ], + capture_output=True, + text=True, + timeout=60, + check=False, + ) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "MCPServerManager" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("protocol_version", ("auto", "2025-06-18")) +@pytest.mark.parametrize("cancel_mode", ("scope", "task", "wait_for", "read_timeout")) +@pytest.mark.parametrize("concurrency", (1, 5)) +@pytest.mark.parametrize("termination", ("ok", "hang", "hang_body")) +@pytest.mark.parametrize("raise_on_error", (False, True)) +async def test_cancellation_delivers_termination_over_tcp( + cancel_mode: str, concurrency: int, termination: str, raise_on_error: bool, protocol_version: str +) -> None: + started: Final = asyncio.Event() + scope_ready: Final[asyncio.Future[anyio.CancelScope]] = asyncio.get_running_loop().create_future() + terminations: Final[list[bytes]] = [] + starts: Final[list[bytes]] = [] + stop: Final = asyncio.Event() + connections: Final[list[asyncio.Task[None]]] = [] + + async def handle_connection(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + connection: Final = asyncio.current_task() + assert connection is not None + connections.append(connection) + try: + request_line: Final = await reader.readline() + if not request_line: + return + method: Final = request_line.split()[0] + headers: Final = await reader.readuntil(b"\r\n\r\n") + length: Final = next( + ( + int(line.split(b":", 1)[1]) + for line in headers.splitlines() + if line.lower().startswith(b"content-length:") + ), + 0, + ) + try: + body: Final = await reader.readexactly(length) + except asyncio.IncompleteReadError: + return + if method == b"DELETE": + terminations.append(body) + if termination != "ok": + await stop.wait() + return + writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") + elif method == b"GET": + writer.write(b"HTTP/1.1 405 Method Not Allowed\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") + else: + payload: Final = json.loads(body) + if payload["method"] == "tools/call": + if termination == "hang_body": + writer.write( + b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n" + b"Content-Length: 200\r\nConnection: close\r\n\r\n" + ) + await writer.drain() + starts.append(body) + if len(starts) == concurrency: + started.set() + await stop.wait() + return + if payload["method"] == "initialize": + if cancel_mode != "read_timeout": + await asyncio.sleep(0.75) + response: Final = json.dumps( + { + "jsonrpc": "2.0", + "id": payload["id"], + "result": { + "protocolVersion": "2025-06-18", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "tcp-peer", "version": "1"}, + }, + } + ).encode() + writer.write( + b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nMcp-Session-Id: tcp-session\r\n" + + f"Content-Length: {len(response)}\r\nConnection: close\r\n\r\n".encode() + + response + ) + else: + writer.write(b"HTTP/1.1 202 Accepted\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") + await writer.drain() + finally: + writer.close() + await writer.wait_closed() + + listener: Final = await asyncio.start_server(handle_connection, "127.0.0.1", 0) + port: Final = listener.sockets[0].getsockname()[1] + client_timeout: Final = 2 if cancel_mode == "read_timeout" else 30 + client: Final = MCPClient( + server_url=f"http://127.0.0.1:{port}/mcp", + protocol_version=protocol_version, + timeout=client_timeout, + ) + + async def calls(): + results: Final = await asyncio.gather( + *( + client.call_tool(CallToolRequestParams(name="slow", arguments={}), raise_on_error=raise_on_error) + for _ in range(concurrency) + ), + return_exceptions=cancel_mode == "read_timeout", + ) + if cancel_mode == "read_timeout": + if raise_on_error: + assert all(isinstance(result, TimeoutError) for result in results) + else: + assert all(isinstance(result, CallToolResult) and result.is_error for result in results) + return results + + async def invoke(): + if cancel_mode == "scope": + with anyio.fail_after(None) as scope: + scope_ready.set_result(scope) + return await calls() + return await calls() + + try: + task: Final = asyncio.create_task(invoke()) + await asyncio.wait_for(started.wait(), 30) + if cancel_mode == "scope": + (await scope_ready).deadline = anyio.current_time() + 0.2 + if cancel_mode == "task": + task.cancel() + expected_error: Final = ( + TimeoutError + if cancel_mode == "read_timeout" + else asyncio.CancelledError + if cancel_mode == "task" + else TimeoutError + ) + if cancel_mode == "read_timeout": + done, _ = await asyncio.wait((task,), timeout=8) + assert task in done, "Read timeout and bounded cleanup must complete without external cancellation" + await task + elif cancel_mode == "wait_for": + with pytest.raises(expected_error): + await asyncio.wait_for(task, 0.2) + else: + with pytest.raises(expected_error): + await task + assert len(starts) == concurrency + assert len(terminations) == concurrency, "Each cancelled call must send DELETE over a fresh TCP connection" + finally: + stop.set() + if not task.done(): + task.cancel() + await asyncio.wait((task,), timeout=8) + listener.close() + for connection in connections: + connection.cancel() + closed: Final = await asyncio.wait_for(asyncio.gather(*connections, return_exceptions=True), 2) + assert all(result is None or isinstance(result, asyncio.CancelledError) for result in closed), closed + await asyncio.wait_for(listener.wait_closed(), 2) + + +def test_modern_call_uses_discovery_deadline_for_multiple_pages() -> None: + from mcp.types import ToolsCapability + + loop: Final = AutojumpClockLoop() + + async def respond(request: httpx2.Request) -> httpx2.Response: + payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) + assert isinstance(payload, JSONRPCRequest) + if payload.method == "server/discover": + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": DiscoverResult( + supported_versions=["2026-07-28"], capabilities=ServerCapabilities(tools=ToolsCapability()) + ).model_dump(by_alias=True, exclude_none=True), + }, + ) + if payload.method == "tools/list": + await asyncio.sleep(0.06) + assert payload.params is not None + page: Final = ( + {"tools": [], "nextCursor": "second"} + if "cursor" not in payload.params + else { + "tools": [ + { + "name": "quote", + "inputSchema": { + "type": "object", + "properties": {"workspace": {"type": "string", "x-mcp-header": "Workspace"}}, + }, + } + ] + } + ) + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": {"resultType": "complete", "cacheScope": "private", "ttlMs": 0, **page}, + }, + ) + assert payload.method == "tools/call" + assert request.headers["mcp-param-workspace"] == "engineering" + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": {"resultType": "complete", "content": [{"type": "text", "text": "quoted"}], "isError": False}, + }, + ) + + async def run() -> None: + client: Final = _MockTransportClient( + respond, server_url="https://example.com/mcp", protocol_version="2026-07-28", timeout=0.1 + ) + assert [tool.name for tool in await client.list_tools(raise_on_error=True)] == ["quote"] + result: Final = await client.call_tool( + CallToolRequestParams(name="quote", arguments={"workspace": "engineering"}), raise_on_error=True + ) + assert result.is_error is False + assert result.content[0].text == "quoted" + + try: + loop.run_until_complete(run()) + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + + +class _StatefulUpstreamClient(MCPClient): + """An MCPClient whose streamable-HTTP transport talks to an in-process stateful MCP server.""" + + def __init__(self, app: ASGIApp, *, server_url: str, transport_type: MCPTransport) -> None: + super().__init__(server_url=server_url, transport_type=transport_type) + self._app = app + + def _create_transport_context(self) -> tuple[_TransportContext, httpx2.AsyncClient]: + http_client: Final = self._create_httpx_client_factory(transport=httpx2.ASGITransport(app=self._app))( + headers=self._get_auth_headers(), timeout=httpx2.Timeout(self.timeout) + ) + return streamable_http_client(self.server_url, http_client=http_client), http_client + + +class _CleanupHoldingASGITransport(httpx2.ASGITransport): + def __init__(self, app: ASGIApp, cleanup_started: asyncio.Event, release_cleanup: asyncio.Event) -> None: + super().__init__(app=app) + self._cleanup_started: Final = cleanup_started + self._release_cleanup: Final = release_cleanup + + async def aclose(self) -> None: + self._cleanup_started.set() + await self._release_cleanup.wait() + await super().aclose() + + +class _CleanupHoldingStatefulUpstreamClient(_StatefulUpstreamClient): + def __init__( + self, + app: ASGIApp, + *, + server_url: str, + transport_type: MCPTransport, + cleanup_started: asyncio.Event, + release_cleanup: asyncio.Event, + ) -> None: + super().__init__(app, server_url=server_url, transport_type=transport_type) + self._cleanup_started: Final = cleanup_started + self._release_cleanup: Final = release_cleanup + + def _create_transport_context(self) -> tuple[_TransportContext, httpx2.AsyncClient]: + transport: Final = _CleanupHoldingASGITransport( + self._app, + self._cleanup_started, + self._release_cleanup, + ) + http_client: Final = self._create_httpx_client_factory(transport=transport)( + headers=self._get_auth_headers(), + timeout=httpx2.Timeout(self.timeout), + ) + return streamable_http_client(self.server_url, http_client=http_client), http_client + + +def _stateful_upstream() -> Starlette: + service: Final = UpstreamServer("stateful") + selected: Final[dict[str, str]] = {} + + @service.tool() + def select_project(name: str, ctx: Context) -> str: + selected[ctx.headers["mcp-session-id"]] = name + return f"selected {name}" + + @service.tool() + def create_feature(title: str, ctx: Context) -> str: + return f"{selected[ctx.headers['mcp-session-id']]}/{title}" + + return service.streamable_http_app( + stateless_http=False, + json_response=True, + transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False), + ) + + +def _upstream_with_status_for_method(method: str, status_code: int) -> Starlette: + service: Final = UpstreamServer("status upstream") + + @service.tool() + def protected_tool() -> str: + return "protected" + + app: Final = service.streamable_http_app( + stateless_http=False, + json_response=True, + transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False), + ) + + async def reject_method(request: Request, call_next: RequestResponseEndpoint) -> Response: + if request.method == "POST": + payload: Final = await request.json() + if isinstance(payload, dict) and payload.get("method") == method: + return Response(status_code=status_code, headers={"WWW-Authenticate": "Bearer"}) + return await call_next(request) + + app.add_middleware(BaseHTTPMiddleware, dispatch=reject_method) + return app + + +def _upstream_with_malformed_response_for_method(method: str) -> Starlette: + service: Final = UpstreamServer("malformed response upstream") + + @service.tool() + def malformed_tool() -> str: + return "unreachable" + + app: Final = service.streamable_http_app( + stateless_http=False, + json_response=True, + transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False), + ) + + async def malformed_method(request: Request, call_next: RequestResponseEndpoint) -> Response: + if request.method == "GET": + return Response(content="data: {}\n\n", media_type="text/event-stream") + if request.method == "POST": + payload: Final = await request.json() + if isinstance(payload, dict) and payload.get("method") == method: + return Response(content="not a JSON-RPC message", media_type="text/event-stream") + return await call_next(request) + + app.add_middleware(BaseHTTPMiddleware, dispatch=malformed_method) + return app + + +def _client_with_session( + app: Starlette, + admission: Callable[[], AbstractAsyncContextManager[None]] | None = None, +) -> tuple[_StatefulUpstreamClient, PersistentMCPSession]: + client: Final = _StatefulUpstreamClient(app, server_url="http://upstream/mcp", transport_type=MCPTransport.http) + return client, client.open_persistent_session(admission) + + +@pytest.mark.asyncio +async def test_cancelling_one_waiter_for_session_readiness_does_not_cancel_other_waiters() -> None: + app: Final = _stateful_upstream() + async with app.router.lifespan_context(app): + initialize_started: Final = asyncio.Event() + release_initialize: Final = asyncio.Event() + + async def gated_server(scope: Scope, receive: Receive, send: Send) -> None: + async def gated_send(message: Message) -> None: + if message["type"] == "http.response.body" and not initialize_started.is_set(): + initialize_started.set() + await release_initialize.wait() + await send(message) + + await app(scope, receive, gated_send) + + gated_app: Final[ASGIApp] = gated_server + client: Final = _StatefulUpstreamClient( + gated_app, server_url="http://upstream/mcp", transport_type=MCPTransport.http + ) + session: Final = client.open_persistent_session() + + async def operation(_: ClientSession) -> str: + return "ready" + + cancelled_waiter: Final = asyncio.create_task(session.run(operation)) + remaining_waiter: Final = asyncio.create_task(session.run(operation)) + try: + await asyncio.wait_for(initialize_started.wait(), 5) + await asyncio.sleep(0) + cancelled_waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled_waiter + release_initialize.set() + assert await asyncio.wait_for(remaining_waiter, 5) == "ready" + assert not session.closed + finally: + release_initialize.set() + session.close() + await asyncio.wait_for(session.wait_closed(), 5) + + +@pytest.mark.asyncio +async def test_persistent_session_survives_recoverable_mcp_errors_and_closes_on_connection_closed() -> None: + app: Final = _stateful_upstream() + async with app.router.lifespan_context(app): + client, session = _client_with_session(app) + + async def invalid_operation(_: ClientSession) -> str: + raise MCPError(code=-32602, message="invalid params") + + try: + with pytest.raises(MCPError): + await session.run(invalid_operation) + assert not session.closed + selected: Final = await client.call_tool( + CallToolRequestParams(name="select_project", arguments={"name": "recoverable"}), + persistent_session=session, + ) + created: Final = await client.call_tool( + CallToolRequestParams(name="create_feature", arguments={"title": "still-open"}), + persistent_session=session, + ) + finally: + session.close() + assert selected.is_error is False + assert created.content[0].text == "recoverable/still-open" + await asyncio.wait_for(session.wait_closed(), 5) + + connection_app: Final = _stateful_upstream() + async with connection_app.router.lifespan_context(connection_app): + _, session = _client_with_session(connection_app) + started: Final = asyncio.Event() + release: Final = asyncio.Event() + + async def disconnected_operation(_: ClientSession) -> str: + started.set() + await release.wait() + raise MCPError(code=CONNECTION_CLOSED, message="upstream connection closed") + + async def pending_operation(_: ClientSession) -> str: + return "unused" + + active: Final = asyncio.create_task(session.run(disconnected_operation)) + await asyncio.wait_for(started.wait(), 5) + pending: Final = asyncio.create_task(session.run(pending_operation)) + await asyncio.sleep(0) + release.set() + with pytest.raises(MCPError) as caught: + await asyncio.wait_for(active, 5) + assert caught.value.error.code == CONNECTION_CLOSED + with pytest.raises((MCPError, UpstreamSessionClosedError)): + await asyncio.wait_for(pending, 5) + await asyncio.wait_for(session.wait_closed(), 5) + assert session.closed + + +@pytest.mark.asyncio +async def test_ended_persistent_session_is_not_reusable_during_transport_cleanup() -> None: + app: Final = _stateful_upstream() + async with app.router.lifespan_context(app): + cleanup_started: Final = asyncio.Event() + release_cleanup: Final = asyncio.Event() + client: Final = _CleanupHoldingStatefulUpstreamClient( + app, + server_url="http://upstream/mcp", + transport_type=MCPTransport.http, + cleanup_started=cleanup_started, + release_cleanup=release_cleanup, + ) + session: Final = client.open_persistent_session() + + async def disconnected_operation(_: ClientSession) -> str: + raise MCPError(code=CONNECTION_CLOSED, message="upstream connection closed") + + try: + with pytest.raises(MCPError) as caught: + await asyncio.wait_for(session.run(disconnected_operation), 5) + assert caught.value.error.code == CONNECTION_CLOSED + await asyncio.wait_for(cleanup_started.wait(), 5) + + assert not session.reusable + assert not session.closed + session.retire() + + release_cleanup.set() + await asyncio.wait_for(session.wait_closed(), 5) + assert session.closed + finally: + release_cleanup.set() + if not session.closed: + session.close() + await asyncio.wait_for(session.wait_closed(), 5) + + +@pytest.mark.asyncio +async def test_persistent_session_remaps_stream_error_over_mcp_error() -> None: + app: Final = _stateful_upstream() + async with app.router.lifespan_context(app): + client: Final = _StatefulUpstreamClient( + app, + server_url="http://upstream/mcp", + transport_type=MCPTransport.http, + ) + session: Final = client.open_persistent_session() + request: Final = httpx2.Request("POST", "http://upstream/mcp") + response: Final = httpx2.Response(401, request=request, headers={"WWW-Authenticate": "Bearer"}) + http_error: Final = httpx2.HTTPStatusError("401 Unauthorized", request=request, response=response) + + async def operation(_: ClientSession) -> str: + stream_error: Final[asyncio.Future[Exception] | None] = session._stream_error + assert stream_error is not None + stream_error.set_result(http_error) + raise MCPError(code=-32603, message="response stream closed") + + try: + with pytest.raises(httpx2.HTTPStatusError) as caught: + await session.run(operation) + assert caught.value is http_error + assert not session.reusable + finally: + session.close() + await asyncio.wait_for(session.wait_closed(), 5) + + +@pytest.mark.asyncio +async def test_persistent_session_ends_when_a_malformed_response_is_remapped() -> None: + app: Final = _upstream_with_malformed_response_for_method("tools/call") + async with app.router.lifespan_context(app): + client: Final = _StatefulUpstreamClient( + app, + server_url="http://upstream/mcp", + transport_type=MCPTransport.http, + ) + session: Final = client.open_persistent_session() + try: + + async def ready(_: ClientSession) -> str: + return "ready" + + assert await asyncio.wait_for(session.run(ready), 5) == "ready" + stream_error: Final[asyncio.Future[Exception] | None] = session._stream_error + assert stream_error is not None + recorded_stream_error: Final = await asyncio.wait_for(asyncio.shield(stream_error), 5) + assert isinstance(recorded_stream_error, ValidationError) + + with pytest.raises(ValidationError) as caught: + await client.call_tool( + CallToolRequestParams(name="malformed_tool", arguments={}), + raise_on_error=True, + persistent_session=session, + ) + assert caught.value is recorded_stream_error + assert not session.reusable + await asyncio.wait_for(session.wait_closed(), 5) + assert session.closed + finally: + if not session.closed: + session.close() + await asyncio.wait_for(session.wait_closed(), 5) + + +@pytest.mark.asyncio +async def test_failed_persistent_initialize_is_not_reusable_during_transport_cleanup() -> None: + app: Final = _upstream_with_status_for_method("initialize", 500) + async with app.router.lifespan_context(app): + cleanup_started: Final = asyncio.Event() + release_cleanup: Final = asyncio.Event() + client: Final = _CleanupHoldingStatefulUpstreamClient( + app, + server_url="http://upstream/mcp", + transport_type=MCPTransport.http, + cleanup_started=cleanup_started, + release_cleanup=release_cleanup, + ) + session: Final = client.open_persistent_session() + pending: Final = asyncio.create_task(session.run(lambda _: "ready")) + try: + await asyncio.wait_for(cleanup_started.wait(), 5) + assert not session.reusable + assert not session.closed + release_cleanup.set() + with pytest.raises(httpx2.HTTPStatusError): + await asyncio.wait_for(pending, 5) + await asyncio.wait_for(session.wait_closed(), 5) + assert session.closed + finally: + release_cleanup.set() + if not session.closed: + session.close() + await asyncio.wait_for(session.wait_closed(), 5) + + +@pytest.mark.asyncio +async def test_cancelled_queued_operation_does_not_run_and_session_accepts_a_later_operation() -> None: + app: Final = _stateful_upstream() + async with app.router.lifespan_context(app): + _, session = _client_with_session(app) + first_started: Final = asyncio.Event() + release_first: Final = asyncio.Event() + cancelled_operation_ran: Final = asyncio.Event() + + async def first_operation(_: ClientSession) -> str: + first_started.set() + await release_first.wait() + return "first" + + async def cancelled_operation(_: ClientSession) -> str: + cancelled_operation_ran.set() + return "cancelled" + + async def third_operation(_: ClientSession) -> str: + return "third" + + first_task: Final = asyncio.create_task(session.run(first_operation)) + try: + await asyncio.wait_for(first_started.wait(), 5) + second_task: Final = asyncio.create_task(session.run(cancelled_operation)) + await asyncio.sleep(0) + assert session._queue.qsize() == 1 + second_task.cancel() + with pytest.raises(asyncio.CancelledError): + await second_task + third_task: Final = asyncio.create_task(session.run(third_operation)) + await asyncio.sleep(0) + assert session._queue.qsize() == 2 + release_first.set() + assert await asyncio.wait_for(first_task, 5) == "first" + assert await asyncio.wait_for(third_task, 5) == "third" + assert not cancelled_operation_ran.is_set() + assert not session.closed + finally: + release_first.set() + session.close() + await asyncio.wait_for(session.wait_closed(), 5) + + +@pytest.mark.asyncio +async def test_retiring_persistent_session_drains_in_flight_and_queued_operations() -> None: + app: Final = _stateful_upstream() + async with app.router.lifespan_context(app): + _, session = _client_with_session(app) + first_started: Final = asyncio.Event() + release_first: Final = asyncio.Event() + + async def first_operation(_: ClientSession) -> str: + first_started.set() + await release_first.wait() + return "first" + + async def queued_operation(_: ClientSession) -> str: + return "queued" + + first_waiter: Final = asyncio.create_task(session.run(first_operation)) + try: + await asyncio.wait_for(first_started.wait(), 5) + queued_waiter: Final = asyncio.create_task(session.run(queued_operation)) + await asyncio.sleep(0) + assert session._queue.qsize() == 1 + session.retire() + release_first.set() + + results: Final = await asyncio.wait_for(asyncio.gather(first_waiter, queued_waiter), 5) + assert results == ["first", "queued"] + await asyncio.wait_for(session.wait_closed(), 5) + assert session.closed + finally: + release_first.set() + if not session.closed: + session.close() + await asyncio.wait_for(session.wait_closed(), 5) + + +@pytest.mark.asyncio +async def test_timed_out_persistent_operation_waiting_for_admission_is_skipped() -> None: + app: Final = _stateful_upstream() + async with app.router.lifespan_context(app): + admission_slots: Final = asyncio.Semaphore(1) + admission_requested: Final = asyncio.Event() + + def admission() -> AbstractAsyncContextManager[None]: + admission_requested.set() + return admission_slots + + _, session = _client_with_session(app, admission) + operation_ran: Final = asyncio.Event() + + async def warmup(_: ClientSession) -> str: + return "ready" + + async def timed_out_operation(_: ClientSession) -> str: + operation_ran.set() + return "expired" + + async def later_operation(_: ClientSession) -> str: + return "later" + + try: + assert await session.run(warmup) == "ready" + admission_requested.clear() + async with admission_slots: + timed_out_call: Final = asyncio.create_task(session.run(timed_out_operation)) + await asyncio.wait_for(admission_requested.wait(), 1) + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(timed_out_call, 0.05) + assert not operation_ran.is_set() + + later_result: Final = await asyncio.wait_for(session.run(later_operation), 1) + + assert later_result == "later" + assert session.reusable + assert not operation_ran.is_set() + finally: + session.close() + await asyncio.wait_for(session.wait_closed(), 5) + + +@pytest.mark.asyncio +async def test_persistent_session_holds_admission_until_initialize_completes() -> None: + app: Final = _stateful_upstream() + async with app.router.lifespan_context(app): + initialize_seen: Final = asyncio.Event() + admission_slots: Final = asyncio.Semaphore(1) + + async def observing_app(scope: Scope, receive: Receive, send: Send) -> None: + async def observe_send(message: Message) -> None: + if message["type"] == "http.response.body": + initialize_seen.set() + await send(message) + + await app(scope, receive, observe_send) + + client: Final = _StatefulUpstreamClient( + observing_app, + server_url="http://upstream/mcp", + transport_type=MCPTransport.http, + ) + session: Final = client.open_persistent_session(lambda: admission_slots) + + async def operation(_: ClientSession) -> str: + return "ready" + + run_task: Final = asyncio.create_task(session.run(operation)) + try: + async with admission_slots: + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(initialize_seen.wait(), 0.05) + assert not session._ready.done() + assert not run_task.done() + + result: Final = await asyncio.wait_for(run_task, 5) + assert result == "ready" + assert initialize_seen.is_set() + finally: + session.close() + await asyncio.wait_for(session.wait_closed(), 5) + await asyncio.gather(run_task, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_persistent_session_releases_admission_when_connect_fails() -> None: + async def failing_app(_: Scope, __: Receive, ___: Send) -> None: + raise httpx2.ConnectError("connect failed") + + client: Final = _StatefulUpstreamClient( + failing_app, + server_url="http://upstream/mcp", + transport_type=MCPTransport.http, + ) + admission_slots: Final = asyncio.Semaphore(1) + session: Final = client.open_persistent_session(lambda: admission_slots) + + async def operation(_: ClientSession) -> str: + return "unreachable" + + try: + with pytest.raises(httpx2.ConnectError, match="connect failed"): + await asyncio.wait_for(session.run(operation), 1) + await asyncio.wait_for(session.wait_closed(), 1) + + await asyncio.wait_for(admission_slots.acquire(), 0.1) + admission_slots.release() + finally: + session.close() + await asyncio.wait_for(session.wait_closed(), 5) + + +@pytest.mark.asyncio +async def test_failed_connect_is_not_reusable_during_admission_teardown() -> None: + connect_failed: Final = asyncio.Event() + admission_exit_started: Final = asyncio.Event() + release_admission_exit: Final = asyncio.Event() + + async def failing_app(_: Scope, __: Receive, ___: Send) -> None: + connect_failed.set() + raise httpx2.ConnectError("connect failed") + + @asynccontextmanager + async def blocking_admission() -> AsyncIterator[None]: + try: + yield + finally: + admission_exit_started.set() + await release_admission_exit.wait() + + client: Final = _StatefulUpstreamClient( + failing_app, + server_url="http://upstream/mcp", + transport_type=MCPTransport.http, + ) + session: Final = client.open_persistent_session(blocking_admission) + + async def operation(_: ClientSession) -> str: + return "unreachable" + + run_task: Final = asyncio.create_task(session.run(operation)) + try: + await asyncio.wait_for(connect_failed.wait(), 1) + await asyncio.wait_for(admission_exit_started.wait(), 1) + + assert not run_task.done() + assert not session.reusable + assert not session.closed + + release_admission_exit.set() + with pytest.raises(httpx2.ConnectError, match="connect failed"): + await asyncio.wait_for(run_task, 1) + await asyncio.wait_for(session.wait_closed(), 1) + assert session.closed + finally: + release_admission_exit.set() + session.close() + await asyncio.wait_for(session.wait_closed(), 5) + await asyncio.gather(run_task, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_persistent_session_admission_is_acquired_only_when_operation_runs() -> None: + app: Final = _stateful_upstream() + async with app.router.lifespan_context(app): + admission_slots: Final = asyncio.Semaphore(2) + _, session = _client_with_session(app, lambda: admission_slots) + first_started: Final = asyncio.Event() + release_first: Final = asyncio.Event() + operation_order: Final[list[str]] = [] + + async def first_operation(_: ClientSession) -> str: + operation_order.append("first") + first_started.set() + await release_first.wait() + return "first" + + async def second_operation(_: ClientSession) -> str: + operation_order.append("second") + return "second" + + async def third_operation(_: ClientSession) -> str: + operation_order.append("third") + return "third" + + first_task: Final = asyncio.create_task(session.run(first_operation)) + second_task: Final = asyncio.create_task(session.run(second_operation)) + third_task: Final = asyncio.create_task(session.run(third_operation)) + try: + await asyncio.wait_for(first_started.wait(), 5) + await asyncio.sleep(0) + assert session._queue.qsize() == 2 + await asyncio.wait_for(admission_slots.acquire(), 0.1) + try: + assert admission_slots.locked() + finally: + admission_slots.release() + assert operation_order == ["first"] + + release_first.set() + results: Final = await asyncio.wait_for( + asyncio.gather(first_task, second_task, third_task), + 5, + ) + + assert results == ["first", "second", "third"] + assert operation_order == ["first", "second", "third"] + finally: + release_first.set() + session.close() + await asyncio.wait_for(session.wait_closed(), 5) + await asyncio.gather(first_task, second_task, third_task, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_retiring_idle_persistent_session_closes_and_rejects_new_operations() -> None: + app: Final = _stateful_upstream() + async with app.router.lifespan_context(app): + _, session = _client_with_session(app) + + async def operation(_: ClientSession) -> str: + return "ready" + + try: + assert await session.run(operation) == "ready" + session.retire() + await asyncio.wait_for(session.wait_closed(), 1) + + assert session.closed + with pytest.raises(UpstreamSessionClosedError): + await session.run(operation) + finally: + if not session.closed: + session.close() + await asyncio.wait_for(session.wait_closed(), 5) + + +@pytest.mark.asyncio +async def test_closing_persistent_session_mid_operation_fails_the_waiter_instead_of_hanging() -> None: + app: Final = _stateful_upstream() + async with app.router.lifespan_context(app): + _, session = _client_with_session(app) + started: Final = asyncio.Event() + + async def slow_operation(_: object) -> str: + started.set() + await asyncio.sleep(30) + return "never" + + waiter: Final = asyncio.ensure_future(session.run(slow_operation)) + await asyncio.wait_for(started.wait(), 5) + session.close() + with pytest.raises(RuntimeError, match="upstream MCP session closed"): + await asyncio.wait_for(waiter, 5) + await asyncio.wait_for(session.wait_closed(), 5) + + +@pytest.mark.asyncio +async def test_persistent_session_reports_an_upstream_cancellation_as_a_runtime_error_not_a_cancelled_caller() -> None: + app: Final = _stateful_upstream() + async with app.router.lifespan_context(app): + client, session = _client_with_session(app) + + async def cancelled_upstream(_: object) -> str: + raise asyncio.CancelledError() + + try: + with pytest.raises(RuntimeError, match="cancelled"): + await asyncio.wait_for(session.run(cancelled_upstream), 5) + selected: Final = await client.call_tool( + CallToolRequestParams(name="select_project", arguments={"name": "e"}), persistent_session=session + ) + finally: + session.close() + assert selected.is_error is False, "an upstream cancellation must not be mistaken for a cancelled caller" + await asyncio.wait_for(session.wait_closed(), 5) + + +@pytest.mark.asyncio +async def test_closing_persistent_session_fails_a_caller_blocked_on_a_full_queue_instead_of_hanging() -> None: + from litellm.mcp_client.client import _MAX_PENDING_OPERATIONS + + app: Final = _stateful_upstream() + async with app.router.lifespan_context(app): + _, session = _client_with_session(app) + started: Final = asyncio.Event() + + async def slow_operation(_: object) -> str: + started.set() + await asyncio.sleep(30) + return "never" + + waiters: Final = tuple( + asyncio.ensure_future(session.run(slow_operation)) for _ in range(_MAX_PENDING_OPERATIONS + 2) + ) + await asyncio.wait_for(started.wait(), 5) + await asyncio.sleep(0) + session.close() + outcomes: Final = await asyncio.wait_for(asyncio.gather(*waiters, return_exceptions=True), 5) + assert all(isinstance(outcome, RuntimeError) for outcome in outcomes), outcomes + await asyncio.wait_for(session.wait_closed(), 5) + + +@pytest.mark.asyncio +async def test_persistent_session_ends_after_a_broken_stream_so_the_next_call_gets_a_fresh_one() -> None: + app: Final = _stateful_upstream() + async with app.router.lifespan_context(app): + _, session = _client_with_session(app) + + async def broken_stream(_: object) -> str: + raise anyio.BrokenResourceError() + + with pytest.raises(anyio.BrokenResourceError): + await asyncio.wait_for(session.run(broken_stream), 5) + await asyncio.wait_for(session.wait_closed(), 5) + assert session.closed, "a dead transport must end the session instead of being reused" + +@pytest.mark.asyncio +@pytest.mark.parametrize("legacy_callback", [False, True]) +async def test_tool_progress_preserves_message_and_legacy_callback(legacy_callback: bool) -> None: + received = AsyncMock() + + async def legacy(progress: float, total: float | None) -> None: + await received(progress, total) + + async def modern(progress: float, total: float | None, message: str | None) -> None: + await received(progress, total, message) + + async def call_tool(**kwargs): + await kwargs["progress_callback"](1.0, 2.0, "halfway") + return CallToolResult(content=[]) + + session = AsyncMock(spec=ClientSession, call_tool=AsyncMock(side_effect=call_tool)) + + async def run(operation, **kwargs): + return await operation(session) + + persistent = MagicMock(spec=PersistentMCPSession, run=AsyncMock(side_effect=run)) + client = MCPClient(server_url="https://example.com/mcp") + result = await client.call_tool( + CallToolRequestParams(name="slow"), + host_progress_callback=legacy if legacy_callback else modern, + persistent_session=persistent, + ) + assert not result.is_error + if legacy_callback: + received.assert_awaited_once_with(1.0, 2.0) + else: + received.assert_awaited_once_with(1.0, 2.0, "halfway") + + +@pytest.mark.asyncio +async def test_cancelling_persistent_operation_stops_work_and_keeps_session_usable() -> None: + app: Final = _stateful_upstream() + async with app.router.lifespan_context(app): + client, session = _client_with_session(app) + started: Final = asyncio.Event() + stopped: Final = asyncio.Event() + + async def slow_operation(_: ClientSession) -> None: + started.set() + try: + await asyncio.Event().wait() + finally: + stopped.set() + + waiter: Final = asyncio.create_task(session.run(slow_operation)) + try: + await asyncio.wait_for(started.wait(), 5) + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + await asyncio.sleep(0.1) + assert stopped.is_set(), "Caller cancellation must stop the running upstream operation" + result: Final = await asyncio.wait_for( + client.call_tool( + CallToolRequestParams(name="select_project", arguments={"name": "after-cancellation"}), + persistent_session=session, + ), + 5, + ) + assert not result.is_error + finally: + session.close() + await asyncio.wait_for(session.wait_closed(), 5) + await asyncio.gather(waiter, return_exceptions=True) diff --git a/tests/integration/mcp/test_openapi_to_mcp_generator_real_io.py b/tests/integration/mcp/test_openapi_to_mcp_generator_real_io.py new file mode 100644 index 00000000000..acdeef06898 --- /dev/null +++ b/tests/integration/mcp/test_openapi_to_mcp_generator_real_io.py @@ -0,0 +1,44 @@ +""" +Tests for OpenAPI to MCP generator, focusing on security and edge cases. + +This test suite ensures that: +1. Parameter names with invalid Python identifiers are handled safely +2. No exec() is used (security) +3. All edge cases (hyphens, dots, keywords, special chars) work correctly +4. Path traversal attacks are prevented +5. Path parameters are properly URL encoded +""" + +import pytest + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) + + +# --------------------------------------------------------------------------- +# Tool name sanitization for OpenAPI -> MCP +# Repro: GitHub's REST OpenAPI uses tag-namespaced operationIds like +# "actions/download-job-logs-for-workflow-run". Without sanitization the +# generated MCP tool name contains '/', which Anthropic/OpenAI/Bedrock all +# reject (^[a-zA-Z0-9_-]+$). This block guards the registration + preview +# paths against that. +# --------------------------------------------------------------------------- + + +def test_openapi_generator_import_does_not_require_mcp_sdk() -> None: + import subprocess + import sys + + script = """ +import builtins +original_import = builtins.__import__ +def without_mcp(name, *args, **kwargs): + if name == 'mcp' or name.startswith('mcp.'): + raise ModuleNotFoundError('MCP SDK unavailable') + return original_import(name, *args, **kwargs) +builtins.__import__ = without_mcp +import litellm.proxy.mcp_server.openapi_to_mcp_generator +""" + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True) + assert result.returncode == 0, result.stderr diff --git a/tests/integration/observability/test_anthropic_cache_control_hook_real_io.py b/tests/integration/observability/test_anthropic_cache_control_hook_real_io.py new file mode 100644 index 00000000000..ab412de705e --- /dev/null +++ b/tests/integration/observability/test_anthropic_cache_control_hook_real_io.py @@ -0,0 +1,65 @@ +import json +import os +import subprocess +import sys +import textwrap +from typing import Optional, Tuple + +import pytest + +import litellm + + +@pytest.fixture(autouse=True) +def _no_openai_api_base_override(monkeypatch): + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", None) + + +class TestAnthropicPromptCachingEnvVars: + """Both settings are read from the environment at import, so an admin can enable + auto-caching without a config file. Each case re-imports litellm in a subprocess + so the env is read fresh without contaminating this process's module graph. + """ + + @staticmethod + def _import_litellm_with_env(env_override: dict) -> Tuple[bool, Optional[str]]: + env = os.environ.copy() + env.pop("LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING", None) + env.pop("LITELLM_ANTHROPIC_PROMPT_CACHING_TTL", None) + env.update(env_override) + script = textwrap.dedent( + """ + import json, litellm + print(json.dumps([litellm.enable_anthropic_prompt_caching, litellm.anthropic_prompt_caching_ttl])) + """ + ) + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, env=env, timeout=300) + assert result.returncode == 0, result.stderr + enabled, ttl = json.loads(result.stdout.strip().splitlines()[-1]) + return enabled, ttl + + def test_unset_env_leaves_auto_caching_off(self): + assert self._import_litellm_with_env({}) == (False, None) + + @pytest.mark.parametrize("value", ["true", "True", "TRUE"]) + def test_env_enables_auto_caching_case_insensitively(self, value): + enabled, _ = self._import_litellm_with_env({"LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING": value}) + assert enabled is True + + @pytest.mark.parametrize("value", ["false", "0", "yes", ""]) + def test_env_only_enables_on_true(self, value): + enabled, _ = self._import_litellm_with_env({"LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING": value}) + assert enabled is False + + @pytest.mark.parametrize("value", ["5m", "1h"]) + def test_ttl_env_is_applied(self, value): + _, ttl = self._import_litellm_with_env({"LITELLM_ANTHROPIC_PROMPT_CACHING_TTL": value}) + assert ttl == value + + @pytest.mark.parametrize("value", ["10m", "1H", "3600", "ephemeral"]) + def test_unsupported_ttl_env_falls_back_to_provider_default(self, value): + """An unparseable TTL must fall back to Anthropic's 5m default, never reach the provider verbatim.""" + _, ttl = self._import_litellm_with_env({"LITELLM_ANTHROPIC_PROMPT_CACHING_TTL": value}) + assert ttl is None diff --git a/tests/integration/sdk/test_anthropic_common_utils_real_io.py b/tests/integration/sdk/test_anthropic_common_utils_real_io.py new file mode 100644 index 00000000000..c19dcbb51d1 --- /dev/null +++ b/tests/integration/sdk/test_anthropic_common_utils_real_io.py @@ -0,0 +1,41 @@ +import json +import subprocess +import sys +from typing import Final + + +_MASKED_HOSTED_WEB_CALL_PROBE: Final = """ +import json + +from litellm.llms.anthropic.common_utils import write_request_hosted_web_call_arguments + +message = { + "role": "assistant", + "provider_specific_fields": { + "web_search_calls": [ + { + "id": "ws_srvtoolu_1", + "type": "web_search_call", + "status": "completed", + "action": {"type": "search", "query": "deploy token 12345", "queries": ["deploy token 12345"]}, + } + ] + }, +} +write_request_hosted_web_call_arguments(message, {"srvtoolu_1": json.dumps({"query": "deploy token [REDACTED]"})}) +print(json.dumps(message["provider_specific_fields"]["web_search_calls"])) +""" + + +def test_a_masked_replayed_hosted_web_call_serializes_before_its_openai_model_was_ever_built() -> None: + result: Final = subprocess.run( + (sys.executable, "-I", "-c", _MASKED_HOSTED_WEB_CALL_PROBE), + capture_output=True, + text=True, + check=True, + timeout=120, + ) + hosted: Final = json.loads(result.stdout.strip().splitlines()[-1]) + assert [call["action"] for call in hosted] == [ + {"type": "search", "query": "deploy token [REDACTED]", "queries": ["deploy token [REDACTED]"]} + ] diff --git a/tests/integration/sdk/test_azure_common_utils_real_io.py b/tests/integration/sdk/test_azure_common_utils_real_io.py new file mode 100644 index 00000000000..5c988435194 --- /dev/null +++ b/tests/integration/sdk/test_azure_common_utils_real_io.py @@ -0,0 +1,204 @@ +import os +from collections.abc import Mapping +from typing import Final, Optional +from unittest.mock import MagicMock +import openai +import requests + +import litellm +from litellm.llms.azure.common_utils import BaseAzureLLM + + +def test_default_max_retries_env_var_reaches_azure_sdk_client(): + import subprocess + import sys + + code = ( + "from litellm.llms.azure.common_utils import BaseAzureLLM\n" + "client = BaseAzureLLM().get_azure_openai_client(" + "api_key='test-api-key', api_base='https://test.openai.azure.com', api_version='2024-02-01'," + " client=None, _is_async=True, litellm_params={}, model='gpt-4')\n" + "print(client.max_retries)" + ) + completed = subprocess.run( + [sys.executable, "-c", code], + env={**os.environ, "DEFAULT_MAX_RETRIES": "0"}, + capture_output=True, + text=True, + check=True, + ) + + assert completed.stdout.strip() == "0" + + +# Test parameters for different API functions with Azure models +AZURE_API_FUNCTION_PARAMS = [ + # (function_name, is_async, args) + ( + "completion", + False, + { + "model": "azure/gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 10, + "api_key": "test-api-key", + "api_base": "https://test.openai.azure.com", + "api_version": "2023-05-15", + }, + ), + ( + "completion", + True, + { + "model": "azure/gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 10, + "stream": True, + "api_key": "test-api-key", + "api_base": "https://test.openai.azure.com", + "api_version": "2023-05-15", + }, + ), + ( + "embedding", + False, + { + "model": "azure/text-embedding-ada-002", + "input": "Hello world", + "api_key": "test-api-key", + "api_base": "https://test.openai.azure.com", + "api_version": "2023-05-15", + }, + ), + ( + "embedding", + True, + { + "model": "azure/text-embedding-ada-002", + "input": "Hello world", + "api_key": "test-api-key", + "api_base": "https://test.openai.azure.com", + "api_version": "2023-05-15", + }, + ), + ( + "speech", + False, + { + "model": "azure/tts", + "input": "Hello, this is a test of text to speech", + "voice": "alloy", + "api_key": "test-api-key", + "api_base": "https://test.openai.azure.com", + "api_version": "2023-05-15", + }, + ), + ( + "speech", + True, + { + "model": "azure/tts-1", + "input": "Hello, this is a test of text to speech", + "voice": "alloy", + "api_key": "test-api-key", + "api_base": "https://test.openai.azure.com", + "api_version": "2023-05-15", + }, + ), + ( + "transcription", + False, + { + "model": "azure/whisper-1", + "file": MagicMock(), + "api_key": "test-api-key", + "api_base": "https://test.openai.azure.com", + "api_version": "2023-05-15", + }, + ), + ( + "transcription", + True, + { + "model": "azure/whisper-1", + "file": MagicMock(), + "api_key": "test-api-key", + "api_base": "https://test.openai.azure.com", + "api_version": "2023-05-15", + }, + ), +] +_ENTRA_TRACE_ID: Final = "f0c149d3-ef20-4194-8b81-c81a5e002500" + + +def _entra_error(aadsts_code: str, description: str) -> Mapping[str, object]: + return { + "error": "invalid_client", + "error_description": f"{aadsts_code}: {description} Trace ID: {_ENTRA_TRACE_ID}", + "error_codes": [int(aadsts_code.removeprefix("AADSTS"))], + } + + +_REAL_CREDENTIAL_NON_REJECTIONS: Final = { + "unreachable_through_a_credential_chain": ("default_azure_credential_chain", None, ""), + "html_503_through_a_credential_chain": ( + "default_azure_credential_chain", + 503, + "Service Unavailable", + ), + "json_503_through_a_credential_chain": ( + "default_azure_credential_chain", + 503, + _entra_error("AADSTS90033", "A transient error has occurred."), + ), + "throttled_client_secret": ( + "client_secret", + 429, + _entra_error("AADSTS50196", "The server terminated an operation because it encountered a client request loop."), + ), + "tenant_throttled_client_secret": ( + "client_secret", + 429, + _entra_error("AADSTS90055", "There are too many incoming requests."), + ), + "request_budget_exceeded_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS90024", "A transient error has occurred. Try again."), + ), + "retryable_error_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS50033", "A transient error has occurred."), + ), + "external_server_retryable_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS90006", "The service is temporarily unavailable."), + ), + "graph_retryable_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS90090", "The service is temporarily unavailable."), + ), + "token_issuance_error_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS50000", "There was an error issuing a token or an issue with our sign-in service."), + ), + "throttled_without_a_code_through_a_credential_chain": ( + "default_azure_credential_chain", + 429, + {"error": "temporarily_unavailable", "error_description": "Too many requests."}, + ), + "json_503_through_workload_identity": ( + "workload_identity", + 503, + _entra_error("AADSTS90033", "A transient error has occurred."), + ), + "throttled_workload_identity_naming_a_rejection_code": ( + "workload_identity", + 429, + _entra_error("AADSTS7000215", "Invalid client secret provided."), + ), +} diff --git a/tests/integration/sdk/test_configuration_real_io.py b/tests/integration/sdk/test_configuration_real_io.py new file mode 100644 index 00000000000..2821bbab722 --- /dev/null +++ b/tests/integration/sdk/test_configuration_real_io.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import os +import subprocess +import sys +from collections.abc import Generator +from typing import Final + +import pytest + +from litellm.rust_bridge import configuration + + +@pytest.fixture(autouse=True) +def _isolated_configuration( # pyright: ignore[reportUnusedFunction] # pytest discovers fixtures dynamically + monkeypatch: pytest.MonkeyPatch, +) -> Generator[None]: + configuration.reset_rust_configuration() + monkeypatch.delenv("LITELLM_RUST", raising=False) + yield + configuration.reset_rust_configuration() + + +@pytest.mark.parametrize(("value", "expected"), (("1", "True"), ("0", "False"))) +def test_environment_controls_startup(value: str, expected: str) -> None: + environment: Final = {**os.environ, "LITELLM_RUST": value} + result: Final = subprocess.run( + ( + sys.executable, + "-c", + "from litellm.rust_bridge.configuration import rust_enabled; print(rust_enabled())", + ), + check=True, + capture_output=True, + text=True, + env=environment, + ) + + assert result.stdout.strip() == expected diff --git a/tests/integration/sdk/test_context_real_io.py b/tests/integration/sdk/test_context_real_io.py new file mode 100644 index 00000000000..233aeda04f6 --- /dev/null +++ b/tests/integration/sdk/test_context_real_io.py @@ -0,0 +1,146 @@ +import asyncio +import base64 +import uuid +from collections.abc import Generator +from contextlib import contextmanager +from dataclasses import dataclass, field +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from threading import Event, Thread +from typing import Final, Literal + +import pytest + +import litellm +from litellm.llms.anthropic.messages.context_management.inference_view import CountContextEdits +from litellm.llms.vertex_ai.count_tokens.context import VertexGeminiCountContextError, gemini_count_contents +_NO_EDITS: Final = CountContextEdits(spec=None, model="gemini/gemini-2.5-flash", tools=None) + + +_PNG_BYTES: Final = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" +) +_LOOP_PROOF_TIMEOUT_SECONDS: Final = 5 + + +@dataclass(frozen=True, slots=True) +class _LoopProbe: + request_arrived: Event = field(default_factory=Event) + loop_ran: Event = field(default_factory=Event) + loop_ran_while_download_pending: Event = field(default_factory=Event) + + +def _png_handler_waiting_for_loop(probe: _LoopProbe) -> type[BaseHTTPRequestHandler]: + class _PngHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + probe.request_arrived.set() + if probe.loop_ran.wait(_LOOP_PROOF_TIMEOUT_SECONDS): + probe.loop_ran_while_download_pending.set() + self.send_response(200) + self.send_header("Content-Type", "image/png") + self.send_header("Content-Length", str(len(_PNG_BYTES))) + self.end_headers() + self.wfile.write(_PNG_BYTES) + + def log_message(self, format: str, *args: object) -> None: + return None + + return _PngHandler + + +@contextmanager +def _png_server_waiting_for_loop(probe: _LoopProbe) -> Generator[str, None, None]: + with ThreadingHTTPServer(("127.0.0.1", 0), _png_handler_waiting_for_loop(probe)) as server: + thread: Final = Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + + +async def _run_loop_once_download_starts(probe: _LoopProbe) -> None: + assert await asyncio.to_thread(probe.request_arrived.wait, _LOOP_PROOF_TIMEOUT_SECONDS) + probe.loop_ran.set() + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_remote_image_downloads_without_blocking_the_event_loop( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + probe: Final = _LoopProbe() + with _png_server_waiting_for_loop(probe) as base_url: + image_url: Final = f"{base_url}/{uuid.uuid4().hex}.png" + _, contents = await asyncio.gather( + _run_loop_once_download_starts(probe), + gemini_count_contents( + ( + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe it"}, + {"type": "image", "source": {"type": "url", "url": image_url}}, + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ), + ) + + assert contents == [ + { + "role": "user", + "parts": [ + {"text": "Describe it"}, + {"inline_data": {"data": base64.b64encode(_PNG_BYTES).decode(), "mime_type": "image/png"}}, + ], + } + ] + assert probe.loop_ran_while_download_pending.is_set() + + +class _MissingImageHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + self.send_response(404) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, format: str, *args: object) -> None: + return None + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_image_that_cannot_be_fetched_returns_an_error_value( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + with ThreadingHTTPServer(("127.0.0.1", 0), _MissingImageHandler) as server: + Thread(target=server.serve_forever, daemon=True).start() + try: + contents = await gemini_count_contents( + ( + { + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "url", + "url": f"http://127.0.0.1:{server.server_address[1]}/{uuid.uuid4().hex}.png", + }, + } + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ) + finally: + server.shutdown() + + assert isinstance(contents, VertexGeminiCountContextError) diff --git a/tests/integration/sdk/test_docker_real_io.py b/tests/integration/sdk/test_docker_real_io.py new file mode 100644 index 00000000000..d0d6d9bedf2 --- /dev/null +++ b/tests/integration/sdk/test_docker_real_io.py @@ -0,0 +1,50 @@ +import hashlib +import shutil +import subprocess + +import pytest + +from litellm import sandbox + +DOCKER_IMAGE = "alpine:3.20" + + +def _docker_usable() -> bool: + if shutil.which("docker") is None: + return False + try: + return ( + subprocess.run( + ["docker", "info"], capture_output=True, timeout=20 + ).returncode + == 0 + ) + except (OSError, subprocess.SubprocessError): + return False + + +@pytest.mark.skipif(not _docker_usable(), reason="docker daemon not available") +async def test_real_docker_roundtrip(): + box = sandbox.docker(DOCKER_IMAGE, workdir="/workspace") + try: + result = await box.run(["echo", "hello"], timeout=120) + assert result.stdout.strip() == "hello" + assert result.exit_code == 0 + + await box.write("seed.txt", b"seed") + await box.write("sub/out.txt", b"from host") + assert await box.read("sub/out.txt") == b"from host" + await box.write("node_modules/skip.js", b"x") + + assert await box.which("sh") is not None + assert await box.which("definitely-not-a-binary-xyz") is None + tmp = await box.tempdir() + assert tmp.startswith("/") + + snap = await box.snapshot() + assert snap == { + "seed.txt": hashlib.sha256(b"seed").hexdigest(), + "sub/out.txt": hashlib.sha256(b"from host").hexdigest(), + } + finally: + await box.close() diff --git a/tests/unit/test_drop_params_env_var.py b/tests/integration/sdk/test_drop_params_env_var_real_io.py similarity index 100% rename from tests/unit/test_drop_params_env_var.py rename to tests/integration/sdk/test_drop_params_env_var_real_io.py diff --git a/tests/unit/test_eager_tiktoken_load.py b/tests/integration/sdk/test_eager_tiktoken_load_real_io.py similarity index 100% rename from tests/unit/test_eager_tiktoken_load.py rename to tests/integration/sdk/test_eager_tiktoken_load_real_io.py diff --git a/tests/integration/sdk/test_http_handler_real_io.py b/tests/integration/sdk/test_http_handler_real_io.py new file mode 100644 index 00000000000..b7c0e9d748b --- /dev/null +++ b/tests/integration/sdk/test_http_handler_real_io.py @@ -0,0 +1,296 @@ +import asyncio +import importlib +import json +import os +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final +import httpx +import pytest + +import litellm +from litellm.exceptions import Timeout as LitellmTimeout +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.model_cost_table import publish +from litellm.llms.base_llm.auth.types import TokenExchangeSpec +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + MaskedHTTPStatusError, + get_httpx_client, +) +from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome + + +async def _read_http_request(reader: asyncio.StreamReader) -> None: + raw = b"" + while b"\r\n\r\n" not in raw: + chunk = await reader.read(1024) + if not chunk: + return + raw += chunk + head, _, body = raw.partition(b"\r\n\r\n") + content_length = next( + (int(line.split(b":", 1)[1]) for line in head.split(b"\r\n") if line.lower().startswith(b"content-length")), + 0, + ) + while len(body) < content_length: + body += await reader.read(content_length - len(body)) + + +@pytest.mark.asyncio +async def test_init_held_async_handler_survives_external_client_close(): + handler = AsyncHTTPHandler(timeout=42.5) + held_client = handler.client + await held_client.aclose() + assert held_client.is_closed + + async def respond(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + await _read_http_request(reader) + writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok") + await writer.drain() + writer.close() + + server = await asyncio.start_server(respond, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + try: + response = await handler.post(f"http://127.0.0.1:{port}/v1/compress", json={"messages": []}) + finally: + server.close() + await server.wait_closed() + + assert response.status_code == 200 + assert handler.client is not held_client + assert handler.client.timeout == httpx.Timeout(42.5) + await handler.close() + + +@pytest.mark.asyncio +async def test_init_held_async_handler_survives_evicted_client_close(): + from litellm.caching.evicted_client_closer import EvictedClientCloser + from litellm.caching.llm_caching_handler import LLMClientCache + + cache = LLMClientCache(evicted_client_closer=EvictedClientCloser(grace_seconds=0)) + handler = AsyncHTTPHandler(timeout=42.5) + held_client = handler.client + cache.set_cache("init-held-handler", handler, litellm_owned_client=True, ttl=0) + await asyncio.sleep(0.02) + assert cache.get_cache("init-held-handler") is None + await asyncio.sleep(0.05) + assert held_client.is_closed + + async def respond(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + await _read_http_request(reader) + writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok") + await writer.drain() + writer.close() + + server = await asyncio.start_server(respond, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + try: + response = await handler.post(f"http://127.0.0.1:{port}/v1/compress", json={"messages": []}) + finally: + server.close() + await server.wait_closed() + + assert response.status_code == 200 + assert handler.client is not held_client + assert handler.client.timeout == httpx.Timeout(42.5) + await handler.close() + + +def test_init_held_sync_handler_recreates_closed_client(): + from http.server import BaseHTTPRequestHandler, HTTPServer + + class OkRequestHandler(BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(200) + self.send_header("Content-Length", "2") + self.end_headers() + self.wfile.write(b"ok") + + def log_message(self, format, *args): + pass + + handler = HTTPHandler(timeout=7) + held_client = handler.client + held_client.close() + + server = HTTPServer(("127.0.0.1", 0), OkRequestHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + response = handler.get(f"http://127.0.0.1:{server.server_port}/") + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + assert response.status_code == 200 + assert handler.client is not held_client + assert handler.client.timeout == httpx.Timeout(7) + handler.close() + + +_WIF_SPEC: Final = TokenExchangeSpec( + token_url="https://api.anthropic.com/v1/oauth/token", + assertion_ref="oidc/env/ANTHROPIC_IDENTITY_TOKEN", + assertion_field="assertion", + static_body={"grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", "federation_rule_id": "fdrl_1"}, + body_encoding="json", + request_headers={}, + cache_key_identity=("fdrl_1", "org-1", "", ""), + timeout_seconds=2.0, +) + + +@pytest.fixture() +def _vcr_outcome_gate(request, vcr): + install_live_call_probe(request, vcr) + yield + record_vcr_outcome(request, vcr) + + +@pytest.fixture(scope="function") +def isolate_litellm_state(): + """ + Per-function isolation fixture. + + Resets litellm globals to their true defaults before each test and + restores them afterward, so tests don't leak side effects. + Works safely under pytest-xdist parallel execution. + """ + original_state = {} + for attr in ( + "callbacks", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", + ): + if hasattr(litellm, attr): + val = getattr(litellm, attr) + original_state[attr] = val.copy() if val else [] + for attr in ("pre_call_rules", "post_call_rules"): + if hasattr(litellm, attr): + val = getattr(litellm, attr) + original_state[attr] = val.copy() if val else [] + for attr in _SCALAR_DEFAULTS: + if hasattr(litellm, attr): + original_state[attr] = getattr(litellm, attr) + if hasattr(litellm, "in_memory_llm_clients_cache"): + litellm.in_memory_llm_clients_cache.flush_cache() + for attr in ( + "callbacks", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", + "pre_call_rules", + "post_call_rules", + ): + if hasattr(litellm, attr): + setattr(litellm, attr, []) + for attr, default_val in _SCALAR_DEFAULTS.items(): + if hasattr(litellm, attr): + setattr(litellm, attr, default_val) + yield + asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue()) + if hasattr(litellm, "in_memory_llm_clients_cache"): + litellm.in_memory_llm_clients_cache.flush_cache() + for attr, original_value in original_state.items(): + if hasattr(litellm, attr): + setattr(litellm, attr, original_value) + publish(litellm.model_cost) + + +_SCALAR_DEFAULTS = { + "num_retries": getattr(litellm, "num_retries", None), + "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), + "request_timeout": getattr(litellm, "request_timeout", None), + "set_verbose": getattr(litellm, "set_verbose", False), + "cache": getattr(litellm, "cache", None), + "allowed_fails": getattr(litellm, "allowed_fails", 3), + "default_fallbacks": getattr(litellm, "default_fallbacks", None), + "enable_azure_ad_token_refresh": getattr(litellm, "enable_azure_ad_token_refresh", None), + "tag_budget_config": getattr(litellm, "tag_budget_config", None), + "model_cost": getattr(litellm, "model_cost", None), + "token_counter": getattr(litellm, "token_counter", None), + "disable_aiohttp_transport": getattr(litellm, "disable_aiohttp_transport", False), + "force_ipv4": getattr(litellm, "force_ipv4", False), + "drop_params": getattr(litellm, "drop_params", None), + "modify_params": getattr(litellm, "modify_params", False), + "api_base": getattr(litellm, "api_base", None), + "api_key": getattr(litellm, "api_key", None), +} + + +@pytest.fixture(scope="module") +def setup_and_teardown(): + """ + Module-scoped setup. Reloads litellm only in single-process mode + (skipped under xdist to avoid cross-worker interference). + """ + import litellm + + worker_id = os.environ.get("PYTEST_XDIST_WORKER", None) + if worker_id is None: + importlib.reload(litellm) + try: + if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"): + import litellm.proxy.proxy_server + + importlib.reload(litellm.proxy.proxy_server) + except Exception: + pass + if hasattr(litellm, "in_memory_llm_clients_cache"): + litellm.in_memory_llm_clients_cache.flush_cache() + yield + + +_SERVER_DELAY_S = 5 + +_PER_REQUEST_TIMEOUT_S = 1.0 + +_CLIENT_DEFAULT_TIMEOUT_S = 60.0 + + +class _SlowHandler(BaseHTTPRequestHandler): + def do_POST(self): + time.sleep(_SERVER_DELAY_S) + try: + self.send_response(200) + self.end_headers() + self.wfile.write(b"{}") + except OSError: + pass + + def log_message(self, *args): + pass + + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +def test_post_delay_exceeds_per_request_timeout_raises(): + server = ThreadingHTTPServer(("127.0.0.1", 0), _SlowHandler) + threading.Thread(target=server.serve_forever, daemon=True).start() + host, port = server.server_address + + handler = get_httpx_client(params={"timeout": _CLIENT_DEFAULT_TIMEOUT_S}) + try: + with pytest.raises(LitellmTimeout): + handler.post( + f"http://{host}:{port}/delay", + headers={"content-type": "application/json"}, + data=json.dumps({"model": "claude", "messages": []}), + timeout=_PER_REQUEST_TIMEOUT_S, + ) + except MaskedHTTPStatusError as e: + pytest.skip(f"httpbin.org unavailable: {e}") + finally: + handler.close() + server.shutdown() + server.server_close() + + +_TIMEOUT_CLIENT_LIMITS: Final = httpx.Timeout(42.0, connect=7.0) diff --git a/tests/integration/sdk/test_jwt_signing_real_io.py b/tests/integration/sdk/test_jwt_signing_real_io.py new file mode 100644 index 00000000000..37c60fd55ff --- /dev/null +++ b/tests/integration/sdk/test_jwt_signing_real_io.py @@ -0,0 +1,37 @@ +import subprocess +import sys +import textwrap +from typing import Final + +from litellm.llms.base_llm.auth.jwt_signing import MISSING_SIGNING_DEPENDENCIES_MESSAGE + + +class TestBaseSdkImport: + """A base ``pip install litellm`` has neither PyJWT nor cryptography (both are proxy extras), + and ``litellm/__init__`` reaches this module through the Anthropic provider, so a + module-level import of either would break ``import litellm`` for every base SDK user.""" + + def test_module_imports_with_pyjwt_and_cryptography_absent(self): + script: Final = textwrap.dedent( + """ + import sys + + class Blocker: + def find_spec(self, name, path=None, target=None): + if name.split(".")[0] in {"jwt", "cryptography"}: + raise ModuleNotFoundError(f"No module named {name!r}") + + sys.meta_path.insert(0, Blocker()) + import litellm + from litellm.llms.base_llm.auth.jwt_signing import jwks_document_json + try: + jwks_document_json("not a key") + except ImportError as e: + print(e) + """ + ) + result: Final = subprocess.run( + [sys.executable, "-I", "-c", script], capture_output=True, text=True, check=False + ) + assert result.returncode == 0, result.stderr[-2000:] + assert result.stdout.strip() == MISSING_SIGNING_DEPENDENCIES_MESSAGE diff --git a/tests/integration/sdk/test_lazy_imports_real_io.py b/tests/integration/sdk/test_lazy_imports_real_io.py new file mode 100644 index 00000000000..6c1b54eb704 --- /dev/null +++ b/tests/integration/sdk/test_lazy_imports_real_io.py @@ -0,0 +1,23 @@ +import os +import subprocess +import sys +from typing import Final + +import litellm +import litellm.utils + + +def test_import_litellm_does_not_load_fastapi_or_bpe_table(): + result: Final = subprocess.run( + [ + sys.executable, + "-c", + "import sys, litellm; print(','.join(m for m in ('fastapi','starlette','litellm.proxy.proxy_cli','litellm.litellm_core_utils.default_encoding','openai.types.realtime') if m in sys.modules))", + ], + check=True, + capture_output=True, + text=True, + env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, + ) + + assert result.stdout.strip() == "" diff --git a/tests/integration/sdk/test_litellm_cli_real_io.py b/tests/integration/sdk/test_litellm_cli_real_io.py new file mode 100644 index 00000000000..9abf19ff9bd --- /dev/null +++ b/tests/integration/sdk/test_litellm_cli_real_io.py @@ -0,0 +1,384 @@ +import json +import os +import subprocess +import sys +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +import pytest + +import litellm +from litellm.cli.main import cli + +if sys.version_info >= (3, 11): + import tomllib +else: + import tomli as tomllib + +_HELP_ARGV: Final = ( + ("--help",), + ("logout", "--help"), + ("login", "--help"), + ("models", "list", "--help"), + ("claude", "--help"), + ("opencode", "--help"), + ("codex", "--help"), + ("pi", "--help"), + ("keys", "list", "--help"), +) + +_COMMAND_ARGV: Final = (("logout",), ("whoami",)) + +_RUN_EACH_ARGV: Final = """ +import contextlib, io, json, sys +from {module} import cli +runs = {{}} +for argv in json.loads(sys.argv[1]): + sys.argv = ["lite", *argv] + output = io.StringIO() + with contextlib.redirect_stdout(output): + try: + cli() + except SystemExit as stopped: + code = stopped.code + else: + code = 0 + runs[" ".join(argv)] = {{"exit": code, "output": output.getvalue()}} +print(json.dumps({{"runs": runs, "sdk_initialized": "completion" in vars(sys.modules["litellm"])}})) +""" + +_USE_SDK_AFTER_HELP: Final = """ +import builtins, importlib, json, sys +interpreter_imports = (builtins.__import__, importlib.import_module) +from {module} import cli +sys.argv = ["lite", "--help"] +try: + cli() +except SystemExit: + pass +before = "completion" in vars(sys.modules["litellm"]) +{first_use} +from litellm.router import Router +import litellm +import litellm.responses.main +from litellm import acompletion +print(json.dumps({{ + "before": before, + "router": Router.__name__, + "completion": callable(litellm.completion), + "acompletion": acompletion is litellm.acompletion, + "responses_is_function": callable(litellm.responses) and not isinstance(litellm.responses, type(litellm)), + "namespace": sorted(vars(litellm)), + "module_type": type(litellm).__name__, + "imports_restored": (builtins.__import__, importlib.import_module) == interpreter_imports, +}})) +""" + +_FIRST_USE_FROM_THREADS: Final = """ +import json, sys, threading +from litellm_cli.main import cli +sys.argv = ["lite", "--help"] +try: + cli() +except SystemExit: + pass +import litellm +start = threading.Barrier(4) +seen = [] +def use_sdk(): + start.wait() + try: + seen.append(callable(litellm.completion)) + except AttributeError as error: + seen.append(str(error)) +threads = [threading.Thread(target=use_sdk) for _ in range(4)] +for thread in threads: + thread.start() +for thread in threads: + thread.join() +print(json.dumps({"seen": seen})) +""" + +_STAR_IMPORT_AFTER_HELP: Final = """ +import json, sys +from litellm_cli.main import cli +sys.argv = ["lite", "--help"] +try: + cli() +except SystemExit: + pass +before = "completion" in vars(sys.modules["litellm"]) +from litellm import * +print(json.dumps({"before": before, "completion": callable(globals().get("completion")), "acompletion": callable(globals().get("acompletion"))})) +""" + +_SDK_WHOSE_SUBMODULE_READS_IT_WHILE_LOADING: Final = { + "sdk_sync.py": """ +import threading +initializing = threading.Event() +submodule_loading = threading.Event() +""", + "litellm/__init__.py": """ +import sdk_sync +sdk_sync.initializing.set() +sdk_sync.submodule_loading.wait(2) +from litellm import submodule +ready = True +""", + "litellm/submodule.py": """ +import litellm +import sdk_sync +sdk_sync.submodule_loading.set() +ready_when_loaded = getattr(litellm, "ready", False) +""", + "litellm/_process_defaults.py": "def apply_process_defaults() -> None: ...\n", + "litellm/proxy/__init__.py": "", + "litellm/proxy/client/__init__.py": "", + "litellm/proxy/client/cli/__init__.py": "", + "litellm/proxy/client/cli/main.py": """ +import json, threading +import litellm +import sdk_sync + +def cli(): + def initialize_sdk(): + litellm.ready + + def import_submodule_once_initialization_starts(): + sdk_sync.initializing.wait() + import litellm.submodule + + threads = [threading.Thread(target=initialize_sdk, daemon=True), + threading.Thread(target=import_submodule_once_initialization_starts, daemon=True)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(10) + print(json.dumps({"finished": [not thread.is_alive() for thread in threads]})) + +litellm_proxy_cli = cli +""", +} + +_SDK_WHOSE_INIT_FAILS: Final = { + "litellm/__init__.py": 'raise ImportError("a dependency of the SDK is broken")\n', + "litellm/utils.py": "", + "litellm/_process_defaults.py": "def apply_process_defaults() -> None: ...\n", + "litellm/proxy/__init__.py": "", + "litellm/proxy/client/__init__.py": "", + "litellm/proxy/client/cli/__init__.py": "", + "litellm/proxy/client/cli/main.py": """ +import json +import litellm + +def error_from(use): + try: + use() + except Exception as error: + return f"{type(error).__name__}: {error}" + return None + +def import_sdk_code(): + import litellm.utils + +def cli(): + print(json.dumps([error_from(lambda: litellm.completion), error_from(lambda: litellm.completion), error_from(import_sdk_code)])) + +litellm_proxy_cli = cli +""", +} + +_RUN_LITE_AGAINST_SDK_AT: Final = """ +import sys +sys.path.insert(0, sys.argv[1]) +from litellm_cli.main import cli +cli() +""" + +_PYPROJECT: Final = Path(__file__).resolve().parents[3] / "pyproject.toml" + +_RUN_ENTRY_POINT_HELP: Final = """ +import contextlib, io, json, sys +from {module} import {function} +imported_sdk = "litellm" in sys.modules +sys.argv = ["litellm", "--help"] +output = io.StringIO() +with contextlib.redirect_stdout(output): + try: + {function}() + except SystemExit as stopped: + code = stopped.code + else: + code = 0 +print(json.dumps({{"imported_sdk": imported_sdk, "exit": code, "help": output.getvalue()}})) +""" + +_WORKER_LOGGERS_UNDER_A_DEFAULT_LOG_CONFIG: Final = """ +import json, logging, os +os.environ["LITELLM_WORKER_LOG_CONFIG"] = json.dumps({{"version": 1, "root": {{"level": "INFO"}}}}) +from litellm_cli.main import proxy_worker_app +proxy_worker_app() +print(json.dumps({{name: logging.getLogger(name).disabled for name in {names!r}}})) +""" + +_COMMAND_MODULES_LOADED_BY: Final = """ +import json, sys +from litellm_cli.main import cli +sys.argv = ["lite", *sys.argv[1:]] +try: + cli() +except SystemExit: + pass +prefix = "litellm.proxy.client.cli.commands." +print(json.dumps({"modules": sorted(name[len(prefix):] for name in sys.modules if name.startswith(prefix))})) +""" + + +def _run(source: str, argv: tuple[str, ...], home: Path) -> dict[str, object]: + checkout: Final = str(_PYPROJECT.parent) + bootstrap: Final = ( + "import pathlib,sys\n" + f"sys.path.insert(0, {checkout!r})\n" + f"sys.path.insert(0, str(pathlib.Path({checkout!r}) / '.github/packaging'))\n" + "from cli_package import editable_support\n" + f"exec(editable_support(pathlib.Path({checkout!r}) / 'litellm/cli')[0][1])\n" + ) + result: Final = subprocess.run( + [sys.executable, "-I", "-c", bootstrap + source, *argv], + capture_output=True, + text=True, + timeout=120, + env={**os.environ, "HOME": str(home), "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, + ) + assert result.returncode == 0, result.stderr + return json.loads(result.stdout.rstrip("\n").rpartition("\n")[2]) + + +@pytest.fixture(scope="module") +def cli_runs(tmp_path_factory: pytest.TempPathFactory) -> Mapping[str, dict[str, object]]: + argv: Final = (json.dumps(_HELP_ARGV + _COMMAND_ARGV),) + return { + "thin": _run(_RUN_EACH_ARGV.format(module="litellm_cli.main"), argv, tmp_path_factory.mktemp("thin")), + "full": _run( + _RUN_EACH_ARGV.format(module="litellm.proxy.client.cli.main"), argv, tmp_path_factory.mktemp("full") + ), + } + + +def test_litellm_command_imports_the_proxy_only_when_it_runs(tmp_path: Path): + module, _, function = tomllib.loads(_PYPROJECT.read_text())["project"]["scripts"]["litellm"].partition(":") + run: Final = _run(_RUN_ENTRY_POINT_HELP.format(module=module, function=function), (), tmp_path) + assert run["imported_sdk"] is False + assert run["exit"] == 0 + assert "--num_workers" in str(run["help"]) + + +def test_a_worker_log_config_silences_the_sdk_loggers_but_not_the_ones_the_proxy_app_creates(tmp_path: Path): + disabled: Final = _run( + _WORKER_LOGGERS_UNDER_A_DEFAULT_LOG_CONFIG.format(names=("LiteLLM Proxy", "litellm.proxy.health_check")), + (), + tmp_path, + ) + assert disabled == {"LiteLLM Proxy": True, "litellm.proxy.health_check": False} + + +def test_lite_runs_commands_without_initializing_the_sdk(cli_runs: Mapping[str, dict[str, object]]): + assert cli_runs["thin"]["sdk_initialized"] is False + assert cli_runs["full"]["sdk_initialized"] is True + + +@pytest.mark.parametrize("argv", _HELP_ARGV + _COMMAND_ARGV, ids=" ".join) +def test_lite_prints_what_the_full_sdk_cli_prints(argv: tuple[str, ...], cli_runs: Mapping[str, dict[str, object]]): + thin: Final = cli_runs["thin"]["runs"] + full: Final = cli_runs["full"]["runs"] + assert isinstance(thin, dict) and isinstance(full, dict) + + assert thin[" ".join(argv)] == full[" ".join(argv)] + assert thin[" ".join(argv)]["output"] + + +@pytest.mark.parametrize( + ("argv", "loaded", "not_loaded"), + [ + (("logout", "--help"), {"auth"}, {"agents", "autoroute.commands", "models", "keys"}), + (("models", "list", "--help"), {"models"}, {"agents", "autoroute.commands", "keys"}), + (("claude", "--help"), {"agents"}, {"autoroute.commands", "models", "keys"}), + ], +) +def test_lite_imports_only_the_command_it_runs( + argv: tuple[str, ...], loaded: set[str], not_loaded: set[str], tmp_path: Path +): + modules: Final = _run(_COMMAND_MODULES_LOADED_BY, argv, tmp_path)["modules"] + assert isinstance(modules, list) + + assert loaded <= set(modules) + assert not not_loaded & set(modules) + + +@pytest.fixture(scope="module") +def sdk_used_after_full_cli(tmp_path_factory: pytest.TempPathFactory) -> dict[str, object]: + return _run( + _USE_SDK_AFTER_HELP.format(module="litellm.proxy.client.cli.main", first_use="pass"), + (), + tmp_path_factory.mktemp("full"), + ) + + +@pytest.mark.parametrize( + "first_use", + [ + "import litellm; litellm.completion", + "import litellm; litellm.types.utils.ModelResponse", + "from litellm import acompletion", + "from litellm.router import Router", + "from litellm.types import utils", + "import litellm.responses.main", + 'importlib.import_module("litellm.router")', + 'importlib.import_module(".utils", "litellm")', + 'exec("from .types import utils", {"__name__": "litellm.lazy_import", "__package__": "litellm"})', + ], +) +def test_sdk_initializes_in_place_on_first_use_after_lite_imports( + first_use: str, sdk_used_after_full_cli: dict[str, object], tmp_path: Path +): + thin: Final = _run(_USE_SDK_AFTER_HELP.format(module="litellm_cli.main", first_use=first_use), (), tmp_path) + + assert thin["before"] is False + assert sdk_used_after_full_cli["before"] is True + assert {**thin, "before": True} == sdk_used_after_full_cli + assert (thin["router"], thin["completion"], thin["acompletion"], thin["responses_is_function"]) == ( + "Router", + True, + True, + True, + ) + + +def test_star_import_after_lite_imports_gets_the_whole_sdk(tmp_path: Path): + assert _run(_STAR_IMPORT_AFTER_HELP, (), tmp_path) == {"before": False, "completion": True, "acompletion": True} + + +def test_threads_that_use_the_sdk_at_once_all_wait_for_one_initialization(tmp_path: Path): + used: Final = _run(_FIRST_USE_FROM_THREADS, (), tmp_path) + + assert used == {"seen": [True, True, True, True]} + + +def _write_sdk(root: Path, files: Mapping[str, str]) -> Path: + for relative_path, source in files.items(): + (root / relative_path).parent.mkdir(parents=True, exist_ok=True) + (root / relative_path).write_text(source) + return root + + +def test_importing_sdk_code_while_another_thread_initializes_the_sdk_does_not_hang(tmp_path: Path): + sdk_root: Final = _write_sdk(tmp_path / "sdk", _SDK_WHOSE_SUBMODULE_READS_IT_WHILE_LOADING) + + assert _run(_RUN_LITE_AGAINST_SDK_AT, (str(sdk_root),), tmp_path) == {"finished": [True, True]} + + +def test_every_use_after_a_failed_sdk_init_raises_the_init_error(tmp_path: Path): + sdk_root: Final = _write_sdk(tmp_path / "sdk", _SDK_WHOSE_INIT_FAILS) + + assert _run(_RUN_LITE_AGAINST_SDK_AT, (str(sdk_root),), tmp_path) == ["ImportError: a dependency of the SDK is broken"] * 3 diff --git a/tests/integration/sdk/test_long_decode_real_io.py b/tests/integration/sdk/test_long_decode_real_io.py new file mode 100644 index 00000000000..13147860cd2 --- /dev/null +++ b/tests/integration/sdk/test_long_decode_real_io.py @@ -0,0 +1,323 @@ +import json +import threading +from collections.abc import Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from itertools import count +from typing import Final + +import httpx +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa + +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +MESSAGES: Final = ({"role": "user", "content": "Prove it"},) + +HOP_1: Final = { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "hello"}]}, + "finishReason": "CONTINUATION", + "continuationToken": "dG9rLTE=", + } + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 30, "totalTokenCount": 40}, +} +HOP_2: Final = { + "candidates": [{"content": {"role": "model", "parts": [{"text": " world"}]}, "finishReason": "STOP"}], + "usageMetadata": { + "promptTokenCount": 10, + "cachedContentTokenCount": 9, + "candidatesTokenCount": 20, + "totalTokenCount": 30, + }, +} +STREAM_HOP_1_END: Final = { + "candidates": [{**HOP_1["candidates"][0], "content": {"role": "model", "parts": [{"text": "lo"}]}}], + "usageMetadata": HOP_1["usageMetadata"], +} + + +def _replay(hops: Sequence[httpx.Response], sent: list[dict[str, object]]) -> Callable[[httpx.Request], httpx.Response]: + def handler(request: httpx.Request) -> httpx.Response: + sent.append(json.loads(request.content)) + return hops[len(sent) - 1] + + return handler + + +def _sse(*chunks: Mapping[str, object]) -> httpx.Response: + return httpx.Response( + 200, + content="".join(f"data: {json.dumps(chunk)}\r\n\r\n" for chunk in chunks).encode(), + headers={"content-type": "text/event-stream"}, + ) + + +NO_CITATIONS_OR_URL_CONTEXT: Final = ( + pytest.param({}, id="missing"), + pytest.param({"citationMetadata": None, "urlContextMetadata": None}, id="null"), + pytest.param({"citationMetadata": {}, "urlContextMetadata": {}}, id="empty"), +) + + +STREAM_FIRST_CHUNK: Final = {"candidates": [{"content": {"role": "model", "parts": [{"text": "hel"}]}}]} + + +def _stream_hops() -> tuple[httpx.Response, httpx.Response]: + return (_sse(STREAM_FIRST_CHUNK, STREAM_HOP_1_END), _sse(HOP_2)) + + +def _stream_text(chunks: list) -> str: + return "".join(c.choices[0].delta.content or "" for c in chunks if c.choices) + + +def _stream_finish_reasons(chunks: list) -> list[str]: + return [c.choices[0].finish_reason for c in chunks if c.choices and c.choices[0].finish_reason] + + +def _delta_fields(chunks: list) -> list[Mapping[str, object]]: + return [ + choice.delta.provider_specific_fields + for chunk in chunks + for choice in chunk.choices # comprehension-ok: flattening stream chunks into their choices + if choice.delta.provider_specific_fields + ] + + +def _client_tokens(chunks: list) -> list[object]: + return [fields["continuation_token"] for fields in _delta_fields(chunks) if "continuation_token" in fields] + + +@dataclass(frozen=True, slots=True) +class _TokenEndpoint: + credentials: str + down: threading.Event + + +def _service_account_json(token_uri: str) -> str: + private_key: Final = ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()) + .decode() + ) + return json.dumps( + { + "type": "service_account", + "project_id": "p", + "private_key_id": "k", + "private_key": private_key, + "client_email": "sa@p.iam.gserviceaccount.com", + "token_uri": token_uri, + } + ) + + +@pytest.fixture +def token_endpoint() -> Iterator[_TokenEndpoint]: + down: Final = threading.Event() + minted: Final = count(1) + + class Handler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + self.rfile.read(int(self.headers["Content-Length"])) + refused: Final = down.is_set() + payload: Final = json.dumps( + {"error": "invalid_grant"} if refused else {"access_token": f"token-{next(minted)}", "expires_in": 0} + ).encode() + self.send_response(400 if refused else 200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, *_args: object) -> None: + return None + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + yield _TokenEndpoint(_service_account_json(f"http://127.0.0.1:{server.server_port}/token"), down) + server.shutdown() + server.server_close() + + +def _google_taking_each_access_token_once( + hops: Sequence[httpx.Response], used: list[str], on_answer: Callable[[], object] +) -> Callable[[httpx.Request], httpx.Response]: + def handler(request: httpx.Request) -> httpx.Response: + token: Final = request.headers["authorization"] + if token in used: + return httpx.Response(401, json={"error": {"code": 401, "status": "UNAUTHENTICATED"}}) + used.append(token) + on_answer() + return hops[len(used) - 1] + + return handler + + +async def _vertex_answer( + is_async: bool, + stream: bool, + handler: Callable[[httpx.Request], httpx.Response], + credentials: str, + extra_headers: Mapping[str, str] | None = None, +) -> tuple[str, list[str], list[object]]: + kwargs: Final = { + "model": "vertex_ai/gemini-3.8-flash", + "extra_headers": extra_headers, + "messages": list(MESSAGES), + "max_tokens": 100_000, + "stream": stream, + "vertex_project": "p", + "vertex_location": "global", + "vertex_credentials": credentials, + } + transport: Final = httpx.MockTransport(handler) + response: Final = ( + await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=transport)) + if is_async + else litellm.completion(**kwargs, client=HTTPHandler(client=httpx.Client(transport=transport))) + ) + if stream: + chunks: Final = [chunk async for chunk in response] if is_async else list(response) + return _stream_text(chunks), _stream_finish_reasons(chunks), _client_tokens(chunks) + choice: Final = response.choices[0] + fields: Final = choice.message.provider_specific_fields or {} + return choice.message.content, [choice.finish_reason], [fields["continuation_token"]] if fields else [] + + +def _two_hops(stream: bool) -> tuple[httpx.Response, httpx.Response]: + return _stream_hops() if stream else (httpx.Response(200, json=HOP_1), httpx.Response(200, json=HOP_2)) + + +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_vertex_long_decode_finishes_when_each_continuation_outlives_the_access_token_before_it( + is_async: bool, stream: bool, token_endpoint: _TokenEndpoint +): + used: Final[list[str]] = [] + google: Final = _google_taking_each_access_token_once(_two_hops(stream), used, lambda: None) + + answer: Final = await _vertex_answer(is_async, stream, google, token_endpoint.credentials) + + assert answer == ("hello world", ["stop"], []) + assert len(used) == 2 + + +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_vertex_long_decode_returns_the_answer_so_far_with_its_token_when_no_new_access_token_can_be_had( + is_async: bool, stream: bool, token_endpoint: _TokenEndpoint +): + used: Final[list[str]] = [] + google: Final = _google_taking_each_access_token_once(_two_hops(stream), used, token_endpoint.down.set) + + answer: Final = await _vertex_answer(is_async, stream, google, token_endpoint.credentials) + + assert answer == ("hello", ["length"], ["dG9rLTE="]) + assert len(used) == 1 + + +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_vertex_long_decode_keeps_the_callers_own_authorization_header_on_every_continuation( + is_async: bool, stream: bool, token_endpoint: _TokenEndpoint +): + hops: Final = _two_hops(stream) + seen: Final[list[str]] = [] + + def google(request: httpx.Request) -> httpx.Response: + seen.append(request.headers["authorization"]) + return hops[len(seen) - 1] + + answer: Final = await _vertex_answer( + is_async, stream, google, token_endpoint.credentials, {"Authorization": "Bearer mine"} + ) + + assert answer == ("hello world", ["stop"], []) + assert seen == ["Bearer mine", "Bearer mine"] + + +BUDGET_ENV: Final = "LITELLM_GEMINI_DEFAULT_MAX_OUTPUT_TOKENS" +ARGON: Final = "gemini/gemini-4-argon" +OUTPUT_LIMIT_KEYS: Final = ("max_output_tokens", "maxOutputTokens") + + +def _output_limits(request: Mapping[str, object]) -> dict[str, object]: + generation_config: Final = request.get("generationConfig", {}) + assert isinstance(generation_config, dict) + return {key: value for key, value in generation_config.items() if key in OUTPUT_LIMIT_KEYS} + + +def _cost_map_max_output_tokens(model: str) -> int: + provider, name = model.split("/", 1) + max_output_tokens: Final = litellm.get_model_info(model=name, custom_llm_provider=provider)["max_output_tokens"] + assert isinstance(max_output_tokens, int) + return max_output_tokens + + +def _gemini_chat_models_answering_no_text() -> tuple[str, ...]: + rows: Final = ( + (name, row) + for name, row in litellm.model_cost.items() + if row.get("mode") == "chat" and "text" not in row.get("supported_output_modalities", ("text",)) + ) + return tuple( + name if name.startswith("gemini/") else f"vertex_ai/{name}" + for name, row in rows + if row.get("litellm_provider") in ("gemini", "vertex_ai-language-models") + ) + + +@pytest.mark.asyncio +async def test_env_budget_skips_every_chat_model_the_cost_map_lists_as_answering_no_text( + monkeypatch: pytest.MonkeyPatch, token_endpoint: _TokenEndpoint +): + monkeypatch.setenv(BUDGET_ENV, "50") + models: Final = _gemini_chat_models_answering_no_text() + sent: Final[list[dict[str, object]]] = [] + replies: Final = tuple(httpx.Response(200, json=HOP_2) for _ in models) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_replay(replies, sent))) + + for model in models: + await litellm.acompletion( + model=model, + messages=list(MESSAGES), + api_key="fake-gemini-key", + vertex_project="p", + vertex_location="us-central1", + vertex_credentials=token_endpoint.credentials, + client=client, + ) + + assert models + assert {model: _output_limits(request) for model, request in zip(models, sent, strict=True)} == dict.fromkeys( + models, {} + ) + + +@pytest.mark.parametrize("model", ["gemini/gemini-2.5-flash", "vertex_ai/gemini-2.5-flash", ARGON]) +@pytest.mark.asyncio +async def test_env_budget_above_a_models_output_limit_sends_that_limit_instead( + monkeypatch: pytest.MonkeyPatch, token_endpoint: _TokenEndpoint, model: str +): + budget: Final = 2 * _cost_map_max_output_tokens(model) + monkeypatch.setenv(BUDGET_ENV, str(budget)) + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=model, + messages=list(MESSAGES), + api_key="fake-gemini-key", + vertex_project="p", + vertex_location="us-central1", + vertex_credentials=token_endpoint.credentials, + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {"max_output_tokens": _cost_map_max_output_tokens(model)} diff --git a/tests/integration/sdk/test_router_fuzz_regressions_rust_only_real_io.py b/tests/integration/sdk/test_router_fuzz_regressions_rust_only_real_io.py new file mode 100644 index 00000000000..13548b34fd3 --- /dev/null +++ b/tests/integration/sdk/test_router_fuzz_regressions_rust_only_real_io.py @@ -0,0 +1,193 @@ +""" +Router behavior the router-code differential fuzz run found changed once the Rust core became the +only core: count settings below zero, an allowed_fails_policy write that is no policy, a model +response object as a deployment's mock_response, and endpoints read or patched on the Router class. +""" + +import asyncio +import json +import socket +import threading +from collections.abc import Iterator, Mapping +from datetime import datetime +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final + +import pytest + +import litellm +from litellm import Router + +MESSAGES: Final = [{"role": "user", "content": "hi"}] + + +class _Upstream(BaseHTTPRequestHandler): + """An OpenAI upstream whose path's first segment names how it answers: ``cut`` sends two + Responses stream events then cuts the chunked body, ``filtered`` answers a chat completion the + provider blocked, ``down`` answers 502, and anything else a plain chat completion.""" + + protocol_version = "HTTP/1.1" + + def log_message(self, format: str, *args: object) -> None: # noqa: A002 # the base class's name + del format, args + + def _json(self, status: int, body: Mapping[str, object]) -> None: + data: Final = json.dumps(body).encode() + self.send_response(status) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(data))) + self.end_headers() + _ = self.wfile.write(data) + + def _chunk(self, event: Mapping[str, object]) -> None: + data: Final = f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + _ = self.wfile.write(f"{len(data):x}\r\n".encode() + data + b"\r\n") + self.wfile.flush() + + def _cut(self) -> None: + self.send_response(200) + self.send_header("content-type", "text/event-stream") + self.send_header("transfer-encoding", "chunked") + self.end_headers() + response: Final = {"id": "r1", "object": "response", "created_at": 1, "status": "in_progress", "output": []} + self._chunk({"type": "response.created", "sequence_number": 0, "response": {**response, "model": "gpt-4o"}}) + self._chunk( + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "m1", + "output_index": 0, + "content_index": 0, + "delta": "hello ", + } + ) + self.connection.shutdown(socket.SHUT_RDWR) + self.close_connection = True + + def do_POST(self) -> None: + _ = self.rfile.read(int(self.headers.get("content-length") or 0)) + mode: Final = self.path.split("/")[1] + if mode == "cut": + self._cut() + return + if mode == "down": + self._json(502, {"error": {"message": "provider error", "type": "api_error"}}) + return + finish: Final = "content_filter" if mode == "filtered" else "stop" + message: Final = {"role": "assistant", "content": mode} + choice: Final = {"index": 0, "finish_reason": finish, "message": message} + usage: Final = {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + self._json(200, {"id": "c1", "object": "chat.completion", "created": 1, "choices": [choice], "usage": usage}) + + +@pytest.fixture +def upstream() -> Iterator[str]: + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), _Upstream) + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield f"http://127.0.0.1:{server.server_address[1]}" + server.shutdown() + thread.join() + + +def _served(name: str, deployment_id: str, model: str, api_base: str) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {"model": model, "api_key": "k", "api_base": api_base}, + "model_info": {"id": deployment_id}, + } + + +def _request_logging(call_type: str, *, stream: bool) -> object: + from litellm.litellm_core_utils.litellm_logging import Logging + + return Logging( + model="g", + messages=MESSAGES, + stream=stream, + call_type=call_type, + start_time=datetime.now(), # noqa: DTZ005 # litellm's own naive start time + litellm_call_id="call-1", + function_id="f", + ) + + +async def _recorded() -> None: + """Let the background recordings of a stream's outcome run.""" + for _ in range(50): + await asyncio.sleep(0.02) + + +@pytest.mark.asyncio +async def test_a_responses_stream_cut_after_it_opened_counts_one_fail_with_a_shared_logging_object( + upstream: str, +) -> None: + router: Final = Router( + model_list=[_served("g", "d1", "openai/gpt-4o", f"{upstream}/cut/v1")], num_retries=0, disable_cooldowns=True + ) + + stream: Final = await router.aresponses( + model="g", input="hi", stream=True, litellm_logging_obj=_request_logging("aresponses", stream=True) + ) + with pytest.raises(Exception): # noqa: B017, PT011 # the transport's read error, whatever class maps it + async for _ in stream: # pyright: ignore[reportGeneralTypeIssues] # a stream when stream=True + pass + await _recorded() + + assert router.health().store.get_cache(key="d1:fails") == 1 + + +@pytest.mark.asyncio +async def test_a_blocked_response_raised_for_a_fallback_counts_no_success_with_a_shared_logging_object( + upstream: str, +) -> None: + router: Final = Router( + model_list=[ + _served("a", "a1", "openai/gpt-4o", f"{upstream}/filtered/v1"), + _served("safe", "s1", "openai/gpt-4o", f"{upstream}/ok/v1"), + ], + content_policy_fallbacks=[{"a": ["safe"]}], + num_retries=0, + disable_cooldowns=True, + ) + + response: Final = await router.acompletion( + model="a", messages=MESSAGES, litellm_logging_obj=_request_logging("acompletion", stream=False) + ) + await _recorded() + + assert response.choices[0].message.content == "ok" + assert router.health().store.get_cache(key="a1:successes") is None + assert router.health().store.get_cache(key="s1:successes") == 1 + + +@pytest.mark.asyncio +async def test_a_text_completion_stream_that_fails_before_its_first_chunk_counts_nothing(upstream: str) -> None: + router: Final = Router( + model_list=[_served("g", "d1", "openai/gpt-3.5-turbo-instruct", f"{upstream}/down/v1")], + num_retries=0, + disable_cooldowns=True, + ) + + stream: Final = await router.atext_completion(model="g", prompt="hi", stream=True) + with pytest.raises(litellm.BadGatewayError): + async for _ in stream: # pyright: ignore[reportGeneralTypeIssues] # a stream when stream=True + pass + await _recorded() + + assert router.health().store.get_cache(key="d1:successes") is None + assert router.health().store.get_cache(key="d1:fails") is None + + +@pytest.mark.asyncio +async def test_a_request_refused_before_the_provider_counts_against_no_deployment(upstream: str) -> None: + router: Final = Router( + model_list=[_served("g", "d1", "openai/gpt-4o", f"{upstream}/ok/v1")], num_retries=0, disable_cooldowns=True + ) + + with pytest.raises(litellm.UnsupportedParamsError): + await router.acompletion(model="g", messages=MESSAGES, reasoning_effort="low") + await _recorded() + + assert router.health().store.get_cache(key="d1:fails") is None + assert not [key for key in router.health().store.in_memory_cache.cache_dict if ":rpm" in key] diff --git a/tests/integration/sdk/test_router_host_real_io.py b/tests/integration/sdk/test_router_host_real_io.py new file mode 100644 index 00000000000..6ad80a90aa1 --- /dev/null +++ b/tests/integration/sdk/test_router_host_real_io.py @@ -0,0 +1,28 @@ +import subprocess +import sys + +import pytest + +import litellm.router_host as router_host_module +from litellm.router_host import router_host + + +def test_router_uses_the_proxy_host_without_importing_the_proxy_first(): + # A fresh interpreter, so no earlier test has imported litellm.proxy. + script = ( + "import sys\n" + "from litellm.router import Router\n" + "from litellm.router_host import router_host\n" + "assert 'litellm.proxy' not in sys.modules, 'importing the router imported the proxy'\n" + "print(type(router_host()).__name__)\n" + ) + result = subprocess.run([sys.executable, "-P", "-c", script], capture_output=True, text=True, check=True) + assert result.stdout.strip() == "ProxyRouterHost" + + +@pytest.fixture(autouse=True) +def nothing_installed(): + saved = router_host_module._installed[0] + router_host_module._installed[0] = None + yield + router_host_module._installed[0] = saved diff --git a/tests/integration/sdk/test_streaming_handler_real_io.py b/tests/integration/sdk/test_streaming_handler_real_io.py new file mode 100644 index 00000000000..5acd6597d99 --- /dev/null +++ b/tests/integration/sdk/test_streaming_handler_real_io.py @@ -0,0 +1,367 @@ +import itertools +import json +import os +import subprocess +import sys +from collections.abc import Mapping +from pathlib import Path +from types import MappingProxyType +from typing import Final +import pytest +from pydantic import TypeAdapter +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage +from litellm.utils import ModelResponseListIterator +from litellm import completion + + +bedrock_chunks = [ + ModelResponseStream( + id="chatcmpl-d249def8-a78b-464c-87b5-3a6f43565292", + created=1742056047, + model=None, + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + provider_specific_fields=None, + content="I'm Claude", + role="assistant", + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields={}, + usage=None, + ), + ModelResponseStream( + id="chatcmpl-fe559823-b383-4249-ab87-52f6ad9d08c2", + created=1742056047, + model=None, + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + provider_specific_fields=None, + content=", an AI", + role="assistant", + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields={}, + usage=None, + ), + ModelResponseStream( + id="chatcmpl-c1c6cc2f-75b9-4a24-88b9-4e5aacd0268b", + created=1742056047, + model=None, + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta( + provider_specific_fields=None, + content="", + role="assistant", + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields={}, + usage=None, + ), +] + + +FINISH_ONLY_LOGPROBS_STREAM_SCRIPT: Final = """ +import json +import sys + +sys.path.insert(0, sys.argv[1]) + +import httpx +from openai import OpenAI + +import litellm + +upstream_choices = ( + {"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None}, + { + "index": 0, + "delta": {}, + "logprobs": {"content": [{"token": "<|return|>", "logprob": -0.5, "top_logprobs": []}]}, + "finish_reason": "stop", + }, +) +body = "".join( + "data: " + + json.dumps({"id": "c", "object": "chat.completion.chunk", "created": 1, "model": "m", "choices": [choice]}) + + "\\n\\n" + for choice in upstream_choices +) + "data: [DONE]\\n\\n" +client = OpenAI( + api_key="unused", + base_url="https://provider.example/v1", + http_client=httpx.Client( + transport=httpx.MockTransport( + lambda request: httpx.Response(200, content=body, headers={"content-type": "text/event-stream"}) + ) + ), +) +stream = litellm.completion( + model="openai/m", messages=[{"role": "user", "content": "hi"}], stream=True, logprobs=True, client=client +) +print(json.dumps([json.loads(chunk.model_dump_json(exclude_none=True)) for chunk in stream])) +""" + + +def test_stream_with_logprobs_only_on_finish_chunk_serializes_in_a_fresh_process() -> None: + result: Final = subprocess.run( + [sys.executable, "-I", "-c", FINISH_ONLY_LOGPROBS_STREAM_SCRIPT, str(Path(litellm.__file__).parent.parent)], + capture_output=True, + text=True, + timeout=120, + env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, + ) + + assert result.returncode == 0, result.stdout + result.stderr + wire_choices: Final = tuple( + itertools.chain.from_iterable(chunk["choices"] for chunk in json.loads(result.stdout.strip().splitlines()[-1])) + ) + finish_choices: Final = [choice for choice in wire_choices if choice.get("finish_reason")] + assert [choice["finish_reason"] for choice in finish_choices] == ["stop"] + assert [entry["token"] for entry in finish_choices[0]["logprobs"]["content"]] == ["<|return|>"] + + +_REPETITION_TEST_CASES = [ + # Basic cases + pytest.param( + ["same"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT, + True, + id="all_identical_raises", + ), + pytest.param( + ["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT - 1), + False, + id="below_threshold_no_raise", + ), + pytest.param( + [None] * litellm.REPEATED_STREAMING_CHUNK_LIMIT, + False, + id="none_content_no_raise", + ), + pytest.param( + [""] * litellm.REPEATED_STREAMING_CHUNK_LIMIT, + False, + id="empty_content_no_raise", + ), + # Short content (len <= 2) should not raise + pytest.param( + ["##"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT, + False, + id="short_content_2chars_no_raise", + ), + pytest.param( + ["{"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT, + False, + id="short_content_1char_no_raise", + ), + pytest.param( + ["ab"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT, + False, + id="short_content_2chars_ab_no_raise", + ), + # All different chunks + pytest.param( + ["diff"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT, + False, + id="all_different_no_raise", + ), + # One chunk different at various positions + pytest.param( + ["different_first"] + ["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT - 1), + False, + id="first_chunk_different_no_raise", + ), + pytest.param( + ["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT - 1) + ["different_last"], + False, + id="last_chunk_different_no_raise", + ), + pytest.param( + ["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT // 2 + 1) + + ["different_mid"] + + ["same"] + * ( + litellm.REPEATED_STREAMING_CHUNK_LIMIT + - litellm.REPEATED_STREAMING_CHUNK_LIMIT // 2 + + 1 + ), + False, + id="middle_chunk_different_no_raise", + ), + pytest.param( + ["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT - 2) + ["diff", "diff"], + False, + id="last_two_different_no_raise", + ), + pytest.param( + ["diff"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT + + ["same"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT + + ["diff"], + True, + id="in_between_same_and_diff_raise", + ), +] + + +# Azure streaming chunks that reproduce issue #24221: +# Azure sends an initial chunk with prompt_filter_results and choices=[], +# then a chunk with role='assistant' and content='', then content chunks. +# With stream_options.include_usage=True, the empty-choices chunk was +# forwarded with an inflated default choice, consuming the sent_first_chunk +# flag and causing strip_role_from_delta to strip the role from the real +# first chunk. +_AZURE_CHUNKS_WITH_PROMPT_FILTER = [ + # Chunk 1: prompt_filter_results, no choices (Azure-specific) + ModelResponseStream( + id="chatcmpl-abc123", + created=1742056047, + model=None, + object="chat.completion.chunk", + choices=[], + usage=None, + ), + # Chunk 2: first real chunk with role='assistant' and empty content + ModelResponseStream( + id="chatcmpl-abc123", + created=1742056047, + model=None, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="", role="assistant"), + ) + ], + usage=None, + ), + # Chunk 3: content + ModelResponseStream( + id="chatcmpl-abc123", + created=1742056047, + model=None, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="Hello!"), + ) + ], + usage=None, + ), + # Chunk 4: finish_reason + ModelResponseStream( + id="chatcmpl-abc123", + created=1742056047, + model=None, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(), + ) + ], + usage=None, + ), + # Chunk 5: final usage chunk, no choices + ModelResponseStream( + id="chatcmpl-abc123", + created=1742056047, + model=None, + object="chat.completion.chunk", + choices=[], + usage=Usage( + completion_tokens=10, + prompt_tokens=20, + total_tokens=30, + ), + ), +] + + +_AZURE_CHOICE_EXTRAS: Final = MappingProxyType( + {"content_filter_results": {"hate": {"filtered": False, "severity": "safe"}}} +) +_OPENAI_AND_AZURE_CHOICE_EXTRAS: Final = ( + pytest.param("openai", MappingProxyType({}), id="openai"), + pytest.param("azure", _AZURE_CHOICE_EXTRAS, id="azure"), +) + + +_TOKEN_IDS: Final = TypeAdapter(list[int] | None) +_CHOICE_INDEX: Final = TypeAdapter(int) +_DELTA: Final = TypeAdapter(Mapping[str, object]) + + +def _streaming_wrapper_for_unit_test( + model: str, + custom_llm_provider: str, + chunks: list[ModelResponseStream], +) -> CustomStreamWrapper: + return CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=chunks), + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=Logging( + model=model, + messages=[{"role": "user", "content": "unit test"}], + stream=True, + call_type="completion", + start_time=0.0, + litellm_call_id="unit-test-call", + function_id="unit-test-function", + ), + ) + + +def _unit_test_streaming_chunk( + content: str | None, + finish_reason: str | None = None, + role: str | None = "assistant", +) -> ModelResponseStream: + return ModelResponseStream( + id="unit-test-stream", + created=1, + model="unit-test-model", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content=content, role=role), + finish_reason=finish_reason, + ) + ], + ) diff --git a/tests/integration/sdk/test_token_counter_real_io.py b/tests/integration/sdk/test_token_counter_real_io.py new file mode 100644 index 00000000000..fd58cb29d24 --- /dev/null +++ b/tests/integration/sdk/test_token_counter_real_io.py @@ -0,0 +1,207 @@ +#### What this tests #### +# This tests litellm.token_counter.token_counter() function +import json +import os +import subprocess +import sys +from collections.abc import Mapping +from pathlib import Path +from typing import Final +import tiktoken +from tokenizers import Regex, Tokenizer, models, pre_tokenizers + +from unittest.mock import AsyncMock + +import litellm +from litellm import decode, encode +import litellm.constants +from litellm.types.utils import SelectTokenizerResponse + + +## Choose which token_counter the test will use. + + +# test_token_counter_normal_plus_function_calling() + + +# test_tokenizers() + + +def test_encoding_and_decoding(tmp_path: Path): + sample_text = "Hellö World, this is my input string!" + + # openai encoding + decoding + openai_tokens = encode(model="gpt-3.5-turbo", text=sample_text) + openai_text = decode(model="gpt-3.5-turbo", tokens=openai_tokens) + + assert openai_text == sample_text + + # claude encoding + decoding + claude_tokens = encode(model="claude-3-5-haiku-20241022", text=sample_text) + + claude_text = decode(model="claude-3-5-haiku-20241022", tokens=claude_tokens) + + assert claude_text == sample_text + + # cohere encoding + decoding + cohere_tokens = encode(model="command-nightly", text=sample_text) + cohere_text = decode(model="command-nightly", tokens=cohere_tokens) + + assert cohere_text == sample_text + + # llama2 encoding + decoding + words = sample_text.split() + result = _run_in_memory_hub( + HUB_ROUND_TRIP_SCRIPT, + { + "hf-internal-testing/llama-tokenizer": _word_level_tokenizer_json( + pre_tokenizers.WhitespaceSplit(), + vocab={"[UNK]": 0, **{word: i + 1 for i, word in enumerate(words)}}, + ) + }, + sample_text, + tmp_path, + ) + + assert result["decoded"] == sample_text + assert result["requested"] == ["hf-internal-testing/llama-tokenizer"] + assert len(result["tokens"]) == len(words) + assert len(result["tokens"]) != len(encode(model="gpt-3.5-turbo", text=sample_text)) + + +# test_encoding_and_decoding() + + +# test_gpt_vision_token_counting() + +from litellm.utils import _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES, claude_json_str + +# Clear the cache at module load to ensure clean state +_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + + +HUB_SETUP_SCRIPT: Final = """ +import json +import sys +sys.path.insert(0, sys.argv[1]) +import httpx +import huggingface_hub +import litellm +served = json.loads(sys.argv[2]) +text = sys.argv[3] +requested = [] +def handle(request): + repo = request.url.path.lstrip("/").split("/resolve/")[0] + if repo not in served or not request.url.path.endswith("/tokenizer.json"): + return httpx.Response(404) + requested.append(repo) + payload = served[repo].encode() + headers = {"content-length": str(len(payload)), "etag": '"fixture"', "x-repo-commit": "a" * 40} + return httpx.Response(200, headers=headers, content=payload if request.method == "GET" else b"") +huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(handle))) +""" + +HUB_TOKENIZER_SCRIPT: Final = HUB_SETUP_SCRIPT + """ +litellm.cohere_models = {"command-r-v1"} +litellm.anthropic_models = {"claude-2"} +custom = litellm.create_pretrained_tokenizer("Xenova/llama-3-tokenizer") +print(json.dumps({ + "llama2": litellm.token_counter(model="meta-llama/Llama-2-7b-chat", text=text), + "llama3": litellm.token_counter(model="meta-llama/llama-3-70b-instruct", text=text), + "cohere": litellm.token_counter(model="command-r-v1", text=text), + "anthropic": litellm.token_counter(model="claude-2", text=text), + "custom": litellm.token_counter(custom_tokenizer=custom, text=text), + "requested": sorted(set(requested)), +})) +""" + +HUB_ROUND_TRIP_SCRIPT: Final = HUB_SETUP_SCRIPT + """ +tokens = litellm.encode(model="meta-llama/Llama-2-7b-chat", text=text) +print(json.dumps({"tokens": tokens, "decoded": litellm.decode(model="meta-llama/Llama-2-7b-chat", tokens=tokens), "requested": sorted(set(requested))})) +""" + +HUB_RETRY_DURING_EXIT_SCRIPT: Final = HUB_SETUP_SCRIPT + """ +import threading +import time +from litellm.utils import _HuggingFaceTokenizerLoadOutcomes, _return_huggingface_tokenizer +retry_reached_hub = threading.Event() +def answer_once_the_process_is_exiting(request): + retry_reached_hub.set() + while threading.main_thread().is_alive(): + time.sleep(0.01) + time.sleep(0.5) + return handle(request) +now = [0.0] +outcomes = _HuggingFaceTokenizerLoadOutcomes(clock=lambda: now[0]) +model = "meta-llama/llama-3-70b-instruct" +_return_huggingface_tokenizer(model, outcomes) +huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(answer_once_the_process_is_exiting))) +now[0] = 3600.0 +_return_huggingface_tokenizer(model, outcomes) +print(json.dumps({"retry_reached_hub": retry_reached_hub.wait(timeout=60)})) +""" + + +def _word_level_tokenizer_json( + pre_tokenizer: pre_tokenizers.PreTokenizer, vocab: Mapping[str, int] | None = None +) -> str: + tokenizer: Final = Tokenizer( + models.WordLevel(vocab=dict(vocab) if vocab is not None else {"[UNK]": 0}, unk_token="[UNK]") + ) + tokenizer.pre_tokenizer = pre_tokenizer + return tokenizer.to_str() + + +def _run_in_memory_hub(script: str, served: dict[str, str], text: str, tmp_path: Path) -> dict: + result: Final = subprocess.run( + [ + sys.executable, + "-I", + "-c", + script, + str(Path(litellm.__file__).parent.parent), + json.dumps(served), + text, + ], + capture_output=True, + text=True, + timeout=60, + env={ + **os.environ, + "HF_HOME": str(tmp_path / "home"), + "HF_HUB_CACHE": str(tmp_path / "cache"), + "HF_ENDPOINT": "http://127.0.0.1:9", + "HF_HUB_OFFLINE": "0", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + }, + ) + + assert result.returncode == 0, result.stdout + result.stderr + return json.loads(result.stdout.strip().splitlines()[-1]) + + +def test_token_counter_uses_the_tokenizer_of_each_model_family_and_of_a_custom_tokenizer(tmp_path: Path) -> None: + sample: Final = "Tokenizers disagree: anthropic, tiktoken; llama-2 & llama-3!" + served: Final = { + "hf-internal-testing/llama-tokenizer": _word_level_tokenizer_json(pre_tokenizers.WhitespaceSplit()), + "Xenova/llama-3-tokenizer": _word_level_tokenizer_json(pre_tokenizers.Split(Regex("."), "isolated")), + "Xenova/c4ai-command-r-v01-tokenizer": _word_level_tokenizer_json(pre_tokenizers.Whitespace()), + } + expected: Final = {repo: len(Tokenizer.from_str(payload).encode(sample).ids) for repo, payload in served.items()} + anthropic_count: Final = len(Tokenizer.from_str(claude_json_str).encode(sample).ids) + tiktoken_count: Final = litellm.token_counter(model="gpt-3.5-turbo", text=sample) + assert len({*expected.values(), anthropic_count, tiktoken_count}) == len(expected) + 2 + + counts: Final = _run_in_memory_hub(HUB_TOKENIZER_SCRIPT, served, sample, tmp_path) + assert counts == { + "llama2": expected["hf-internal-testing/llama-tokenizer"], + "llama3": expected["Xenova/llama-3-tokenizer"], + "cohere": expected["Xenova/c4ai-command-r-v01-tokenizer"], + "anthropic": anthropic_count, + "custom": expected["Xenova/llama-3-tokenizer"], + "requested": sorted(served), + } + + +def test_process_exits_cleanly_while_a_background_tokenizer_retry_is_downloading(tmp_path: Path) -> None: + assert _run_in_memory_hub(HUB_RETRY_DURING_EXIT_SCRIPT, {}, "", tmp_path) == {"retry_reached_hub": True} diff --git a/tests/integration/sdk/test_tokenizer_real_io.py b/tests/integration/sdk/test_tokenizer_real_io.py new file mode 100644 index 00000000000..b6b598b81ea --- /dev/null +++ b/tests/integration/sdk/test_tokenizer_real_io.py @@ -0,0 +1,119 @@ +import os +import subprocess +import sys +from pathlib import Path +from typing import Final + +import pytest + +import litellm +from tests.unit.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON + + +@pytest.mark.parametrize("offline", ("0", "1")) +def test_hub_loader_preserves_environment_auth_cache_and_offline(tmp_path: Path, offline: str) -> None: + script: Final = """ +import json +import sys +from pathlib import Path +sys.path.insert(0, sys.argv[1]) +import httpx +import huggingface_hub +from huggingface_hub.errors import LocalEntryNotFoundError +import litellm +payload = sys.argv[2].encode() +offline = sys.argv[3] == "1" +observed = [] +def handle(request): + assert not offline, "offline loading issued a request" + if request.url.path.endswith("/tokenizer.json"): + observed.append(request.headers.get("authorization")) + if request.headers.get("authorization") != "Bearer audit-fixture-token": + return httpx.Response(401) + return httpx.Response(200, headers={"content-length": str(len(payload)), "etag": '"fixture"', "x-repo-commit": "a" * 40}, content=payload if request.method == "GET" else b"") +if not offline: + huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(handle))) +try: + tokenizer = litellm.create_pretrained_tokenizer("test-fixture/tokenizer")["tokenizer"] +except LocalEntryNotFoundError: + assert offline + assert observed == [] +else: + assert not offline + assert "Bearer audit-fixture-token" in observed + assert tokenizer.decode(tokenizer.encode("Hello World").ids) == "Hello World" + assert tuple(Path(sys.argv[4]).rglob("tokenizer.json")) +print("compatible") +""" + result: Final = subprocess.run( + [ + sys.executable, + "-I", + "-c", + script, + str(Path(litellm.__file__).parent.parent), + TOKENIZER_JSON, + offline, + str(tmp_path / "cache"), + ], + capture_output=True, + text=True, + timeout=30, + env={ + **os.environ, + "HF_HOME": str(tmp_path / "home"), + "HF_HUB_CACHE": str(tmp_path / "cache"), + "HF_ENDPOINT": "http://127.0.0.1:9", + "HF_TOKEN": "audit-fixture-token", + "HF_HUB_OFFLINE": offline, + "HF_HUB_DISABLE_IMPLICIT_TOKEN": "0", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + }, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert result.stdout.strip() == "compatible" + + +@pytest.mark.parametrize("rust", (None, "0", "1")) +def test_tokenization_stays_offline(tmp_path: Path, rust: str | None) -> None: + script: Final = """ +import sys +sys.path.insert(0, sys.argv[1]) +def reject_network(event, args): + if event == "socket.connect": + raise AssertionError("tokenizer attempted a network connection") +sys.addaudithook(reject_network) +import litellm +from litellm.rust_bridge.tokenizer import get_encoding +import tiktoken +from tokenizers import Tokenizer +assert isinstance(litellm.encoding, tiktoken.Encoding) +for name in ("cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit"): + encoding = get_encoding(name) + text = "offline café 漢字 🙂" + " " * 64 + assert encoding.decode(encoding.encode(text)) == text +ids = litellm.encode(text="hello world") +assert litellm.decode(tokens=ids) == "hello world" +assert litellm.token_counter(model=None, text="hello world") == len(ids) +custom = litellm.create_tokenizer(sys.argv[2]) +assert isinstance(custom["tokenizer"], Tokenizer) +custom["tokenizer"].enable_padding(pad_id=0, pad_token="[UNK]") +assert litellm.decode(tokens=litellm.encode(text="Hello World", custom_tokenizer=custom), custom_tokenizer=custom) == "Hello World" +print("compatible") +""" + result: Final = subprocess.run( + [sys.executable, "-I", "-c", script, str(Path(litellm.__file__).parent.parent), TOKENIZER_JSON], + capture_output=True, + text=True, + timeout=30, + cwd=tmp_path, + env={ + **{key: value for key, value in os.environ.items() if key != "LITELLM_RUST"}, + **({"LITELLM_RUST": rust} if rust is not None else {}), + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "TIKTOKEN_CACHE_DIR": str(tmp_path / "unused-tokenizer-cache"), + }, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert result.stdout.strip() == "compatible" + assert not (tmp_path / "unused-tokenizer-cache").exists() diff --git a/tests/unit/caching/test_evicted_client_closer.py b/tests/unit/caching/test_evicted_client_closer.py index 936eb37a733..5c6c9cb92b0 100644 --- a/tests/unit/caching/test_evicted_client_closer.py +++ b/tests/unit/caching/test_evicted_client_closer.py @@ -18,7 +18,6 @@ import pytest from redis.asyncio import ConnectionPool, Redis from litellm.caching.evicted_client_closer import EvictedClientCloser -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from tests.unit.caching._event_loop_reuse import runs_until_a_loop_address_repeats @@ -111,19 +110,6 @@ async def test_redis_client_is_closed_only_after_its_subscription_releases_the_c assert closer.pending_count == 0 -async def _trickling_upstream(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: - """Serves a chunked body slowly, so a request stays on the wire long enough to observe.""" - await reader.read(4096) - writer.write(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n") - await writer.drain() - for _ in range(6): - writer.write(b"5\r\nhello\r\n") - await writer.drain() - await asyncio.sleep(0.1) - writer.write(b"0\r\n\r\n") - await writer.drain() - - @pytest.mark.asyncio async def test_owned_client_is_closed_once_the_grace_window_elapses(): clock = FakeClock() @@ -416,79 +402,6 @@ async def test_a_client_evicted_on_a_collected_loop_is_dropped(): assert closer.pending_count == 0 -@pytest.mark.asyncio -async def test_a_client_serving_a_request_is_not_closed_when_its_grace_window_ends(): - """The grace window on its own cannot promise that a request has finished. - - ``litellm.request_timeout`` defaults to 6000 seconds and a streaming response - is bounded only by how long the upstream keeps sending, so a client past its - deadline is closed only once its own pool reports nothing in flight. - """ - server = await asyncio.start_server(_trickling_upstream, "127.0.0.1", 0) - port = server.sockets[0].getsockname()[1] - clock = FakeClock() - closer = make_closer(clock) - client = httpx.AsyncClient() - - closer.mark_owned(client) - closer.schedule(client) - - async def read_the_stream() -> int: - received = 0 - async with client.stream("GET", f"http://127.0.0.1:{port}/") as response: - async for chunk in response.aiter_bytes(): - received += len(chunk) - return received - - streaming = asyncio.create_task(read_the_stream()) - await asyncio.sleep(0.25) # the request is on the wire - clock.advance(3600.0) # and its grace window is long gone - closer.reap() - await asyncio.sleep(0.05) - - assert client.is_closed is False, "closed a client that was serving a request" - assert await streaming > 0, "the in-flight request did not survive the reap" - - clock.advance(3600.0) - closer.reap() - await asyncio.sleep(0.05) - - assert client.is_closed is True, "an idle client past its grace window must be closed" - assert closer.pending_count == 0 - server.close() - - -@pytest.mark.asyncio -async def test_the_aiohttp_backed_handler_is_not_closed_mid_request(): - """The default async path is aiohttp-backed, whose pool accounts for its own leases.""" - server = await asyncio.start_server(_trickling_upstream, "127.0.0.1", 0) - port = server.sockets[0].getsockname()[1] - clock = FakeClock() - closer = make_closer(clock) - handler = AsyncHTTPHandler() - held_client = handler.client - - closer.mark_owned(handler) - closer.schedule(handler) - - request = asyncio.create_task(handler.get(f"http://127.0.0.1:{port}/")) - await asyncio.sleep(0.25) - clock.advance(3600.0) - closer.reap() - await asyncio.sleep(0.05) - - assert held_client.is_closed is False, "closed a handler that was serving a request" - assert (await request).status_code == 200 - - clock.advance(3600.0) - closer.reap() - await asyncio.sleep(0.05) - - assert held_client.is_closed is True - assert handler.client.is_closed is False, "a held handler must self-heal after its evicted client is closed" - server.close() - - def test_the_pending_queue_cannot_grow_past_its_bound(): """A caller that churns the client cache must not be able to grow this queue.""" clock = FakeClock() diff --git a/tests/unit/caching/test_valkey_semantic_cache.py b/tests/unit/caching/test_valkey_semantic_cache.py index 749658784ac..c4c8827cd49 100644 --- a/tests/unit/caching/test_valkey_semantic_cache.py +++ b/tests/unit/caching/test_valkey_semantic_cache.py @@ -1,9 +1,6 @@ import hashlib import os import struct -import subprocess -import sys -import textwrap from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -508,29 +505,3 @@ async def test_index_info_uses_valkey_ft_info(): assert result == {"num_docs": 3} async_client.ft.assert_called_once_with("test_index") - - -def test_importing_caching_does_not_require_redis(): - # redis is an optional dependency (extra_proxy), so the base SDK can be - # installed without it. Selecting valkey-semantic needs redis, but merely - # importing litellm.caching.caching must not, or `import litellm` breaks for - # every base-SDK user. This runs in a subprocess with redis blocked so the - # check is not polluted by redis already being imported in this session. - code = textwrap.dedent(""" - import sys - for name in ("redis", "redis.asyncio", "redis.commands", - "redis.commands.search"): - sys.modules[name] = None - import litellm.caching.caching # must not import redis at module top - from litellm.types.caching import LiteLLMCacheType - assert LiteLLMCacheType.VALKEY_SEMANTIC == "valkey-semantic" - print("ok") - """) - result = subprocess.run( - [sys.executable, "-c", code], - capture_output=True, - text=True, - env={**os.environ, "PYTHONPATH": _REPO_ROOT}, - ) - assert result.returncode == 0, result.stderr - assert "ok" in result.stdout diff --git a/tests/unit/enterprise/conftest.py b/tests/unit/enterprise/conftest.py index 4c95f967bc4..e815ba13edb 100644 --- a/tests/unit/enterprise/conftest.py +++ b/tests/unit/enterprise/conftest.py @@ -1,12 +1,17 @@ # conftest.py import asyncio -import importlib -import os +from typing import Final import pytest import litellm +import litellm.proxy.proxy_server +from tests._module_state import collect_flushable_caches, flush_caches, restore_mutable_state, snapshot_mutable_state + +_LITELLM_STATE: Final = snapshot_mutable_state(litellm) +_PROXY_SERVER_STATE: Final = snapshot_mutable_state(litellm.proxy.proxy_server) +_FLUSHABLE_CACHES: Final = collect_flushable_caches((litellm, litellm.proxy.proxy_server)) @pytest.fixture(scope="session") @@ -19,38 +24,17 @@ def event_loop(): loop.close() - - @pytest.fixture(scope="function", autouse=True) def setup_and_teardown(): - """ - This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. - """ - curr_dir = os.getcwd() # Get the current working directory - - from litellm import Router - - importlib.reload(litellm) - - try: - if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"): - importlib.reload(litellm.proxy.proxy_server) - except Exception as e: - print(f"Error reloading litellm.proxy.proxy_server: {e}") - - litellm.in_memory_llm_clients_cache.flush_cache() - - import asyncio + restore_mutable_state(litellm, _LITELLM_STATE) + restore_mutable_state(litellm.proxy.proxy_server, _PROXY_SERVER_STATE) + flush_caches(_FLUSHABLE_CACHES) loop = asyncio.get_event_loop_policy().new_event_loop() asyncio.set_event_loop(loop) - print(litellm) - # from litellm import Router, completion, aembedding, acompletion, embedding yield - - # Teardown code (executes after the yield point) - loop.close() # Close the loop created earlier - asyncio.set_event_loop(None) # Remove the reference to the loop + loop.close() + asyncio.set_event_loop(None) def pytest_collection_modifyitems(config, items): diff --git a/tests/unit/harness/sandbox/test_docker.py b/tests/unit/harness/sandbox/test_docker.py index 10cb8a092f3..40e7ce88f45 100644 --- a/tests/unit/harness/sandbox/test_docker.py +++ b/tests/unit/harness/sandbox/test_docker.py @@ -1,7 +1,5 @@ import asyncio -import hashlib import shutil -import subprocess from typing import Optional import pytest @@ -10,8 +8,6 @@ from litellm import sandbox from litellm.harness.errors import SandboxError from litellm.harness.sandbox import DockerSandbox, Sandbox from litellm.harness.sandbox.docker import parse_sha256sum - -DOCKER_IMAGE = "alpine:3.20" CID = "cid123" @@ -198,44 +194,3 @@ async def test_missing_docker_binary(monkeypatch): def test_parse_sha256sum_ignores_junk(): assert parse_sha256sum("garbage\n\n") == {} - - -def _docker_usable() -> bool: - if shutil.which("docker") is None: - return False - try: - return ( - subprocess.run( - ["docker", "info"], capture_output=True, timeout=20 - ).returncode - == 0 - ) - except (OSError, subprocess.SubprocessError): - return False - - -@pytest.mark.skipif(not _docker_usable(), reason="docker daemon not available") -async def test_real_docker_roundtrip(): - box = sandbox.docker(DOCKER_IMAGE, workdir="/workspace") - try: - result = await box.run(["echo", "hello"], timeout=120) - assert result.stdout.strip() == "hello" - assert result.exit_code == 0 - - await box.write("seed.txt", b"seed") - await box.write("sub/out.txt", b"from host") - assert await box.read("sub/out.txt") == b"from host" - await box.write("node_modules/skip.js", b"x") - - assert await box.which("sh") is not None - assert await box.which("definitely-not-a-binary-xyz") is None - tmp = await box.tempdir() - assert tmp.startswith("/") - - snap = await box.snapshot() - assert snap == { - "seed.txt": hashlib.sha256(b"seed").hexdigest(), - "sub/out.txt": hashlib.sha256(b"from host").hexdigest(), - } - finally: - await box.close() diff --git a/tests/unit/integrations/test_anthropic_cache_control_hook.py b/tests/unit/integrations/test_anthropic_cache_control_hook.py index 375f3d47cd7..a8da1ad512d 100644 --- a/tests/unit/integrations/test_anthropic_cache_control_hook.py +++ b/tests/unit/integrations/test_anthropic_cache_control_hook.py @@ -4,12 +4,10 @@ import json import logging import os import pathlib -import subprocess import sys -import textwrap from collections.abc import Mapping from types import MappingProxyType -from typing import Final, List, Optional, Tuple, cast +from typing import Final, List, cast from unittest.mock import MagicMock, patch import pytest @@ -3501,54 +3499,6 @@ class TestConfiguredInjectionPointsSurviveClientMarks: assert sys2 == sys1 -class TestAnthropicPromptCachingEnvVars: - """Both settings are read from the environment at import, so an admin can enable - auto-caching without a config file. Each case re-imports litellm in a subprocess - so the env is read fresh without contaminating this process's module graph. - """ - - @staticmethod - def _import_litellm_with_env(env_override: dict) -> Tuple[bool, Optional[str]]: - env = os.environ.copy() - env.pop("LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING", None) - env.pop("LITELLM_ANTHROPIC_PROMPT_CACHING_TTL", None) - env.update(env_override) - script = textwrap.dedent( - """ - import json, litellm - print(json.dumps([litellm.enable_anthropic_prompt_caching, litellm.anthropic_prompt_caching_ttl])) - """ - ) - result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, env=env, timeout=300) - assert result.returncode == 0, result.stderr - enabled, ttl = json.loads(result.stdout.strip().splitlines()[-1]) - return enabled, ttl - - def test_unset_env_leaves_auto_caching_off(self): - assert self._import_litellm_with_env({}) == (False, None) - - @pytest.mark.parametrize("value", ["true", "True", "TRUE"]) - def test_env_enables_auto_caching_case_insensitively(self, value): - enabled, _ = self._import_litellm_with_env({"LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING": value}) - assert enabled is True - - @pytest.mark.parametrize("value", ["false", "0", "yes", ""]) - def test_env_only_enables_on_true(self, value): - enabled, _ = self._import_litellm_with_env({"LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING": value}) - assert enabled is False - - @pytest.mark.parametrize("value", ["5m", "1h"]) - def test_ttl_env_is_applied(self, value): - _, ttl = self._import_litellm_with_env({"LITELLM_ANTHROPIC_PROMPT_CACHING_TTL": value}) - assert ttl == value - - @pytest.mark.parametrize("value", ["10m", "1H", "3600", "ephemeral"]) - def test_unsupported_ttl_env_falls_back_to_provider_default(self, value): - """An unparseable TTL must fall back to Anthropic's 5m default, never reach the provider verbatim.""" - _, ttl = self._import_litellm_with_env({"LITELLM_ANTHROPIC_PROMPT_CACHING_TTL": value}) - assert ttl is None - - def _contains_key(value, key) -> bool: if isinstance(value, dict): return key in value or any(_contains_key(v, key) for v in value.values()) diff --git a/tests/unit/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py b/tests/unit/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py index 90bf7ba2920..37606dd4a0b 100644 --- a/tests/unit/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py +++ b/tests/unit/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py @@ -215,36 +215,6 @@ async def test_hook_searches_through_the_injected_router_with_the_request_metada assert messages[0]["content"] == "Context:\n\ncontext from vs-router\n\n" -@pytest.mark.asyncio -@pytest.mark.parametrize("sdk_verbose", [False, True]) -async def test_hook_falls_back_to_the_sdk_when_the_runtime_has_no_router( - registry_with: RegisterStores, - warnings: list[logging.LogRecord], - monkeypatch: pytest.MonkeyPatch, - sdk_verbose: bool, -) -> None: - monkeypatch.setattr(litellm, "set_verbose", sdk_verbose) - registry_with("vs-sdk", custom_llm_provider="lit6752-not-a-provider") - - _, messages, _ = await _run_hook( - VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=None)), - ["vs-sdk"], - FakeLoggingObj({"user_api_key_team_id": "team-a"}), - ) - - assert messages == [{"role": "user", "content": "what is litellm?"}] - search_warnings: Final = tuple( - record.getMessage() - for record in warnings - if record.getMessage().startswith("Vector store search failed for vector_store_id=") - ) - assert len(search_warnings) == 1 - assert search_warnings[0].startswith( - "Vector store search failed for vector_store_id=vs-sdk, continuing without its context: " - ) - assert "is not a valid LlmProviders" in search_warnings[0] - - @pytest.mark.asyncio async def test_every_healthy_vector_store_contributes_its_own_context(registry_with: RegisterStores) -> None: """Regression (LIT-6752): each store appended its context to the original messages, so only the last one survived.""" diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 6ad72521d09..616e79ffac6 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -2,16 +2,13 @@ import asyncio import datetime import itertools import json -import os import queue -import subprocess import sys import threading import time from collections.abc import Mapping, Sequence from dataclasses import dataclass from itertools import chain -from pathlib import Path from types import CodeType, FrameType, MappingProxyType, TracebackType from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -2071,66 +2068,6 @@ async def test_stream_logprobs_without_top_logprobs_reach_caller_as_litellm_logp ] -FINISH_ONLY_LOGPROBS_STREAM_SCRIPT: Final = """ -import json -import sys - -sys.path.insert(0, sys.argv[1]) - -import httpx -from openai import OpenAI - -import litellm - -upstream_choices = ( - {"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None}, - { - "index": 0, - "delta": {}, - "logprobs": {"content": [{"token": "<|return|>", "logprob": -0.5, "top_logprobs": []}]}, - "finish_reason": "stop", - }, -) -body = "".join( - "data: " - + json.dumps({"id": "c", "object": "chat.completion.chunk", "created": 1, "model": "m", "choices": [choice]}) - + "\\n\\n" - for choice in upstream_choices -) + "data: [DONE]\\n\\n" -client = OpenAI( - api_key="unused", - base_url="https://provider.example/v1", - http_client=httpx.Client( - transport=httpx.MockTransport( - lambda request: httpx.Response(200, content=body, headers={"content-type": "text/event-stream"}) - ) - ), -) -stream = litellm.completion( - model="openai/m", messages=[{"role": "user", "content": "hi"}], stream=True, logprobs=True, client=client -) -print(json.dumps([json.loads(chunk.model_dump_json(exclude_none=True)) for chunk in stream])) -""" - - -def test_stream_with_logprobs_only_on_finish_chunk_serializes_in_a_fresh_process() -> None: - result: Final = subprocess.run( - [sys.executable, "-I", "-c", FINISH_ONLY_LOGPROBS_STREAM_SCRIPT, str(Path(litellm.__file__).parent.parent)], - capture_output=True, - text=True, - timeout=120, - env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, - ) - - assert result.returncode == 0, result.stdout + result.stderr - wire_choices: Final = tuple( - itertools.chain.from_iterable(chunk["choices"] for chunk in json.loads(result.stdout.strip().splitlines()[-1])) - ) - finish_choices: Final = [choice for choice in wire_choices if choice.get("finish_reason")] - assert [choice["finish_reason"] for choice in finish_choices] == ["stop"] - assert [entry["token"] for entry in finish_choices[0]["logprobs"]["content"]] == ["<|return|>"] - - def _hosted_vllm_stream_wrapper(logging_obj: Logging, error_payload: dict) -> CustomStreamWrapper: """A CustomStreamWrapper over the real OpenAI-compatible line iterator, fed an HTTP 200 SSE body that carries an in-body error payload the way @@ -2398,7 +2335,6 @@ async def test_mid_stream_error_leaves_out_an_api_base_it_cannot_resolve(): assert "API Base:" not in debug_info - @pytest.mark.asyncio @pytest.mark.parametrize( "routed_api_base", @@ -7169,7 +7105,6 @@ def test_chunk_creator_skips_full_dump_of_model_response_stream_chunks( assert getattr(emitted_chunks[1], "custom_chunk_attribute") == "custom value" - def test_openai_chunk_creator_keeps_multiple_choice_logprobs_and_extras( logging_obj: Logging, ) -> None: @@ -9765,7 +9700,6 @@ async def test_logged_response_keeps_finish_reason_from_last_content_chunk( assert logged.choices[0].message.content == "Hello, this reply is cut off" - @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_no_synthetic_finish_reason_logged_when_provider_sent_none(sync_mode: bool, logging_obj: Logging): diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 4696203c914..b65f2fc383d 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -6,15 +6,12 @@ import copy import importlib import io import json -import os -import subprocess import sys import threading import time from collections.abc import Callable, Mapping from concurrent.futures import Future, ThreadPoolExecutor, wait from itertools import accumulate, chain -from pathlib import Path from queue import SimpleQueue from typing import Final, cast from unittest.mock import MagicMock, sentinel @@ -23,7 +20,7 @@ import anyio.to_thread import pytest import tiktoken from PIL import Image -from tokenizers import Regex, Tokenizer, models, pre_tokenizers +from tokenizers import Tokenizer from unittest.mock import AsyncMock, patch @@ -922,48 +919,6 @@ class NeedsToleranceUpdateError(Exception): # test_tokenizers() -def test_encoding_and_decoding(tmp_path: Path): - sample_text = "Hellö World, this is my input string!" - - # openai encoding + decoding - openai_tokens = encode(model="gpt-3.5-turbo", text=sample_text) - openai_text = decode(model="gpt-3.5-turbo", tokens=openai_tokens) - - assert openai_text == sample_text - - # claude encoding + decoding - claude_tokens = encode(model="claude-3-5-haiku-20241022", text=sample_text) - - claude_text = decode(model="claude-3-5-haiku-20241022", tokens=claude_tokens) - - assert claude_text == sample_text - - # cohere encoding + decoding - cohere_tokens = encode(model="command-nightly", text=sample_text) - cohere_text = decode(model="command-nightly", tokens=cohere_tokens) - - assert cohere_text == sample_text - - # llama2 encoding + decoding - words = sample_text.split() - result = _run_in_memory_hub( - HUB_ROUND_TRIP_SCRIPT, - { - "hf-internal-testing/llama-tokenizer": _word_level_tokenizer_json( - pre_tokenizers.WhitespaceSplit(), - vocab={"[UNK]": 0, **{word: i + 1 for i, word in enumerate(words)}}, - ) - }, - sample_text, - tmp_path, - ) - - assert result["decoded"] == sample_text - assert result["requested"] == ["hf-internal-testing/llama-tokenizer"] - assert len(result["tokens"]) == len(words) - assert len(result["tokens"]) != len(encode(model="gpt-3.5-turbo", text=sample_text)) - - # test_encoding_and_decoding() @@ -1162,7 +1117,6 @@ from litellm.utils import ( _HuggingFaceTokenizerLoadOutcomes, _return_huggingface_tokenizer, _select_tokenizer_helper, - claude_json_str, encoding, ) @@ -3564,129 +3518,6 @@ def test_image_estimate_bounds_header_decoding_for_large_inline_payloads() -> No assert len(decode_header.call_args.args[0]) <= 65_536 -HUB_SETUP_SCRIPT: Final = """ -import json -import sys -sys.path.insert(0, sys.argv[1]) -import httpx -import huggingface_hub -import litellm -served = json.loads(sys.argv[2]) -text = sys.argv[3] -requested = [] -def handle(request): - repo = request.url.path.lstrip("/").split("/resolve/")[0] - if repo not in served or not request.url.path.endswith("/tokenizer.json"): - return httpx.Response(404) - requested.append(repo) - payload = served[repo].encode() - headers = {"content-length": str(len(payload)), "etag": '"fixture"', "x-repo-commit": "a" * 40} - return httpx.Response(200, headers=headers, content=payload if request.method == "GET" else b"") -huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(handle))) -""" - -HUB_TOKENIZER_SCRIPT: Final = HUB_SETUP_SCRIPT + """ -litellm.cohere_models = {"command-r-v1"} -litellm.anthropic_models = {"claude-2"} -custom = litellm.create_pretrained_tokenizer("Xenova/llama-3-tokenizer") -print(json.dumps({ - "llama2": litellm.token_counter(model="meta-llama/Llama-2-7b-chat", text=text), - "llama3": litellm.token_counter(model="meta-llama/llama-3-70b-instruct", text=text), - "cohere": litellm.token_counter(model="command-r-v1", text=text), - "anthropic": litellm.token_counter(model="claude-2", text=text), - "custom": litellm.token_counter(custom_tokenizer=custom, text=text), - "requested": sorted(set(requested)), -})) -""" - -HUB_ROUND_TRIP_SCRIPT: Final = HUB_SETUP_SCRIPT + """ -tokens = litellm.encode(model="meta-llama/Llama-2-7b-chat", text=text) -print(json.dumps({"tokens": tokens, "decoded": litellm.decode(model="meta-llama/Llama-2-7b-chat", tokens=tokens), "requested": sorted(set(requested))})) -""" - -HUB_RETRY_DURING_EXIT_SCRIPT: Final = HUB_SETUP_SCRIPT + """ -import threading -import time -from litellm.utils import _HuggingFaceTokenizerLoadOutcomes, _return_huggingface_tokenizer -retry_reached_hub = threading.Event() -def answer_once_the_process_is_exiting(request): - retry_reached_hub.set() - while threading.main_thread().is_alive(): - time.sleep(0.01) - time.sleep(0.5) - return handle(request) -now = [0.0] -outcomes = _HuggingFaceTokenizerLoadOutcomes(clock=lambda: now[0]) -model = "meta-llama/llama-3-70b-instruct" -_return_huggingface_tokenizer(model, outcomes) -huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(answer_once_the_process_is_exiting))) -now[0] = 3600.0 -_return_huggingface_tokenizer(model, outcomes) -print(json.dumps({"retry_reached_hub": retry_reached_hub.wait(timeout=60)})) -""" - - -def _word_level_tokenizer_json( - pre_tokenizer: pre_tokenizers.PreTokenizer, vocab: Mapping[str, int] | None = None -) -> str: - tokenizer: Final = Tokenizer( - models.WordLevel(vocab=dict(vocab) if vocab is not None else {"[UNK]": 0}, unk_token="[UNK]") - ) - tokenizer.pre_tokenizer = pre_tokenizer - return tokenizer.to_str() - - -def _run_in_memory_hub(script: str, served: dict[str, str], text: str, tmp_path: Path) -> dict: - result: Final = subprocess.run( - [ - sys.executable, - "-I", - "-c", - script, - str(Path(litellm.__file__).parent.parent), - json.dumps(served), - text, - ], - capture_output=True, - text=True, - timeout=60, - env={ - **os.environ, - "HF_HOME": str(tmp_path / "home"), - "HF_HUB_CACHE": str(tmp_path / "cache"), - "HF_ENDPOINT": "http://127.0.0.1:9", - "HF_HUB_OFFLINE": "0", - "LITELLM_LOCAL_MODEL_COST_MAP": "True", - }, - ) - - assert result.returncode == 0, result.stdout + result.stderr - return json.loads(result.stdout.strip().splitlines()[-1]) - - -def test_token_counter_uses_the_tokenizer_of_each_model_family_and_of_a_custom_tokenizer(tmp_path: Path) -> None: - sample: Final = "Tokenizers disagree: anthropic, tiktoken; llama-2 & llama-3!" - served: Final = { - "hf-internal-testing/llama-tokenizer": _word_level_tokenizer_json(pre_tokenizers.WhitespaceSplit()), - "Xenova/llama-3-tokenizer": _word_level_tokenizer_json(pre_tokenizers.Split(Regex("."), "isolated")), - "Xenova/c4ai-command-r-v01-tokenizer": _word_level_tokenizer_json(pre_tokenizers.Whitespace()), - } - expected: Final = {repo: len(Tokenizer.from_str(payload).encode(sample).ids) for repo, payload in served.items()} - anthropic_count: Final = len(Tokenizer.from_str(claude_json_str).encode(sample).ids) - tiktoken_count: Final = litellm.token_counter(model="gpt-3.5-turbo", text=sample) - assert len({*expected.values(), anthropic_count, tiktoken_count}) == len(expected) + 2 - - counts: Final = _run_in_memory_hub(HUB_TOKENIZER_SCRIPT, served, sample, tmp_path) - assert counts == { - "llama2": expected["hf-internal-testing/llama-tokenizer"], - "llama3": expected["Xenova/llama-3-tokenizer"], - "cohere": expected["Xenova/c4ai-command-r-v01-tokenizer"], - "anthropic": anthropic_count, - "custom": expected["Xenova/llama-3-tokenizer"], - "requested": sorted(served), - } - - def test_custom_openai_tokenizer_without_a_model_counts_with_the_default_encoding() -> None: sample: Final = "A custom tokenizer passed with no model still counts." @@ -3695,10 +3526,6 @@ def test_custom_openai_tokenizer_without_a_model_counts_with_the_default_encodin ) -def test_process_exits_cleanly_while_a_background_tokenizer_retry_is_downloading(tmp_path: Path) -> None: - assert _run_in_memory_hub(HUB_RETRY_DURING_EXIT_SCRIPT, {}, "", tmp_path) == {"retry_reached_hub": True} - - def test_empty_custom_tokenizer_uses_model_tokenizer() -> None: text_value: Final = "A tokenizer fallback should preserve the model encoding." custom_count: Final = _get_exact_count_function("gpt-3.5-turbo", {})(text_value) diff --git a/tests/unit/litellm_core_utils/test_tokenizer.py b/tests/unit/litellm_core_utils/test_tokenizer.py index 343b67699d3..fd605e5f2d1 100644 --- a/tests/unit/litellm_core_utils/test_tokenizer.py +++ b/tests/unit/litellm_core_utils/test_tokenizer.py @@ -1,7 +1,5 @@ import copy -import os import pickle -import subprocess import sys from pathlib import Path from typing import Final, Literal @@ -176,115 +174,6 @@ def test_public_tokenizer_objects_survive_pickle_and_deepcopy(tmp_path: Path) -> assert pickle.loads(pickle.dumps(litellm.encoding)).encode("hello") == litellm.encoding.encode("hello") -@pytest.mark.parametrize("offline", ("0", "1")) -def test_hub_loader_preserves_environment_auth_cache_and_offline(tmp_path: Path, offline: str) -> None: - script: Final = """ -import json -import sys -from pathlib import Path -sys.path.insert(0, sys.argv[1]) -import httpx -import huggingface_hub -from huggingface_hub.errors import LocalEntryNotFoundError -import litellm -payload = sys.argv[2].encode() -offline = sys.argv[3] == "1" -observed = [] -def handle(request): - assert not offline, "offline loading issued a request" - if request.url.path.endswith("/tokenizer.json"): - observed.append(request.headers.get("authorization")) - if request.headers.get("authorization") != "Bearer audit-fixture-token": - return httpx.Response(401) - return httpx.Response(200, headers={"content-length": str(len(payload)), "etag": '"fixture"', "x-repo-commit": "a" * 40}, content=payload if request.method == "GET" else b"") -if not offline: - huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(handle))) -try: - tokenizer = litellm.create_pretrained_tokenizer("test-fixture/tokenizer")["tokenizer"] -except LocalEntryNotFoundError: - assert offline - assert observed == [] -else: - assert not offline - assert "Bearer audit-fixture-token" in observed - assert tokenizer.decode(tokenizer.encode("Hello World").ids) == "Hello World" - assert tuple(Path(sys.argv[4]).rglob("tokenizer.json")) -print("compatible") -""" - result: Final = subprocess.run( - [ - sys.executable, - "-I", - "-c", - script, - str(Path(litellm.__file__).parent.parent), - TOKENIZER_JSON, - offline, - str(tmp_path / "cache"), - ], - capture_output=True, - text=True, - timeout=30, - env={ - **os.environ, - "HF_HOME": str(tmp_path / "home"), - "HF_HUB_CACHE": str(tmp_path / "cache"), - "HF_ENDPOINT": "http://127.0.0.1:9", - "HF_TOKEN": "audit-fixture-token", - "HF_HUB_OFFLINE": offline, - "HF_HUB_DISABLE_IMPLICIT_TOKEN": "0", - "LITELLM_LOCAL_MODEL_COST_MAP": "True", - }, - ) - assert result.returncode == 0, result.stdout + result.stderr - assert result.stdout.strip() == "compatible" - - -@pytest.mark.parametrize("rust", (None, "0", "1")) -def test_tokenization_stays_offline(tmp_path: Path, rust: str | None) -> None: - script: Final = """ -import sys -sys.path.insert(0, sys.argv[1]) -def reject_network(event, args): - if event == "socket.connect": - raise AssertionError("tokenizer attempted a network connection") -sys.addaudithook(reject_network) -import litellm -from litellm.rust_bridge.tokenizer import get_encoding -import tiktoken -from tokenizers import Tokenizer -assert isinstance(litellm.encoding, tiktoken.Encoding) -for name in ("cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit"): - encoding = get_encoding(name) - text = "offline café 漢字 🙂" + " " * 64 - assert encoding.decode(encoding.encode(text)) == text -ids = litellm.encode(text="hello world") -assert litellm.decode(tokens=ids) == "hello world" -assert litellm.token_counter(model=None, text="hello world") == len(ids) -custom = litellm.create_tokenizer(sys.argv[2]) -assert isinstance(custom["tokenizer"], Tokenizer) -custom["tokenizer"].enable_padding(pad_id=0, pad_token="[UNK]") -assert litellm.decode(tokens=litellm.encode(text="Hello World", custom_tokenizer=custom), custom_tokenizer=custom) == "Hello World" -print("compatible") -""" - result: Final = subprocess.run( - [sys.executable, "-I", "-c", script, str(Path(litellm.__file__).parent.parent), TOKENIZER_JSON], - capture_output=True, - text=True, - timeout=30, - cwd=tmp_path, - env={ - **{key: value for key, value in os.environ.items() if key != "LITELLM_RUST"}, - **({"LITELLM_RUST": rust} if rust is not None else {}), - "LITELLM_LOCAL_MODEL_COST_MAP": "True", - "TIKTOKEN_CACHE_DIR": str(tmp_path / "unused-tokenizer-cache"), - }, - ) - assert result.returncode == 0, result.stdout + result.stderr - assert result.stdout.strip() == "compatible" - assert not (tmp_path / "unused-tokenizer-cache").exists() - - @pytest.mark.parametrize("is_pretokenized", (False, True)) def test_huggingface_batch_sequence_containers_match_python(is_pretokenized: bool) -> None: reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON) diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index d105a7958b4..68e2260814d 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -14,7 +14,6 @@ import itertools import json import logging import os -import subprocess import sys import threading from collections.abc import Mapping @@ -4910,40 +4909,3 @@ def test_a_register_model_opt_out_for_one_claude_model_leaves_every_other_routed litellm.register_model({f"anthropic/{opted_out}": {"supports_prompt_caching": False}}) assert supports_anthropic_cache_control(untouched, "anthropic", {"id": "plain-id"}) is True - - -_MASKED_HOSTED_WEB_CALL_PROBE: Final = """ -import json - -from litellm.llms.anthropic.common_utils import write_request_hosted_web_call_arguments - -message = { - "role": "assistant", - "provider_specific_fields": { - "web_search_calls": [ - { - "id": "ws_srvtoolu_1", - "type": "web_search_call", - "status": "completed", - "action": {"type": "search", "query": "deploy token 12345", "queries": ["deploy token 12345"]}, - } - ] - }, -} -write_request_hosted_web_call_arguments(message, {"srvtoolu_1": json.dumps({"query": "deploy token [REDACTED]"})}) -print(json.dumps(message["provider_specific_fields"]["web_search_calls"])) -""" - - -def test_a_masked_replayed_hosted_web_call_serializes_before_its_openai_model_was_ever_built() -> None: - result: Final = subprocess.run( - (sys.executable, "-I", "-c", _MASKED_HOSTED_WEB_CALL_PROBE), - capture_output=True, - text=True, - check=True, - timeout=120, - ) - hosted: Final = json.loads(result.stdout.strip().splitlines()[-1]) - assert [call["action"] for call in hosted] == [ - {"type": "search", "query": "deploy token [REDACTED]", "queries": ["deploy token [REDACTED]"]} - ] diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index cafd00cb592..618aa15bd91 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -437,28 +437,6 @@ def test_initialize_honors_explicit_max_retries(setup_mocks, configured, expecte assert result["max_retries"] == expected -def test_default_max_retries_env_var_reaches_azure_sdk_client(): - import subprocess - import sys - - code = ( - "from litellm.llms.azure.common_utils import BaseAzureLLM\n" - "client = BaseAzureLLM().get_azure_openai_client(" - "api_key='test-api-key', api_base='https://test.openai.azure.com', api_version='2024-02-01'," - " client=None, _is_async=True, litellm_params={}, model='gpt-4')\n" - "print(client.max_retries)" - ) - completed = subprocess.run( - [sys.executable, "-c", code], - env={**os.environ, "DEFAULT_MAX_RETRIES": "0"}, - capture_output=True, - text=True, - check=True, - ) - - assert completed.stdout.strip() == "0" - - @pytest.mark.parametrize( "call_type", [ @@ -3268,4 +3246,3 @@ def test_non_chat_calls_on_a_foundry_project_v1_base_reach_azure( respx_mock.post().mock(return_value=httpx.Response(200, json=upstream_body)) assert call() == "ok" - diff --git a/tests/unit/llms/base_llm/auth/test_jwt_signing.py b/tests/unit/llms/base_llm/auth/test_jwt_signing.py index f1fbc698c82..62994fc0735 100644 --- a/tests/unit/llms/base_llm/auth/test_jwt_signing.py +++ b/tests/unit/llms/base_llm/auth/test_jwt_signing.py @@ -1,9 +1,7 @@ import base64 import hashlib import json -import subprocess import sys -import textwrap import time from typing import Final @@ -13,7 +11,6 @@ from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import ec, rsa from litellm.llms.base_llm.auth.jwt_signing import ( - MISSING_SIGNING_DEPENDENCIES_MESSAGE, build_jwk, build_jwks, jwks_document_json, @@ -183,33 +180,7 @@ class TestBaseSdkImport: and ``litellm/__init__`` reaches this module through the Anthropic provider, so a module-level import of either would break ``import litellm`` for every base SDK user.""" - def test_module_imports_with_pyjwt_and_cryptography_absent(self): - script: Final = textwrap.dedent( - """ - import sys - - class Blocker: - def find_spec(self, name, path=None, target=None): - if name.split(".")[0] in {"jwt", "cryptography"}: - raise ModuleNotFoundError(f"No module named {name!r}") - - sys.meta_path.insert(0, Blocker()) - import litellm - from litellm.llms.base_llm.auth.jwt_signing import jwks_document_json - try: - jwks_document_json("not a key") - except ImportError as e: - print(e) - """ - ) - result: Final = subprocess.run( - [sys.executable, "-I", "-c", script], capture_output=True, text=True, check=False - ) - assert result.returncode == 0, result.stderr[-2000:] - assert result.stdout.strip() == MISSING_SIGNING_DEPENDENCIES_MESSAGE - def test_signing_reports_the_missing_extra(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setitem(sys.modules, "jwt", None) with pytest.raises(ImportError, match="litellm\\[proxy\\]"): sign_es256_jwt(pem_of(fixed_private_key()), {"sub": "x"}) - diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index 9698c4ac62f..2f7a15e1acf 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -1,19 +1,16 @@ import asyncio import gc -import importlib import json import logging import os import pathlib import ssl import threading -import time import weakref from collections.abc import AsyncIterator, Callable, Mapping from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager from contextvars import Context -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Final from unittest.mock import MagicMock, patch @@ -29,8 +26,6 @@ from yarl import URL import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.exceptions import Timeout as LitellmTimeout -from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.litellm_core_utils.model_cost_table import publish from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine, default_token_exchange_engine from litellm.llms.base_llm.auth.types import MintedToken, TokenExchangeSpec from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport @@ -46,7 +41,6 @@ from litellm.llms.custom_httpx.http_handler import ( realtime_ssl_for_url, ) from litellm.types.llms.custom_http import VerifyTypes -from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @pytest.mark.asyncio @@ -879,116 +873,6 @@ class TestDefaultCachedClientTimeoutHonorsRequestTimeout: assert client.timeout.read == 300.0 -async def _read_http_request(reader: asyncio.StreamReader) -> None: - raw = b"" - while b"\r\n\r\n" not in raw: - chunk = await reader.read(1024) - if not chunk: - return - raw += chunk - head, _, body = raw.partition(b"\r\n\r\n") - content_length = next( - (int(line.split(b":", 1)[1]) for line in head.split(b"\r\n") if line.lower().startswith(b"content-length")), - 0, - ) - while len(body) < content_length: - body += await reader.read(content_length - len(body)) - - -@pytest.mark.asyncio -async def test_init_held_async_handler_survives_external_client_close(): - handler = AsyncHTTPHandler(timeout=42.5) - held_client = handler.client - await held_client.aclose() - assert held_client.is_closed - - async def respond(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: - await _read_http_request(reader) - writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok") - await writer.drain() - writer.close() - - server = await asyncio.start_server(respond, "127.0.0.1", 0) - port = server.sockets[0].getsockname()[1] - try: - response = await handler.post(f"http://127.0.0.1:{port}/v1/compress", json={"messages": []}) - finally: - server.close() - await server.wait_closed() - - assert response.status_code == 200 - assert handler.client is not held_client - assert handler.client.timeout == httpx.Timeout(42.5) - await handler.close() - - -@pytest.mark.asyncio -async def test_init_held_async_handler_survives_evicted_client_close(): - from litellm.caching.evicted_client_closer import EvictedClientCloser - from litellm.caching.llm_caching_handler import LLMClientCache - - cache = LLMClientCache(evicted_client_closer=EvictedClientCloser(grace_seconds=0)) - handler = AsyncHTTPHandler(timeout=42.5) - held_client = handler.client - cache.set_cache("init-held-handler", handler, litellm_owned_client=True, ttl=0) - await asyncio.sleep(0.02) - assert cache.get_cache("init-held-handler") is None - await asyncio.sleep(0.05) - assert held_client.is_closed - - async def respond(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: - await _read_http_request(reader) - writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok") - await writer.drain() - writer.close() - - server = await asyncio.start_server(respond, "127.0.0.1", 0) - port = server.sockets[0].getsockname()[1] - try: - response = await handler.post(f"http://127.0.0.1:{port}/v1/compress", json={"messages": []}) - finally: - server.close() - await server.wait_closed() - - assert response.status_code == 200 - assert handler.client is not held_client - assert handler.client.timeout == httpx.Timeout(42.5) - await handler.close() - - -def test_init_held_sync_handler_recreates_closed_client(): - from http.server import BaseHTTPRequestHandler, HTTPServer - - class OkRequestHandler(BaseHTTPRequestHandler): - def do_GET(self): - self.send_response(200) - self.send_header("Content-Length", "2") - self.end_headers() - self.wfile.write(b"ok") - - def log_message(self, format, *args): - pass - - handler = HTTPHandler(timeout=7) - held_client = handler.client - held_client.close() - - server = HTTPServer(("127.0.0.1", 0), OkRequestHandler) - thread = threading.Thread(target=server.serve_forever, daemon=True) - thread.start() - try: - response = handler.get(f"http://127.0.0.1:{server.server_port}/") - finally: - server.shutdown() - server.server_close() - thread.join(timeout=5) - - assert response.status_code == 200 - assert handler.client is not held_client - assert handler.client.timeout == httpx.Timeout(7) - handler.close() - - def test_caller_supplied_sync_client_is_not_replaced_when_closed(): supplied = httpx.Client() handler = HTTPHandler(client=supplied) @@ -2245,66 +2129,6 @@ async def test_http2_disabled_by_default(monkeypatch: pytest.MonkeyPatch): assert AsyncHTTPHandler.should_use_aiohttp_transport() is True -@pytest.fixture() -def _vcr_outcome_gate(request, vcr): - install_live_call_probe(request, vcr) - yield - record_vcr_outcome(request, vcr) - - -@pytest.fixture(scope="function") -def isolate_litellm_state(): - """ - Per-function isolation fixture. - - Resets litellm globals to their true defaults before each test and - restores them afterward, so tests don't leak side effects. - Works safely under pytest-xdist parallel execution. - """ - original_state = {} - for attr in ( - "callbacks", - "success_callback", - "failure_callback", - "_async_success_callback", - "_async_failure_callback", - ): - if hasattr(litellm, attr): - val = getattr(litellm, attr) - original_state[attr] = val.copy() if val else [] - for attr in ("pre_call_rules", "post_call_rules"): - if hasattr(litellm, attr): - val = getattr(litellm, attr) - original_state[attr] = val.copy() if val else [] - for attr in _SCALAR_DEFAULTS: - if hasattr(litellm, attr): - original_state[attr] = getattr(litellm, attr) - if hasattr(litellm, "in_memory_llm_clients_cache"): - litellm.in_memory_llm_clients_cache.flush_cache() - for attr in ( - "callbacks", - "success_callback", - "failure_callback", - "_async_success_callback", - "_async_failure_callback", - "pre_call_rules", - "post_call_rules", - ): - if hasattr(litellm, attr): - setattr(litellm, attr, []) - for attr, default_val in _SCALAR_DEFAULTS.items(): - if hasattr(litellm, attr): - setattr(litellm, attr, default_val) - yield - asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue()) - if hasattr(litellm, "in_memory_llm_clients_cache"): - litellm.in_memory_llm_clients_cache.flush_cache() - for attr, original_value in original_state.items(): - if hasattr(litellm, attr): - setattr(litellm, attr, original_value) - publish(litellm.model_cost) - - _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -2326,73 +2150,6 @@ _SCALAR_DEFAULTS = { } -@pytest.fixture(scope="module") -def setup_and_teardown(): - """ - Module-scoped setup. Reloads litellm only in single-process mode - (skipped under xdist to avoid cross-worker interference). - """ - import litellm - - worker_id = os.environ.get("PYTEST_XDIST_WORKER", None) - if worker_id is None: - importlib.reload(litellm) - try: - if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"): - import litellm.proxy.proxy_server - - importlib.reload(litellm.proxy.proxy_server) - except Exception: - pass - if hasattr(litellm, "in_memory_llm_clients_cache"): - litellm.in_memory_llm_clients_cache.flush_cache() - yield - - -_SERVER_DELAY_S = 5 - -_PER_REQUEST_TIMEOUT_S = 1.0 - -_CLIENT_DEFAULT_TIMEOUT_S = 60.0 - - -class _SlowHandler(BaseHTTPRequestHandler): - def do_POST(self): - time.sleep(_SERVER_DELAY_S) - try: - self.send_response(200) - self.end_headers() - self.wfile.write(b"{}") - except OSError: - pass - - def log_message(self, *args): - pass - - -@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") -def test_post_delay_exceeds_per_request_timeout_raises(): - server = ThreadingHTTPServer(("127.0.0.1", 0), _SlowHandler) - threading.Thread(target=server.serve_forever, daemon=True).start() - host, port = server.server_address - - handler = get_httpx_client(params={"timeout": _CLIENT_DEFAULT_TIMEOUT_S}) - try: - with pytest.raises(LitellmTimeout): - handler.post( - f"http://{host}:{port}/delay", - headers={"content-type": "application/json"}, - data=json.dumps({"model": "claude", "messages": []}), - timeout=_PER_REQUEST_TIMEOUT_S, - ) - except MaskedHTTPStatusError as e: - pytest.skip(f"httpbin.org unavailable: {e}") - finally: - handler.close() - server.shutdown() - server.server_close() - - def test_realtime_ssl_for_url_sends_no_tls_argument_for_a_plain_ws_endpoint() -> None: assert ( realtime_ssl_for_url("ws://127.0.0.1:8080/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent") diff --git a/tests/unit/llms/opencode/chat/test_transformation.py b/tests/unit/llms/opencode/chat/test_transformation.py index 576f2201130..a0a20daebf8 100644 --- a/tests/unit/llms/opencode/chat/test_transformation.py +++ b/tests/unit/llms/opencode/chat/test_transformation.py @@ -606,26 +606,6 @@ def cost_recorder(monkeypatch: pytest.MonkeyPatch) -> CostRecorder: return recorder -@pytest.mark.asyncio -@pytest.mark.parametrize("upstream", ["chat/completions", "messages", "responses", "gemini"]) -@pytest.mark.parametrize("ingress", ["chat", "responses"]) -async def test_reported_list_cost_is_billed_instead_of_rows( - go_catalog: None, cost_recorder: CostRecorder, upstream: str, ingress: str -) -> None: - model: Final = f"fixture-{upstream.split('/')[0]}" - with respx.mock(using="httpx") as transport: - transport.post().mock( - return_value=httpx.Response(200, json=go_priced_reply(upstream, model, GO_REPORTED_USAGE)) - ) - kwargs: Final = {"model": f"opencode_go/{model}", "api_key": "fake-go-key"} - if ingress == "chat": - response: Final = await litellm.acompletion(messages=[{"role": "user", "content": "hello"}], **kwargs) - else: - response = await litellm.aresponses(input="hello", **kwargs) - assert response._hidden_params["additional_headers"][RESPONSE_COST_HEADER] == pytest.approx(0.00412) - assert await logged_cost(cost_recorder) == pytest.approx(0.00412) - - @pytest.mark.asyncio @pytest.mark.parametrize("usage_extra", [{}, {"cost": 0}, {"opencode": {"billing_source": "go"}}]) async def test_rows_price_the_request_when_no_list_cost_is_reported( diff --git a/tests/unit/llms/vertex_ai/count_tokens/test_context.py b/tests/unit/llms/vertex_ai/count_tokens/test_context.py index 233e5d888fa..26cb15cd27e 100644 --- a/tests/unit/llms/vertex_ai/count_tokens/test_context.py +++ b/tests/unit/llms/vertex_ai/count_tokens/test_context.py @@ -1,11 +1,4 @@ -import asyncio import base64 -import uuid -from collections.abc import Generator -from contextlib import contextmanager -from dataclasses import dataclass, field -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from threading import Event, Thread from typing import Final, Literal from unittest.mock import patch @@ -298,131 +291,6 @@ async def test_openai_messages_keep_their_chat_conversion(custom_llm_provider: L _PNG_BYTES: Final = base64.b64decode( "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" ) -_LOOP_PROOF_TIMEOUT_SECONDS: Final = 5 - - -@dataclass(frozen=True, slots=True) -class _LoopProbe: - request_arrived: Event = field(default_factory=Event) - loop_ran: Event = field(default_factory=Event) - loop_ran_while_download_pending: Event = field(default_factory=Event) - - -def _png_handler_waiting_for_loop(probe: _LoopProbe) -> type[BaseHTTPRequestHandler]: - class _PngHandler(BaseHTTPRequestHandler): - def do_GET(self) -> None: - probe.request_arrived.set() - if probe.loop_ran.wait(_LOOP_PROOF_TIMEOUT_SECONDS): - probe.loop_ran_while_download_pending.set() - self.send_response(200) - self.send_header("Content-Type", "image/png") - self.send_header("Content-Length", str(len(_PNG_BYTES))) - self.end_headers() - self.wfile.write(_PNG_BYTES) - - def log_message(self, format: str, *args: object) -> None: - return None - - return _PngHandler - - -@contextmanager -def _png_server_waiting_for_loop(probe: _LoopProbe) -> Generator[str, None, None]: - with ThreadingHTTPServer(("127.0.0.1", 0), _png_handler_waiting_for_loop(probe)) as server: - thread: Final = Thread(target=server.serve_forever, daemon=True) - thread.start() - try: - yield f"http://127.0.0.1:{server.server_address[1]}" - finally: - server.shutdown() - - -async def _run_loop_once_download_starts(probe: _LoopProbe) -> None: - assert await asyncio.to_thread(probe.request_arrived.wait, _LOOP_PROOF_TIMEOUT_SECONDS) - probe.loop_ran.set() - - -@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) -async def test_remote_image_downloads_without_blocking_the_event_loop( - custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr(litellm, "user_url_validation", False) - probe: Final = _LoopProbe() - with _png_server_waiting_for_loop(probe) as base_url: - image_url: Final = f"{base_url}/{uuid.uuid4().hex}.png" - _, contents = await asyncio.gather( - _run_loop_once_download_starts(probe), - gemini_count_contents( - ( - { - "role": "user", - "content": [ - {"type": "text", "text": "Describe it"}, - {"type": "image", "source": {"type": "url", "url": image_url}}, - ], - }, - ), - "anthropic", - "gemini-2.5-flash", - custom_llm_provider, - {}, - ), - ) - - assert contents == [ - { - "role": "user", - "parts": [ - {"text": "Describe it"}, - {"inline_data": {"data": base64.b64encode(_PNG_BYTES).decode(), "mime_type": "image/png"}}, - ], - } - ] - assert probe.loop_ran_while_download_pending.is_set() - - -class _MissingImageHandler(BaseHTTPRequestHandler): - def do_GET(self) -> None: - self.send_response(404) - self.send_header("Content-Length", "0") - self.end_headers() - - def log_message(self, format: str, *args: object) -> None: - return None - - -@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) -async def test_image_that_cannot_be_fetched_returns_an_error_value( - custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr(litellm, "user_url_validation", False) - with ThreadingHTTPServer(("127.0.0.1", 0), _MissingImageHandler) as server: - Thread(target=server.serve_forever, daemon=True).start() - try: - contents = await gemini_count_contents( - ( - { - "role": "user", - "content": [ - { - "type": "image", - "source": { - "type": "url", - "url": f"http://127.0.0.1:{server.server_address[1]}/{uuid.uuid4().hex}.png", - }, - } - ], - }, - ), - "anthropic", - "gemini-2.5-flash", - custom_llm_provider, - {}, - ) - finally: - server.shutdown() - - assert isinstance(contents, VertexGeminiCountContextError) async def test_gcs_media_type_lookup_uses_the_deployment_vertex_project() -> None: diff --git a/tests/unit/llms/vertex_ai/gemini/test_long_decode.py b/tests/unit/llms/vertex_ai/gemini/test_long_decode.py index 3fe9a5c7a92..91f099fae02 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_long_decode.py +++ b/tests/unit/llms/vertex_ai/gemini/test_long_decode.py @@ -1,16 +1,12 @@ import json import sys -import threading from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence from dataclasses import dataclass -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from itertools import chain, count, repeat +from itertools import chain, repeat from typing import Final import httpx import pytest -from cryptography.hazmat.primitives import serialization -from cryptography.hazmat.primitives.asymmetric import rsa import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -1185,158 +1181,6 @@ async def test_streaming_continuation_request_is_sent_with_the_callers_timeout(i assert read_timeouts == [7.5, 7.5] -@dataclass(frozen=True, slots=True) -class _TokenEndpoint: - credentials: str - down: threading.Event - - -def _service_account_json(token_uri: str) -> str: - private_key: Final = ( - rsa.generate_private_key(public_exponent=65537, key_size=2048) - .private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()) - .decode() - ) - return json.dumps( - { - "type": "service_account", - "project_id": "p", - "private_key_id": "k", - "private_key": private_key, - "client_email": "sa@p.iam.gserviceaccount.com", - "token_uri": token_uri, - } - ) - - -@pytest.fixture -def token_endpoint() -> Iterator[_TokenEndpoint]: - down: Final = threading.Event() - minted: Final = count(1) - - class Handler(BaseHTTPRequestHandler): - def do_POST(self) -> None: - self.rfile.read(int(self.headers["Content-Length"])) - refused: Final = down.is_set() - payload: Final = json.dumps( - {"error": "invalid_grant"} if refused else {"access_token": f"token-{next(minted)}", "expires_in": 0} - ).encode() - self.send_response(400 if refused else 200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(payload))) - self.end_headers() - self.wfile.write(payload) - - def log_message(self, *_args: object) -> None: - return None - - server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - threading.Thread(target=server.serve_forever, daemon=True).start() - yield _TokenEndpoint(_service_account_json(f"http://127.0.0.1:{server.server_port}/token"), down) - server.shutdown() - server.server_close() - - -def _google_taking_each_access_token_once( - hops: Sequence[httpx.Response], used: list[str], on_answer: Callable[[], object] -) -> Callable[[httpx.Request], httpx.Response]: - def handler(request: httpx.Request) -> httpx.Response: - token: Final = request.headers["authorization"] - if token in used: - return httpx.Response(401, json={"error": {"code": 401, "status": "UNAUTHENTICATED"}}) - used.append(token) - on_answer() - return hops[len(used) - 1] - - return handler - - -async def _vertex_answer( - is_async: bool, - stream: bool, - handler: Callable[[httpx.Request], httpx.Response], - credentials: str, - extra_headers: Mapping[str, str] | None = None, -) -> tuple[str, list[str], list[object]]: - kwargs: Final = { - "model": "vertex_ai/gemini-3.8-flash", - "extra_headers": extra_headers, - "messages": list(MESSAGES), - "max_tokens": 100_000, - "stream": stream, - "vertex_project": "p", - "vertex_location": "global", - "vertex_credentials": credentials, - } - transport: Final = httpx.MockTransport(handler) - response: Final = ( - await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=transport)) - if is_async - else litellm.completion(**kwargs, client=HTTPHandler(client=httpx.Client(transport=transport))) - ) - if stream: - chunks: Final = [chunk async for chunk in response] if is_async else list(response) - return _stream_text(chunks), _stream_finish_reasons(chunks), _client_tokens(chunks) - choice: Final = response.choices[0] - fields: Final = choice.message.provider_specific_fields or {} - return choice.message.content, [choice.finish_reason], [fields["continuation_token"]] if fields else [] - - -def _two_hops(stream: bool) -> tuple[httpx.Response, httpx.Response]: - return _stream_hops() if stream else (httpx.Response(200, json=HOP_1), httpx.Response(200, json=HOP_2)) - - -@pytest.mark.parametrize("stream", [True, False]) -@pytest.mark.parametrize("is_async", [True, False]) -@pytest.mark.asyncio -async def test_vertex_long_decode_finishes_when_each_continuation_outlives_the_access_token_before_it( - is_async: bool, stream: bool, token_endpoint: _TokenEndpoint -): - used: Final[list[str]] = [] - google: Final = _google_taking_each_access_token_once(_two_hops(stream), used, lambda: None) - - answer: Final = await _vertex_answer(is_async, stream, google, token_endpoint.credentials) - - assert answer == ("hello world", ["stop"], []) - assert len(used) == 2 - - -@pytest.mark.parametrize("stream", [True, False]) -@pytest.mark.parametrize("is_async", [True, False]) -@pytest.mark.asyncio -async def test_vertex_long_decode_returns_the_answer_so_far_with_its_token_when_no_new_access_token_can_be_had( - is_async: bool, stream: bool, token_endpoint: _TokenEndpoint -): - used: Final[list[str]] = [] - google: Final = _google_taking_each_access_token_once(_two_hops(stream), used, token_endpoint.down.set) - - answer: Final = await _vertex_answer(is_async, stream, google, token_endpoint.credentials) - - assert answer == ("hello", ["length"], ["dG9rLTE="]) - assert len(used) == 1 - - -@pytest.mark.parametrize("stream", [True, False]) -@pytest.mark.parametrize("is_async", [True, False]) -@pytest.mark.asyncio -async def test_vertex_long_decode_keeps_the_callers_own_authorization_header_on_every_continuation( - is_async: bool, stream: bool, token_endpoint: _TokenEndpoint -): - hops: Final = _two_hops(stream) - seen: Final[list[str]] = [] - - def google(request: httpx.Request) -> httpx.Response: - seen.append(request.headers["authorization"]) - return hops[len(seen) - 1] - - answer: Final = await _vertex_answer( - is_async, stream, google, token_endpoint.credentials, {"Authorization": "Bearer mine"} - ) - - assert answer == ("hello world", ["stop"], []) - assert seen == ["Bearer mine", "Bearer mine"] - - BUDGET_ENV: Final = "LITELLM_GEMINI_DEFAULT_MAX_OUTPUT_TOKENS" ARGON: Final = "gemini/gemini-4-argon" UNMAPPED_GEMINI: Final = "gemini/gemini-99-unreleased" @@ -1555,46 +1399,6 @@ async def test_env_budget_leaves_a_snake_case_generation_config_from_extra_body_ assert _output_limits(sent[0]) == {} -def _gemini_chat_models_answering_no_text() -> tuple[str, ...]: - rows: Final = ( - (name, row) - for name, row in litellm.model_cost.items() - if row.get("mode") == "chat" and "text" not in row.get("supported_output_modalities", ("text",)) - ) - return tuple( - name if name.startswith("gemini/") else f"vertex_ai/{name}" - for name, row in rows - if row.get("litellm_provider") in ("gemini", "vertex_ai-language-models") - ) - - -@pytest.mark.asyncio -async def test_env_budget_skips_every_chat_model_the_cost_map_lists_as_answering_no_text( - monkeypatch: pytest.MonkeyPatch, token_endpoint: _TokenEndpoint -): - monkeypatch.setenv(BUDGET_ENV, "50") - models: Final = _gemini_chat_models_answering_no_text() - sent: Final[list[dict[str, object]]] = [] - replies: Final = tuple(httpx.Response(200, json=HOP_2) for _ in models) - client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_replay(replies, sent))) - - for model in models: - await litellm.acompletion( - model=model, - messages=list(MESSAGES), - api_key="fake-gemini-key", - vertex_project="p", - vertex_location="us-central1", - vertex_credentials=token_endpoint.credentials, - client=client, - ) - - assert models - assert {model: _output_limits(request) for model, request in zip(models, sent, strict=True)} == dict.fromkeys( - models, {} - ) - - @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.asyncio async def test_env_budget_leaves_a_continuation_token_resume_with_no_limit_alone( @@ -1619,28 +1423,6 @@ async def test_env_budget_leaves_a_continuation_token_resume_with_no_limit_alone assert _output_limits(sent[0]) == {} -@pytest.mark.parametrize("model", ["gemini/gemini-2.5-flash", "vertex_ai/gemini-2.5-flash", ARGON]) -@pytest.mark.asyncio -async def test_env_budget_above_a_models_output_limit_sends_that_limit_instead( - monkeypatch: pytest.MonkeyPatch, token_endpoint: _TokenEndpoint, model: str -): - budget: Final = 2 * _cost_map_max_output_tokens(model) - monkeypatch.setenv(BUDGET_ENV, str(budget)) - sent: Final[list[dict[str, object]]] = [] - - await litellm.acompletion( - model=model, - messages=list(MESSAGES), - api_key="fake-gemini-key", - vertex_project="p", - vertex_location="us-central1", - vertex_credentials=token_endpoint.credentials, - client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), - ) - - assert _output_limits(sent[0]) == {"max_output_tokens": _cost_map_max_output_tokens(model)} - - @pytest.mark.asyncio async def test_env_budget_never_goes_above_the_cost_map_limit_when_a_deployment_overrides_it( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/mcp_client/test_mcp_client.py b/tests/unit/mcp_client/test_mcp_client.py index 866a29b70e9..c657e17695c 100644 --- a/tests/unit/mcp_client/test_mcp_client.py +++ b/tests/unit/mcp_client/test_mcp_client.py @@ -15,7 +15,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import anyio import httpx2 import pytest -from mcp import ClientSession, MCPError +from mcp import MCPError from mcp.client.session import ClientRequestContext, ElicitationFnT from mcp.client.streamable_http import streamable_http_client from mcp.server.mcpserver import Context @@ -47,7 +47,7 @@ from starlette.applications import Starlette from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint from starlette.requests import Request from starlette.responses import Response -from starlette.types import ASGIApp, Message, Receive, Scope, Send +from starlette.types import ASGIApp # Add the parent directory to the path so we can import litellm import litellm.mcp_client.client as mcp_client_module @@ -55,7 +55,6 @@ from litellm.mcp_client.client import ( MCPClient, MCPSigV4Auth, PersistentMCPSession, - UpstreamSessionClosedError, _first_non_cancelled_cause, _TransportContext, _TransportStreams, @@ -1511,31 +1510,6 @@ def test_mcp_extra_matches_proxy_extra_and_supports_streamable_http(): assert specifier.contains(mcp_versions[0]) -@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( "auth_type, default_header", [ @@ -2094,220 +2068,6 @@ async def test_invalid_tool_list_schema_is_identified_as_an_upstream_response() assert "secret" not in message -class _DiagnosticSSEStream(httpx2.AsyncByteStream): - def __init__(self, messages: asyncio.Queue[bytes | Exception | None]) -> None: - self.messages = messages - - async def __aiter__(self) -> AsyncIterator[bytes]: - yield b"event: endpoint\ndata: /messages\n\n" - while True: - message: Final = await self.messages.get() - if message is None: - return - if isinstance(message, Exception): - raise message - yield b"event: message\ndata: " + message + b"\n\n" - - -_DIAGNOSTIC_STDIO_SERVER: Final = """ -import json, sys -mode, failure_method = sys.argv[1:] -for line in sys.stdin: - request = json.loads(line) - if "method" not in request or "id" not in request: - continue - if request["method"] == failure_method: - if mode == "bad-json": - print("secret-invalid-json", flush=True) - continue - if mode == "closed": - sys.exit(0) - if mode == "silent": - print(json.dumps({"jsonrpc": "2.0", "method": "notifications/message", "params": {"level": "info", "data": "Waiting"}}), flush=True) - continue - if request["method"] == "initialize": - result = {"protocolVersion": request["params"]["protocolVersion"], "capabilities": {"tools": {}, "logging": {}}, "serverInfo": {"name": "diagnostic", "version": "1"}} - elif request["method"] == "tools/list": - print(json.dumps({"jsonrpc": "2.0", "method": "notifications/message", "params": {"level": "info", "data": "Listing tools"}}), flush=True) - print(json.dumps({"jsonrpc": "2.0", "id": "unmatched", "result": {}}), flush=True) - print(json.dumps({"jsonrpc": "2.0", "id": "server-ping", "method": "ping"}), flush=True) - result = {"tools": [{"name": "ping", "inputSchema": {"type": "object"}}]} - else: - result = {"content": [{"type": "text", "text": "pong"}], "isError": False} - print(json.dumps({"jsonrpc": "2.0", "id": request["id"], "result": result}), flush=True) -""" - - -def _diagnostic_transport(transport: MCPTransport, mode: str, failure_method: str) -> _TransportContext: - from mcp import StdioServerParameters - from mcp.client.sse import sse_client - from mcp.client.stdio import stdio_client - - if transport == MCPTransport.stdio: - return stdio_client( - StdioServerParameters( - command=sys.executable, args=["-u", "-c", _DIAGNOSTIC_STDIO_SERVER, mode, failure_method] - ) - ) - messages: Final[asyncio.Queue[bytes | Exception | None]] = asyncio.Queue() - - async def respond(request: httpx2.Request) -> httpx2.Response: - if request.method == "GET": - 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 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(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 httpx2.Response(202) - if payload["method"] == "tools/list": - for message in ( - { - "jsonrpc": "2.0", - "method": "notifications/message", - "params": {"level": "info", "data": "Listing tools"}, - }, - {"jsonrpc": "2.0", "id": "unmatched", "result": {}}, - {"jsonrpc": "2.0", "id": "server-ping", "method": "ping"}, - ): - await messages.put(json.dumps(message).encode()) - result: Final = ( - { - "protocolVersion": payload["params"]["protocolVersion"], - "capabilities": {"tools": {}, "logging": {}}, - "serverInfo": {"name": "diagnostic", "version": "1"}, - } - if payload["method"] == "initialize" - else {"tools": [{"name": "ping", "inputSchema": {"type": "object"}}]} - if payload["method"] == "tools/list" - else {"content": [{"type": "text", "text": "pong"}], "isError": False} - ) - await messages.put(json.dumps({"jsonrpc": "2.0", "id": payload["id"], "result": result}).encode()) - return httpx2.Response(202) - - def factory( - headers: dict[str, str] | None = None, - 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) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) -@pytest.mark.parametrize("failure_method", ["initialize", "tools/list"]) -async def test_transport_parsing_failure_is_preserved(transport: MCPTransport, failure_method: str) -> None: - client: Final = MCPClient(server_url="https://example.com/sse", transport_type=transport, timeout=0.2) - with pytest.raises(ValidationError): - await asyncio.wait_for( - client._execute_session_operation( - _diagnostic_transport(transport, "bad-json", failure_method), lambda session: session.list_tools() - ), - timeout=3, - ) - - -def test_sse_read_failure_is_preserved() -> None: - loop: Final = AutojumpClockLoop() - - async def run() -> None: - client: Final = MCPClient(server_url="https://example.com/sse", transport_type=MCPTransport.sse, timeout=0.2) - 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(), - ), - timeout=3, - ) - - try: - loop.run_until_complete(run()) - finally: - loop.run_until_complete(loop.shutdown_asyncgens()) - loop.run_until_complete(loop.shutdown_default_executor()) - loop.close() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("protocol_version", ["auto", "2025-06-18"]) -@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) -@pytest.mark.parametrize("mode", ["ok", "closed", "silent"]) -async def test_transport_completion_and_normal_messages( - transport: MCPTransport, mode: str, protocol_version: str -) -> None: - from mcp import ClientSession - - from litellm.proxy.mcp_server.rest_endpoints import _connection_error_message - - logging_callback: Final = AsyncMock() - read_timeout: Final = 0.2 if mode == "silent" else 30 - client: Final = MCPClient( - server_url="https://example.com/sse", - transport_type=transport, - timeout=read_timeout, - logging_callback=logging_callback, - protocol_version=protocol_version, - ) - - async def operation(session: ClientSession) -> CallToolResult: - tools: Final = await session.list_tools() - assert [tool.name for tool in tools.tools] == ["ping"] - return await session.call_tool("ping", {}) - - 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.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: - await asyncio.wait_for(pending, timeout=3) - if mode == "closed": - assert "connection was closed" in _connection_error_message(caught.value, client.server_url, read_timeout) - else: - assert isinstance(as_mcp_read_timeout(caught.value), TimeoutError) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) -async def test_transport_cancellation_cleans_up_a_pending_request(transport: MCPTransport) -> None: - ready: Final = asyncio.Event() - - async def on_log(message: LoggingMessageNotificationParams) -> None: - if message.data == "Waiting": - ready.set() - - client: Final = MCPClient( - server_url="https://example.com/sse", transport_type=transport, timeout=30, logging_callback=on_log - ) - task: Final = asyncio.create_task( - client._execute_session_operation( - _diagnostic_transport(transport, "silent", "tools/list"), lambda session: session.list_tools() - ) - ) - try: - await asyncio.wait_for(ready.wait(), timeout=3) - finally: - task.cancel() - with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(task, timeout=3) - - class _InterruptedHTTPBody(httpx2.AsyncByteStream): async def __aiter__(self) -> AsyncIterator[bytes]: yield b'{"jsonrpc":' @@ -2670,125 +2430,6 @@ async def test_optional_discovery_collects_all_pages(method: str, session_id: st assert sum(call.args[0].method == "DELETE" for call in responder.call_args_list) == (1 if session_id else 0) -@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list")) -@pytest.mark.parametrize( - "failure", ("repeat", "cycle", "cap", "method_not_found", "internal_error", "unauthorized", "deadline") -) -@pytest.mark.parametrize("strict", (False, True)) -def test_optional_discovery_rejects_incomplete_walks( - method: str, failure: str, strict: bool, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture -) -> None: - monkeypatch.setattr( - mcp_client_module, - "MCP_TOOL_LISTING_MAX_PAGES", - 3 if failure in ("cycle", "repeat") else 2, - raising=False, - ) - monkeypatch.setattr(mcp_client_module, "MCP_TOOL_LISTING_TIMEOUT", 0.05) - field: Final = { - "prompts/list": "prompts", - "resources/list": "resources", - "resources/templates/list": "resourceTemplates", - }[method] - entry: Final = { - "prompts/list": {"name": "first"}, - "resources/list": {"name": "first", "uri": "test://first"}, - "resources/templates/list": {"name": "first", "uriTemplate": "test://{name}"}, - }[method] - loop: Final = ManualClockLoop() - - async def run() -> None: - cancelled: Final = asyncio.Event() - - async def respond(request: httpx2.Request) -> httpx2.Response: - payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) - if not isinstance(payload, JSONRPCRequest): - return httpx2.Response(202) - if payload.method == "initialize": - return httpx2.Response( - 200, - json={ - "jsonrpc": "2.0", - "id": payload.id, - "result": { - "protocolVersion": (payload.params or {})["protocolVersion"], - "capabilities": {"prompts": {}, "resources": {}}, - "serverInfo": {"name": "interrupted", "version": "1"}, - }, - }, - ) - assert payload.method == method - cursor: Final = (payload.params or {}).get("cursor") - if cursor is not None: - if failure == "deadline": - loop.advance(0.15) - try: - for _ in range(1_000): - await asyncio.sleep(0) - except asyncio.CancelledError: - cancelled.set() - raise - return httpx2.Response(500) - if failure == "unauthorized": - return httpx2.Response(401) - if failure in ("method_not_found", "internal_error"): - return httpx2.Response( - 200, - json={ - "jsonrpc": "2.0", - "id": payload.id, - "error": { - "code": -32601 if failure == "method_not_found" else -32603, - "message": "Later page unavailable", - }, - }, - ) - if failure == "deadline" and cursor is None: - loop.advance(0.1) - next_cursor: Final = ( - "private-cursor-2" if cursor == "private-cursor-1" and failure != "repeat" else "private-cursor-1" - ) - return httpx2.Response( - 200, json={"jsonrpc": "2.0", "id": payload.id, "result": {field: [entry], "nextCursor": next_cursor}} - ) - - responder: Final = AsyncMock(side_effect=respond) - client: Final = _MockTransportClient(responder, server_url="https://example.com/mcp", timeout=0.2) - operation: Final = { - "prompts/list": client.list_prompts, - "resources/list": client.list_resources, - "resources/templates/list": client.list_resource_templates, - }[method] - if strict: - error_type: Final = { - "internal_error": MCPError, - "unauthorized": httpx2.HTTPStatusError, - "deadline": TimeoutError, - }.get(failure, RuntimeError) - with pytest.raises(error_type): - await operation(raise_on_error=True) - else: - assert await operation() == [] - assert len( - tuple( - payload - for call in responder.call_args_list - if isinstance(payload := _JSONRPC_MESSAGE_ADAPTER.validate_json(call.args[0].content), JSONRPCRequest) - and payload.method == method - ) - ) == (3 if failure == "cycle" else 2) - assert "private-cursor" not in caplog.text - if failure == "deadline": - assert cancelled.is_set() - - try: - loop.run_until_complete(run()) - finally: - loop.run_until_complete(loop.shutdown_asyncgens()) - loop.run_until_complete(loop.shutdown_default_executor()) - loop.close() - - @pytest.mark.asyncio @pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list")) async def test_optional_discovery_allows_exhaustion_at_page_cap(method: str, monkeypatch: pytest.MonkeyPatch) -> None: @@ -2867,24 +2508,6 @@ def test_standalone_client_request_does_not_require_proxy_dependencies(): assert result.stdout.strip() == "standalone request succeeded" -def test_client_import_before_proxy_credentials_succeeds_in_fresh_process(): - import subprocess - - result = subprocess.run( - [ - sys.executable, - "-c", - "import litellm.mcp_client.client; from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager; print(MCPServerManager.__name__)", - ], - capture_output=True, - text=True, - timeout=60, - check=False, - ) - assert result.returncode == 0, result.stderr - assert result.stdout.strip() == "MCPServerManager" - - @pytest.mark.asyncio @pytest.mark.parametrize("resolved", (False, True)) async def test_discovery_auth_fingerprint_tracks_effective_credentials(resolved: bool) -> None: @@ -3227,162 +2850,6 @@ async def test_http_close_cancellation_cannot_turn_into_success(original_error: assert task.cancelled() -@pytest.mark.asyncio -@pytest.mark.parametrize("protocol_version", ("auto", "2025-06-18")) -@pytest.mark.parametrize("cancel_mode", ("scope", "task", "wait_for", "read_timeout")) -@pytest.mark.parametrize("concurrency", (1, 5)) -@pytest.mark.parametrize("termination", ("ok", "hang", "hang_body")) -@pytest.mark.parametrize("raise_on_error", (False, True)) -async def test_cancellation_delivers_termination_over_tcp( - cancel_mode: str, concurrency: int, termination: str, raise_on_error: bool, protocol_version: str -) -> None: - started: Final = asyncio.Event() - scope_ready: Final[asyncio.Future[anyio.CancelScope]] = asyncio.get_running_loop().create_future() - terminations: Final[list[bytes]] = [] - starts: Final[list[bytes]] = [] - stop: Final = asyncio.Event() - connections: Final[list[asyncio.Task[None]]] = [] - - async def handle_connection(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: - connection: Final = asyncio.current_task() - assert connection is not None - connections.append(connection) - try: - request_line: Final = await reader.readline() - if not request_line: - return - method: Final = request_line.split()[0] - headers: Final = await reader.readuntil(b"\r\n\r\n") - length: Final = next( - ( - int(line.split(b":", 1)[1]) - for line in headers.splitlines() - if line.lower().startswith(b"content-length:") - ), - 0, - ) - try: - body: Final = await reader.readexactly(length) - except asyncio.IncompleteReadError: - return - if method == b"DELETE": - terminations.append(body) - if termination != "ok": - await stop.wait() - return - writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") - elif method == b"GET": - writer.write(b"HTTP/1.1 405 Method Not Allowed\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") - else: - payload: Final = json.loads(body) - if payload["method"] == "tools/call": - if termination == "hang_body": - writer.write( - b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n" - b"Content-Length: 200\r\nConnection: close\r\n\r\n" - ) - await writer.drain() - starts.append(body) - if len(starts) == concurrency: - started.set() - await stop.wait() - return - if payload["method"] == "initialize": - if cancel_mode != "read_timeout": - await asyncio.sleep(0.75) - response: Final = json.dumps( - { - "jsonrpc": "2.0", - "id": payload["id"], - "result": { - "protocolVersion": "2025-06-18", - "capabilities": {"tools": {}}, - "serverInfo": {"name": "tcp-peer", "version": "1"}, - }, - } - ).encode() - writer.write( - b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nMcp-Session-Id: tcp-session\r\n" - + f"Content-Length: {len(response)}\r\nConnection: close\r\n\r\n".encode() - + response - ) - else: - writer.write(b"HTTP/1.1 202 Accepted\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") - await writer.drain() - finally: - writer.close() - await writer.wait_closed() - - listener: Final = await asyncio.start_server(handle_connection, "127.0.0.1", 0) - port: Final = listener.sockets[0].getsockname()[1] - client_timeout: Final = 2 if cancel_mode == "read_timeout" else 30 - client: Final = MCPClient( - server_url=f"http://127.0.0.1:{port}/mcp", - protocol_version=protocol_version, - timeout=client_timeout, - ) - - async def calls(): - results: Final = await asyncio.gather( - *( - client.call_tool(CallToolRequestParams(name="slow", arguments={}), raise_on_error=raise_on_error) - for _ in range(concurrency) - ), - return_exceptions=cancel_mode == "read_timeout", - ) - if cancel_mode == "read_timeout": - if raise_on_error: - assert all(isinstance(result, TimeoutError) for result in results) - else: - assert all(isinstance(result, CallToolResult) and result.is_error for result in results) - return results - - async def invoke(): - if cancel_mode == "scope": - with anyio.fail_after(None) as scope: - scope_ready.set_result(scope) - return await calls() - return await calls() - - try: - task: Final = asyncio.create_task(invoke()) - await asyncio.wait_for(started.wait(), 30) - if cancel_mode == "scope": - (await scope_ready).deadline = anyio.current_time() + 0.2 - if cancel_mode == "task": - task.cancel() - expected_error: Final = ( - TimeoutError - if cancel_mode == "read_timeout" - else asyncio.CancelledError - if cancel_mode == "task" - else TimeoutError - ) - if cancel_mode == "read_timeout": - done, _ = await asyncio.wait((task,), timeout=8) - assert task in done, "Read timeout and bounded cleanup must complete without external cancellation" - await task - elif cancel_mode == "wait_for": - with pytest.raises(expected_error): - await asyncio.wait_for(task, 0.2) - else: - with pytest.raises(expected_error): - await task - assert len(starts) == concurrency - assert len(terminations) == concurrency, "Each cancelled call must send DELETE over a fresh TCP connection" - finally: - stop.set() - if not task.done(): - task.cancel() - await asyncio.wait((task,), timeout=8) - listener.close() - for connection in connections: - connection.cancel() - closed: Final = await asyncio.wait_for(asyncio.gather(*connections, return_exceptions=True), 2) - assert all(result is None or isinstance(result, asyncio.CancelledError) for result in closed), closed - await asyncio.wait_for(listener.wait_closed(), 2) - - @pytest.mark.asyncio @pytest.mark.parametrize("revision", ["2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25", "auto"]) @pytest.mark.parametrize("accepted", [True, False]) @@ -3710,80 +3177,6 @@ async def test_modern_call_requires_schema_from_bounded_listing( assert observed.count("tools/call") == int(schema_available) -def test_modern_call_uses_discovery_deadline_for_multiple_pages() -> None: - from mcp.types import ToolsCapability - - loop: Final = AutojumpClockLoop() - - async def respond(request: httpx2.Request) -> httpx2.Response: - payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) - assert isinstance(payload, JSONRPCRequest) - if payload.method == "server/discover": - return httpx2.Response( - 200, - json={ - "jsonrpc": "2.0", - "id": payload.id, - "result": DiscoverResult( - supported_versions=["2026-07-28"], capabilities=ServerCapabilities(tools=ToolsCapability()) - ).model_dump(by_alias=True, exclude_none=True), - }, - ) - if payload.method == "tools/list": - await asyncio.sleep(0.06) - assert payload.params is not None - page: Final = ( - {"tools": [], "nextCursor": "second"} - if "cursor" not in payload.params - else { - "tools": [ - { - "name": "quote", - "inputSchema": { - "type": "object", - "properties": {"workspace": {"type": "string", "x-mcp-header": "Workspace"}}, - }, - } - ] - } - ) - return httpx2.Response( - 200, - json={ - "jsonrpc": "2.0", - "id": payload.id, - "result": {"resultType": "complete", "cacheScope": "private", "ttlMs": 0, **page}, - }, - ) - assert payload.method == "tools/call" - assert request.headers["mcp-param-workspace"] == "engineering" - return httpx2.Response( - 200, - json={ - "jsonrpc": "2.0", - "id": payload.id, - "result": {"resultType": "complete", "content": [{"type": "text", "text": "quoted"}], "isError": False}, - }, - ) - - async def run() -> None: - client: Final = _MockTransportClient( - respond, server_url="https://example.com/mcp", protocol_version="2026-07-28", timeout=0.1 - ) - assert [tool.name for tool in await client.list_tools(raise_on_error=True)] == ["quote"] - result: Final = await client.call_tool( - CallToolRequestParams(name="quote", arguments={"workspace": "engineering"}), raise_on_error=True - ) - assert result.is_error is False - assert result.content[0].text == "quoted" - - try: - loop.run_until_complete(run()) - finally: - loop.run_until_complete(loop.shutdown_asyncgens()) - loop.close() - - @pytest.mark.asyncio async def test_cancelled_modern_catalog_load_prevents_tool_execution() -> None: from queue import SimpleQueue @@ -4138,45 +3531,6 @@ class _StatefulUpstreamClient(MCPClient): return streamable_http_client(self.server_url, http_client=http_client), http_client -class _CleanupHoldingASGITransport(httpx2.ASGITransport): - def __init__(self, app: ASGIApp, cleanup_started: asyncio.Event, release_cleanup: asyncio.Event) -> None: - super().__init__(app=app) - self._cleanup_started: Final = cleanup_started - self._release_cleanup: Final = release_cleanup - - async def aclose(self) -> None: - self._cleanup_started.set() - await self._release_cleanup.wait() - await super().aclose() - - -class _CleanupHoldingStatefulUpstreamClient(_StatefulUpstreamClient): - def __init__( - self, - app: ASGIApp, - *, - server_url: str, - transport_type: MCPTransport, - cleanup_started: asyncio.Event, - release_cleanup: asyncio.Event, - ) -> None: - super().__init__(app, server_url=server_url, transport_type=transport_type) - self._cleanup_started: Final = cleanup_started - self._release_cleanup: Final = release_cleanup - - def _create_transport_context(self) -> tuple[_TransportContext, httpx2.AsyncClient]: - transport: Final = _CleanupHoldingASGITransport( - self._app, - self._cleanup_started, - self._release_cleanup, - ) - http_client: Final = self._create_httpx_client_factory(transport=transport)( - headers=self._get_auth_headers(), - timeout=httpx2.Timeout(self.timeout), - ) - return streamable_http_client(self.server_url, http_client=http_client), http_client - - class _InProcessUpstreamClient(MCPClient): def __init__(self, server: UpstreamServer, *, elicitation_callback: ElicitationFnT | None = None) -> None: super().__init__( @@ -4258,32 +3612,6 @@ def _upstream_with_status_for_method(method: str, status_code: int) -> Starlette return app -def _upstream_with_malformed_response_for_method(method: str) -> Starlette: - service: Final = UpstreamServer("malformed response upstream") - - @service.tool() - def malformed_tool() -> str: - return "unreachable" - - app: Final = service.streamable_http_app( - stateless_http=False, - json_response=True, - transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False), - ) - - async def malformed_method(request: Request, call_next: RequestResponseEndpoint) -> Response: - if request.method == "GET": - return Response(content="data: {}\n\n", media_type="text/event-stream") - if request.method == "POST": - payload: Final = await request.json() - if isinstance(payload, dict) and payload.get("method") == method: - return Response(content="not a JSON-RPC message", media_type="text/event-stream") - return await call_next(request) - - app.add_middleware(BaseHTTPMiddleware, dispatch=malformed_method) - return app - - @pytest.mark.asyncio async def test_persistent_session_keeps_upstream_state_across_tool_calls() -> None: app: Final = _stateful_upstream() @@ -4319,48 +3647,6 @@ def _client_with_session( return client, client.open_persistent_session(admission) -@pytest.mark.asyncio -async def test_cancelling_one_waiter_for_session_readiness_does_not_cancel_other_waiters() -> None: - app: Final = _stateful_upstream() - async with app.router.lifespan_context(app): - initialize_started: Final = asyncio.Event() - release_initialize: Final = asyncio.Event() - - async def gated_server(scope: Scope, receive: Receive, send: Send) -> None: - async def gated_send(message: Message) -> None: - if message["type"] == "http.response.body" and not initialize_started.is_set(): - initialize_started.set() - await release_initialize.wait() - await send(message) - - await app(scope, receive, gated_send) - - gated_app: Final[ASGIApp] = gated_server - client: Final = _StatefulUpstreamClient( - gated_app, server_url="http://upstream/mcp", transport_type=MCPTransport.http - ) - session: Final = client.open_persistent_session() - - async def operation(_: ClientSession) -> str: - return "ready" - - cancelled_waiter: Final = asyncio.create_task(session.run(operation)) - remaining_waiter: Final = asyncio.create_task(session.run(operation)) - try: - await asyncio.wait_for(initialize_started.wait(), 5) - await asyncio.sleep(0) - cancelled_waiter.cancel() - with pytest.raises(asyncio.CancelledError): - await cancelled_waiter - release_initialize.set() - assert await asyncio.wait_for(remaining_waiter, 5) == "ready" - assert not session.closed - finally: - release_initialize.set() - session.close() - await asyncio.wait_for(session.wait_closed(), 5) - - @pytest.mark.asyncio async def test_persistent_session_survives_a_caller_timeout_on_one_operation() -> None: app: Final = _stateful_upstream() @@ -4388,99 +3674,6 @@ async def test_persistent_session_survives_a_caller_timeout_on_one_operation() - await asyncio.wait_for(session.wait_closed(), 5) -@pytest.mark.asyncio -async def test_persistent_session_survives_recoverable_mcp_errors_and_closes_on_connection_closed() -> None: - app: Final = _stateful_upstream() - async with app.router.lifespan_context(app): - client, session = _client_with_session(app) - - async def invalid_operation(_: ClientSession) -> str: - raise MCPError(code=-32602, message="invalid params") - - try: - with pytest.raises(MCPError): - await session.run(invalid_operation) - assert not session.closed - selected: Final = await client.call_tool( - CallToolRequestParams(name="select_project", arguments={"name": "recoverable"}), - persistent_session=session, - ) - created: Final = await client.call_tool( - CallToolRequestParams(name="create_feature", arguments={"title": "still-open"}), - persistent_session=session, - ) - finally: - session.close() - assert selected.is_error is False - assert created.content[0].text == "recoverable/still-open" - await asyncio.wait_for(session.wait_closed(), 5) - - connection_app: Final = _stateful_upstream() - async with connection_app.router.lifespan_context(connection_app): - _, session = _client_with_session(connection_app) - started: Final = asyncio.Event() - release: Final = asyncio.Event() - - async def disconnected_operation(_: ClientSession) -> str: - started.set() - await release.wait() - raise MCPError(code=CONNECTION_CLOSED, message="upstream connection closed") - - async def pending_operation(_: ClientSession) -> str: - return "unused" - - active: Final = asyncio.create_task(session.run(disconnected_operation)) - await asyncio.wait_for(started.wait(), 5) - pending: Final = asyncio.create_task(session.run(pending_operation)) - await asyncio.sleep(0) - release.set() - with pytest.raises(MCPError) as caught: - await asyncio.wait_for(active, 5) - assert caught.value.error.code == CONNECTION_CLOSED - with pytest.raises((MCPError, UpstreamSessionClosedError)): - await asyncio.wait_for(pending, 5) - await asyncio.wait_for(session.wait_closed(), 5) - assert session.closed - - -@pytest.mark.asyncio -async def test_ended_persistent_session_is_not_reusable_during_transport_cleanup() -> None: - app: Final = _stateful_upstream() - async with app.router.lifespan_context(app): - cleanup_started: Final = asyncio.Event() - release_cleanup: Final = asyncio.Event() - client: Final = _CleanupHoldingStatefulUpstreamClient( - app, - server_url="http://upstream/mcp", - transport_type=MCPTransport.http, - cleanup_started=cleanup_started, - release_cleanup=release_cleanup, - ) - session: Final = client.open_persistent_session() - - async def disconnected_operation(_: ClientSession) -> str: - raise MCPError(code=CONNECTION_CLOSED, message="upstream connection closed") - - try: - with pytest.raises(MCPError) as caught: - await asyncio.wait_for(session.run(disconnected_operation), 5) - assert caught.value.error.code == CONNECTION_CLOSED - await asyncio.wait_for(cleanup_started.wait(), 5) - - assert not session.reusable - assert not session.closed - session.retire() - - release_cleanup.set() - await asyncio.wait_for(session.wait_closed(), 5) - assert session.closed - finally: - release_cleanup.set() - if not session.closed: - session.close() - await asyncio.wait_for(session.wait_closed(), 5) - - @pytest.mark.asyncio async def test_persistent_tool_call_preserves_upstream_http_status_error() -> None: app: Final = _upstream_with_status_for_method("tools/call", 401) @@ -4534,104 +3727,6 @@ async def test_persistent_tool_call_classifies_initialized_notification_status(s await asyncio.wait_for(session.wait_closed(), 5) -@pytest.mark.asyncio -async def test_persistent_session_remaps_stream_error_over_mcp_error() -> None: - app: Final = _stateful_upstream() - async with app.router.lifespan_context(app): - client: Final = _StatefulUpstreamClient( - app, - server_url="http://upstream/mcp", - transport_type=MCPTransport.http, - ) - session: Final = client.open_persistent_session() - request: Final = httpx2.Request("POST", "http://upstream/mcp") - response: Final = httpx2.Response(401, request=request, headers={"WWW-Authenticate": "Bearer"}) - http_error: Final = httpx2.HTTPStatusError("401 Unauthorized", request=request, response=response) - - async def operation(_: ClientSession) -> str: - stream_error: Final[asyncio.Future[Exception] | None] = session._stream_error - assert stream_error is not None - stream_error.set_result(http_error) - raise MCPError(code=-32603, message="response stream closed") - - try: - with pytest.raises(httpx2.HTTPStatusError) as caught: - await session.run(operation) - assert caught.value is http_error - assert not session.reusable - finally: - session.close() - await asyncio.wait_for(session.wait_closed(), 5) - - -@pytest.mark.asyncio -async def test_persistent_session_ends_when_a_malformed_response_is_remapped() -> None: - app: Final = _upstream_with_malformed_response_for_method("tools/call") - async with app.router.lifespan_context(app): - client: Final = _StatefulUpstreamClient( - app, - server_url="http://upstream/mcp", - transport_type=MCPTransport.http, - ) - session: Final = client.open_persistent_session() - try: - - async def ready(_: ClientSession) -> str: - return "ready" - - assert await asyncio.wait_for(session.run(ready), 5) == "ready" - stream_error: Final[asyncio.Future[Exception] | None] = session._stream_error - assert stream_error is not None - recorded_stream_error: Final = await asyncio.wait_for(asyncio.shield(stream_error), 5) - assert isinstance(recorded_stream_error, ValidationError) - - with pytest.raises(ValidationError) as caught: - await client.call_tool( - CallToolRequestParams(name="malformed_tool", arguments={}), - raise_on_error=True, - persistent_session=session, - ) - assert caught.value is recorded_stream_error - assert not session.reusable - await asyncio.wait_for(session.wait_closed(), 5) - assert session.closed - finally: - if not session.closed: - session.close() - await asyncio.wait_for(session.wait_closed(), 5) - - -@pytest.mark.asyncio -async def test_failed_persistent_initialize_is_not_reusable_during_transport_cleanup() -> None: - app: Final = _upstream_with_status_for_method("initialize", 500) - async with app.router.lifespan_context(app): - cleanup_started: Final = asyncio.Event() - release_cleanup: Final = asyncio.Event() - client: Final = _CleanupHoldingStatefulUpstreamClient( - app, - server_url="http://upstream/mcp", - transport_type=MCPTransport.http, - cleanup_started=cleanup_started, - release_cleanup=release_cleanup, - ) - session: Final = client.open_persistent_session() - pending: Final = asyncio.create_task(session.run(lambda _: "ready")) - try: - await asyncio.wait_for(cleanup_started.wait(), 5) - assert not session.reusable - assert not session.closed - release_cleanup.set() - with pytest.raises(httpx2.HTTPStatusError): - await asyncio.wait_for(pending, 5) - await asyncio.wait_for(session.wait_closed(), 5) - assert session.closed - finally: - release_cleanup.set() - if not session.closed: - session.close() - await asyncio.wait_for(session.wait_closed(), 5) - - @pytest.mark.asyncio async def test_persistent_session_routes_elicitation_to_the_current_clients_callback() -> None: service: Final = UpstreamServer("elicitation") @@ -4687,406 +3782,6 @@ async def test_persistent_session_without_elicitation_callback_uses_sdk_default( assert "elicitation callback is unavailable" not in str(persistent.model_dump(mode="json")) -@pytest.mark.asyncio -async def test_cancelled_queued_operation_does_not_run_and_session_accepts_a_later_operation() -> None: - app: Final = _stateful_upstream() - async with app.router.lifespan_context(app): - _, session = _client_with_session(app) - first_started: Final = asyncio.Event() - release_first: Final = asyncio.Event() - cancelled_operation_ran: Final = asyncio.Event() - - async def first_operation(_: ClientSession) -> str: - first_started.set() - await release_first.wait() - return "first" - - async def cancelled_operation(_: ClientSession) -> str: - cancelled_operation_ran.set() - return "cancelled" - - async def third_operation(_: ClientSession) -> str: - return "third" - - first_task: Final = asyncio.create_task(session.run(first_operation)) - try: - await asyncio.wait_for(first_started.wait(), 5) - second_task: Final = asyncio.create_task(session.run(cancelled_operation)) - await asyncio.sleep(0) - assert session._queue.qsize() == 1 - second_task.cancel() - with pytest.raises(asyncio.CancelledError): - await second_task - third_task: Final = asyncio.create_task(session.run(third_operation)) - await asyncio.sleep(0) - assert session._queue.qsize() == 2 - release_first.set() - assert await asyncio.wait_for(first_task, 5) == "first" - assert await asyncio.wait_for(third_task, 5) == "third" - assert not cancelled_operation_ran.is_set() - assert not session.closed - finally: - release_first.set() - session.close() - await asyncio.wait_for(session.wait_closed(), 5) - - -@pytest.mark.asyncio -async def test_retiring_persistent_session_drains_in_flight_and_queued_operations() -> None: - app: Final = _stateful_upstream() - async with app.router.lifespan_context(app): - _, session = _client_with_session(app) - first_started: Final = asyncio.Event() - release_first: Final = asyncio.Event() - - async def first_operation(_: ClientSession) -> str: - first_started.set() - await release_first.wait() - return "first" - - async def queued_operation(_: ClientSession) -> str: - return "queued" - - first_waiter: Final = asyncio.create_task(session.run(first_operation)) - try: - await asyncio.wait_for(first_started.wait(), 5) - queued_waiter: Final = asyncio.create_task(session.run(queued_operation)) - await asyncio.sleep(0) - assert session._queue.qsize() == 1 - session.retire() - release_first.set() - - results: Final = await asyncio.wait_for(asyncio.gather(first_waiter, queued_waiter), 5) - assert results == ["first", "queued"] - await asyncio.wait_for(session.wait_closed(), 5) - assert session.closed - finally: - release_first.set() - if not session.closed: - session.close() - await asyncio.wait_for(session.wait_closed(), 5) - - -@pytest.mark.asyncio -async def test_timed_out_persistent_operation_waiting_for_admission_is_skipped() -> None: - app: Final = _stateful_upstream() - async with app.router.lifespan_context(app): - admission_slots: Final = asyncio.Semaphore(1) - admission_requested: Final = asyncio.Event() - - def admission() -> AbstractAsyncContextManager[None]: - admission_requested.set() - return admission_slots - - _, session = _client_with_session(app, admission) - operation_ran: Final = asyncio.Event() - - async def warmup(_: ClientSession) -> str: - return "ready" - - async def timed_out_operation(_: ClientSession) -> str: - operation_ran.set() - return "expired" - - async def later_operation(_: ClientSession) -> str: - return "later" - - try: - assert await session.run(warmup) == "ready" - admission_requested.clear() - async with admission_slots: - timed_out_call: Final = asyncio.create_task(session.run(timed_out_operation)) - await asyncio.wait_for(admission_requested.wait(), 1) - with pytest.raises(asyncio.TimeoutError): - await asyncio.wait_for(timed_out_call, 0.05) - assert not operation_ran.is_set() - - later_result: Final = await asyncio.wait_for(session.run(later_operation), 1) - - assert later_result == "later" - assert session.reusable - assert not operation_ran.is_set() - finally: - session.close() - await asyncio.wait_for(session.wait_closed(), 5) - - -@pytest.mark.asyncio -async def test_persistent_session_holds_admission_until_initialize_completes() -> None: - app: Final = _stateful_upstream() - async with app.router.lifespan_context(app): - initialize_seen: Final = asyncio.Event() - admission_slots: Final = asyncio.Semaphore(1) - - async def observing_app(scope: Scope, receive: Receive, send: Send) -> None: - async def observe_send(message: Message) -> None: - if message["type"] == "http.response.body": - initialize_seen.set() - await send(message) - - await app(scope, receive, observe_send) - - client: Final = _StatefulUpstreamClient( - observing_app, - server_url="http://upstream/mcp", - transport_type=MCPTransport.http, - ) - session: Final = client.open_persistent_session(lambda: admission_slots) - - async def operation(_: ClientSession) -> str: - return "ready" - - run_task: Final = asyncio.create_task(session.run(operation)) - try: - async with admission_slots: - with pytest.raises(asyncio.TimeoutError): - await asyncio.wait_for(initialize_seen.wait(), 0.05) - assert not session._ready.done() - assert not run_task.done() - - result: Final = await asyncio.wait_for(run_task, 5) - assert result == "ready" - assert initialize_seen.is_set() - finally: - session.close() - await asyncio.wait_for(session.wait_closed(), 5) - await asyncio.gather(run_task, return_exceptions=True) - - -@pytest.mark.asyncio -async def test_persistent_session_releases_admission_when_connect_fails() -> None: - async def failing_app(_: Scope, __: Receive, ___: Send) -> None: - raise httpx2.ConnectError("connect failed") - - client: Final = _StatefulUpstreamClient( - failing_app, - server_url="http://upstream/mcp", - transport_type=MCPTransport.http, - ) - admission_slots: Final = asyncio.Semaphore(1) - session: Final = client.open_persistent_session(lambda: admission_slots) - - async def operation(_: ClientSession) -> str: - return "unreachable" - - try: - with pytest.raises(httpx2.ConnectError, match="connect failed"): - await asyncio.wait_for(session.run(operation), 1) - await asyncio.wait_for(session.wait_closed(), 1) - - await asyncio.wait_for(admission_slots.acquire(), 0.1) - admission_slots.release() - finally: - session.close() - await asyncio.wait_for(session.wait_closed(), 5) - - -@pytest.mark.asyncio -async def test_failed_connect_is_not_reusable_during_admission_teardown() -> None: - connect_failed: Final = asyncio.Event() - admission_exit_started: Final = asyncio.Event() - release_admission_exit: Final = asyncio.Event() - - async def failing_app(_: Scope, __: Receive, ___: Send) -> None: - connect_failed.set() - raise httpx2.ConnectError("connect failed") - - @asynccontextmanager - async def blocking_admission() -> AsyncIterator[None]: - try: - yield - finally: - admission_exit_started.set() - await release_admission_exit.wait() - - client: Final = _StatefulUpstreamClient( - failing_app, - server_url="http://upstream/mcp", - transport_type=MCPTransport.http, - ) - session: Final = client.open_persistent_session(blocking_admission) - - async def operation(_: ClientSession) -> str: - return "unreachable" - - run_task: Final = asyncio.create_task(session.run(operation)) - try: - await asyncio.wait_for(connect_failed.wait(), 1) - await asyncio.wait_for(admission_exit_started.wait(), 1) - - assert not run_task.done() - assert not session.reusable - assert not session.closed - - release_admission_exit.set() - with pytest.raises(httpx2.ConnectError, match="connect failed"): - await asyncio.wait_for(run_task, 1) - await asyncio.wait_for(session.wait_closed(), 1) - assert session.closed - finally: - release_admission_exit.set() - session.close() - await asyncio.wait_for(session.wait_closed(), 5) - await asyncio.gather(run_task, return_exceptions=True) - - -@pytest.mark.asyncio -async def test_persistent_session_admission_is_acquired_only_when_operation_runs() -> None: - app: Final = _stateful_upstream() - async with app.router.lifespan_context(app): - admission_slots: Final = asyncio.Semaphore(2) - _, session = _client_with_session(app, lambda: admission_slots) - first_started: Final = asyncio.Event() - release_first: Final = asyncio.Event() - operation_order: Final[list[str]] = [] - - async def first_operation(_: ClientSession) -> str: - operation_order.append("first") - first_started.set() - await release_first.wait() - return "first" - - async def second_operation(_: ClientSession) -> str: - operation_order.append("second") - return "second" - - async def third_operation(_: ClientSession) -> str: - operation_order.append("third") - return "third" - - first_task: Final = asyncio.create_task(session.run(first_operation)) - second_task: Final = asyncio.create_task(session.run(second_operation)) - third_task: Final = asyncio.create_task(session.run(third_operation)) - try: - await asyncio.wait_for(first_started.wait(), 5) - await asyncio.sleep(0) - assert session._queue.qsize() == 2 - await asyncio.wait_for(admission_slots.acquire(), 0.1) - try: - assert admission_slots.locked() - finally: - admission_slots.release() - assert operation_order == ["first"] - - release_first.set() - results: Final = await asyncio.wait_for( - asyncio.gather(first_task, second_task, third_task), - 5, - ) - - assert results == ["first", "second", "third"] - assert operation_order == ["first", "second", "third"] - finally: - release_first.set() - session.close() - await asyncio.wait_for(session.wait_closed(), 5) - await asyncio.gather(first_task, second_task, third_task, return_exceptions=True) - - -@pytest.mark.asyncio -async def test_retiring_idle_persistent_session_closes_and_rejects_new_operations() -> None: - app: Final = _stateful_upstream() - async with app.router.lifespan_context(app): - _, session = _client_with_session(app) - - async def operation(_: ClientSession) -> str: - return "ready" - - try: - assert await session.run(operation) == "ready" - session.retire() - await asyncio.wait_for(session.wait_closed(), 1) - - assert session.closed - with pytest.raises(UpstreamSessionClosedError): - await session.run(operation) - finally: - if not session.closed: - session.close() - await asyncio.wait_for(session.wait_closed(), 5) - - -@pytest.mark.asyncio -async def test_closing_persistent_session_mid_operation_fails_the_waiter_instead_of_hanging() -> None: - app: Final = _stateful_upstream() - async with app.router.lifespan_context(app): - _, session = _client_with_session(app) - started: Final = asyncio.Event() - - async def slow_operation(_: object) -> str: - started.set() - await asyncio.sleep(30) - return "never" - - waiter: Final = asyncio.ensure_future(session.run(slow_operation)) - await asyncio.wait_for(started.wait(), 5) - session.close() - with pytest.raises(RuntimeError, match="upstream MCP session closed"): - await asyncio.wait_for(waiter, 5) - await asyncio.wait_for(session.wait_closed(), 5) - - -@pytest.mark.asyncio -async def test_persistent_session_reports_an_upstream_cancellation_as_a_runtime_error_not_a_cancelled_caller() -> None: - app: Final = _stateful_upstream() - async with app.router.lifespan_context(app): - client, session = _client_with_session(app) - - async def cancelled_upstream(_: object) -> str: - raise asyncio.CancelledError() - - try: - with pytest.raises(RuntimeError, match="cancelled"): - await asyncio.wait_for(session.run(cancelled_upstream), 5) - selected: Final = await client.call_tool( - CallToolRequestParams(name="select_project", arguments={"name": "e"}), persistent_session=session - ) - finally: - session.close() - assert selected.is_error is False, "an upstream cancellation must not be mistaken for a cancelled caller" - await asyncio.wait_for(session.wait_closed(), 5) - - -@pytest.mark.asyncio -async def test_closing_persistent_session_fails_a_caller_blocked_on_a_full_queue_instead_of_hanging() -> None: - from litellm.mcp_client.client import _MAX_PENDING_OPERATIONS - - app: Final = _stateful_upstream() - async with app.router.lifespan_context(app): - _, session = _client_with_session(app) - started: Final = asyncio.Event() - - async def slow_operation(_: object) -> str: - started.set() - await asyncio.sleep(30) - return "never" - - waiters: Final = tuple( - asyncio.ensure_future(session.run(slow_operation)) for _ in range(_MAX_PENDING_OPERATIONS + 2) - ) - await asyncio.wait_for(started.wait(), 5) - await asyncio.sleep(0) - session.close() - outcomes: Final = await asyncio.wait_for(asyncio.gather(*waiters, return_exceptions=True), 5) - assert all(isinstance(outcome, RuntimeError) for outcome in outcomes), outcomes - await asyncio.wait_for(session.wait_closed(), 5) - - -@pytest.mark.asyncio -async def test_persistent_session_ends_after_a_broken_stream_so_the_next_call_gets_a_fresh_one() -> None: - app: Final = _stateful_upstream() - async with app.router.lifespan_context(app): - _, session = _client_with_session(app) - - async def broken_stream(_: object) -> str: - raise anyio.BrokenResourceError() - - with pytest.raises(anyio.BrokenResourceError): - await asyncio.wait_for(session.run(broken_stream), 5) - await asyncio.wait_for(session.wait_closed(), 5) - assert session.closed, "a dead transport must end the session instead of being reused" - - @pytest.mark.asyncio @pytest.mark.parametrize("session_started_in_request", [True, False]) async def test_persistent_session_releases_the_current_requests_scope_response(session_started_in_request) -> None: @@ -5134,76 +3829,6 @@ async def test_persistent_session_releases_the_current_requests_scope_response(s session.close() await asyncio.wait_for(session.wait_closed(), 5) -@pytest.mark.asyncio -@pytest.mark.parametrize("legacy_callback", [False, True]) -async def test_tool_progress_preserves_message_and_legacy_callback(legacy_callback: bool) -> None: - received = AsyncMock() - - async def legacy(progress: float, total: float | None) -> None: - await received(progress, total) - - async def modern(progress: float, total: float | None, message: str | None) -> None: - await received(progress, total, message) - - async def call_tool(**kwargs): - await kwargs["progress_callback"](1.0, 2.0, "halfway") - return CallToolResult(content=[]) - - session = AsyncMock(spec=ClientSession, call_tool=AsyncMock(side_effect=call_tool)) - - async def run(operation, **kwargs): - return await operation(session) - - persistent = MagicMock(spec=PersistentMCPSession, run=AsyncMock(side_effect=run)) - client = MCPClient(server_url="https://example.com/mcp") - result = await client.call_tool( - CallToolRequestParams(name="slow"), - host_progress_callback=legacy if legacy_callback else modern, - persistent_session=persistent, - ) - assert not result.is_error - if legacy_callback: - received.assert_awaited_once_with(1.0, 2.0) - else: - received.assert_awaited_once_with(1.0, 2.0, "halfway") - - -@pytest.mark.asyncio -async def test_cancelling_persistent_operation_stops_work_and_keeps_session_usable() -> None: - app: Final = _stateful_upstream() - async with app.router.lifespan_context(app): - client, session = _client_with_session(app) - started: Final = asyncio.Event() - stopped: Final = asyncio.Event() - - async def slow_operation(_: ClientSession) -> None: - started.set() - try: - await asyncio.Event().wait() - finally: - stopped.set() - - waiter: Final = asyncio.create_task(session.run(slow_operation)) - try: - await asyncio.wait_for(started.wait(), 5) - waiter.cancel() - with pytest.raises(asyncio.CancelledError): - await waiter - await asyncio.sleep(0.1) - assert stopped.is_set(), "Caller cancellation must stop the running upstream operation" - result: Final = await asyncio.wait_for( - client.call_tool( - CallToolRequestParams(name="select_project", arguments={"name": "after-cancellation"}), - persistent_session=session, - ), - 5, - ) - assert not result.is_error - finally: - session.close() - await asyncio.wait_for(session.wait_closed(), 5) - await asyncio.gather(waiter, return_exceptions=True) - @pytest.mark.asyncio @pytest.mark.parametrize("header_name", ["Authorization", "x-upstream-token"]) diff --git a/tests/unit/proxy/auth/test_master_key_boot_check.py b/tests/unit/proxy/auth/test_master_key_boot_check.py index dc490068c4d..50075a97c6c 100644 --- a/tests/unit/proxy/auth/test_master_key_boot_check.py +++ b/tests/unit/proxy/auth/test_master_key_boot_check.py @@ -1,10 +1,6 @@ import asyncio import re -import shutil -import subprocess -import sys from collections.abc import Mapping -from pathlib import Path from typing import Final import pytest @@ -26,7 +22,6 @@ from litellm.proxy.auth.master_key_boot_check import ( UnsafeMasterKeyError, UnsafeMasterKeyReason, UnsafeMasterKeyRefused, - announce_on_stderr_at_exit, enforce_master_key_boot_verdict, master_key_boot_verdict, render_refusal, @@ -361,24 +356,6 @@ def test_migration_with_an_exported_key_replaces_it_in_place_and_numbers_every_s assert text.index("os.environ/") < text.index(MIGRATE_FROM_MASTER_KEY_ENV_VAR + "=") < text.index("Start the proxy") -@pytest.mark.skipif(shutil.which("openssl") is None, reason="the printed command shells out to openssl") -@pytest.mark.parametrize("from_master_key", [PUBLICLY_KNOWN_KEY, "", " " + PUBLICLY_KNOWN_KEY]) -def test_printed_migration_commands_save_both_keys_to_the_env_file(tmp_path: Path, from_master_key: str): - from dotenv import dotenv_values - - text = render_refusal( - _refusal(migration=StoredSecretsMigration(from_master_key=from_master_key, encrypted_value_count=1)) - ) - commands = [line.strip() for line in text.splitlines() if line.strip().startswith("echo ")] - - subprocess.run(["bash", "-c", "\n".join(commands)], cwd=tmp_path, capture_output=True, text=True, check=True) - - saved = dotenv_values(tmp_path / ".env") - assert saved[MIGRATE_FROM_MASTER_KEY_ENV_VAR] == from_master_key - assert _verdict(saved[MASTER_KEY_ENV_VAR]) == SafeMasterKey() - assert len(commands) == 2 - - @pytest.mark.parametrize("migration", [_MIGRATION, None]) def test_override_hint_is_the_last_paragraph(migration: StoredSecretsMigration | None): text = render_refusal(_refusal(migration=migration)) @@ -388,30 +365,6 @@ def test_override_hint_is_the_last_paragraph(migration: StoredSecretsMigration | assert f"general_settings.{WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_SETTING}" in last_paragraph -@pytest.mark.skipif(shutil.which("openssl") is None, reason="the printed command shells out to openssl") -def test_printed_command_saves_a_key_the_boot_check_accepts(tmp_path: Path): - completed = subprocess.run( - ["bash", "-c", GENERATE_MASTER_KEY_COMMAND], cwd=tmp_path, capture_output=True, text=True, check=True - ) - - saved = (tmp_path / ".env").read_text() - match = re.fullmatch(rf"{MASTER_KEY_ENV_VAR}=(sk-[0-9a-f]{{64}})\n", saved) - assert match is not None - assert completed.stdout == saved - assert _verdict(match.group(1)) == SafeMasterKey() - - -@pytest.mark.skipif(shutil.which("openssl") is None, reason="the printed command shells out to openssl") -def test_rotation_command_prints_a_key_the_boot_check_accepts_and_saves_nothing(tmp_path: Path): - completed = subprocess.run( - ["bash", "-c", PRINT_NEW_MASTER_KEY_COMMAND], cwd=tmp_path, capture_output=True, text=True, check=True - ) - - assert re.fullmatch(r"sk-[0-9a-f]{64}\n", completed.stdout) is not None - assert _verdict(completed.stdout.strip()) == SafeMasterKey() - assert list(tmp_path.iterdir()) == [] - - def test_refusal_announces_the_fix_and_aborts_the_boot(): announced: list[str] = [] refusal = _refusal(source=ConfigFileSource(config_file_path="/app/config.yaml")) @@ -429,19 +382,3 @@ def test_safe_and_overridden_keys_boot_without_announcing(verdict: MasterKeyBoot enforce_master_key_boot_verdict(verdict, announce=announced.append) assert announced == [] - - -def test_announced_fix_is_the_last_thing_a_crashing_process_prints(): - crash_after_announcing = ( - "from litellm.proxy.auth.master_key_boot_check import announce_on_stderr_at_exit\n" - "announce_on_stderr_at_exit('THE FIX')\n" - "print('buffered stdout')\n" - "raise RuntimeError('lifespan failed')\n" - ) - - completed = subprocess.run([sys.executable, "-c", crash_after_announcing], capture_output=True, text=True) - - assert completed.returncode != 0 - assert "RuntimeError: lifespan failed" in completed.stderr - assert completed.stderr.endswith("THE FIX") - assert completed.stdout == "buffered stdout\n" diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index b5dc6042474..ba0f436f3dd 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -3,16 +3,12 @@ import contextvars import json import logging import os -import subprocess -import sys import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from contextlib import contextmanager, nullcontext from datetime import datetime, timedelta, timezone from functools import partial from io import BytesIO -from pathlib import Path -from textwrap import dedent from types import MappingProxyType, SimpleNamespace from typing import Final, Literal from unittest.mock import ANY, AsyncMock, MagicMock, patch @@ -12170,129 +12166,6 @@ def test_the_request_span_parents_datastore_spans_of_the_request(monkeypatch, al assert contextvars.copy_context().run(ambient_parent_after_auth) is request_span -_RECORDING_DDTRACE = dedent( - """ - import functools - import inspect - - - class _Span: - def __enter__(self): - return self - - def __exit__(self, *exc): - return None - - - class _Tracer: - def __init__(self): - self.spans = [] - - def wrap(self, name=None, **kwargs): - def decorator(f): - span_name = name or f"{f.__module__}.{f.__name__}" - if inspect.iscoroutinefunction(f): - - @functools.wraps(f) - async def async_wrapped(*args, **kw): - self.spans.append(span_name) - return await f(*args, **kw) - - return async_wrapped - - @functools.wraps(f) - def wrapped(*args, **kw): - self.spans.append(span_name) - return f(*args, **kw) - - return wrapped - - return decorator - - def trace(self, name, **kwargs): - return _Span() - - def current_span(self): - return None - - def current_root_span(self): - return None - - - tracer = _Tracer() - """ -) - -_DDTRACE_AUTH_PROBE = dedent( - """ - import asyncio - import json - - import ddtrace - from fastapi import Request - from starlette.datastructures import URL - - import litellm.proxy.proxy_server as proxy_server - from litellm.proxy._types import ProxyException - from litellm.proxy.auth.user_api_key_auth import user_api_key_auth - - proxy_server.master_key = "sk-probe" - - - async def auth(api_key): - request = Request(scope={"type": "http", "headers": [], "method": "POST", "path": "/chat/completions"}) - request._url = URL(url="/chat/completions") - try: - await user_api_key_auth( - request=request, - api_key=api_key, - azure_api_key_header="", - anthropic_api_key_header=None, - google_ai_studio_api_key_header=None, - azure_apim_header=None, - custom_litellm_key_header=None, - ) - return "accepted" - except ProxyException: - return "rejected" - - - async def main(): - outcomes = [await auth("Bearer sk-probe"), await auth("Bearer sk-wrong")] - print(json.dumps({"outcomes": outcomes, "spans": ddtrace.tracer.spans})) - - - asyncio.run(main()) - """ -) - - -def test_user_api_key_auth_opens_a_datadog_span_for_accepted_and_rejected_keys(tmp_path: Path): - stub_root = tmp_path / "site" - (stub_root / "ddtrace").mkdir(parents=True) - (stub_root / "ddtrace" / "__init__.py").write_text(_RECORDING_DDTRACE) - probe = tmp_path / "probe.py" - probe.write_text(_DDTRACE_AUTH_PROBE) - repo_root = Path(litellm.__file__).resolve().parent.parent - env = { - **os.environ, - "USE_DDTRACE": "true", - "PYTHONPATH": os.pathsep.join( - [str(stub_root), str(repo_root)] + [p for p in (os.environ.get("PYTHONPATH"),) if p] - ), - } - - result = subprocess.run( - [sys.executable, str(probe)], env=env, cwd=repo_root, capture_output=True, text=True, check=False - ) - - assert result.returncode == 0, result.stderr[-4000:] - report = json.loads(result.stdout.strip().splitlines()[-1]) - assert report["outcomes"] == ["accepted", "rejected"] - auth_span = "litellm.proxy.auth.user_api_key_auth.user_api_key_auth" - assert [span for span in report["spans"] if span == auth_span] == [auth_span, auth_span] - - @pytest.mark.asyncio @pytest.mark.parametrize("is_proxy_admin", [False, True], ids=["standard-return", "proxy-admin-return"]) async def test_jwt_builder_returns_every_team_grant_the_key_path_gets(is_proxy_admin): diff --git a/tests/unit/proxy/client/cli/test_claude_settings.py b/tests/unit/proxy/client/cli/test_claude_settings.py index 80aa86c5058..655d70d4ef9 100644 --- a/tests/unit/proxy/client/cli/test_claude_settings.py +++ b/tests/unit/proxy/client/cli/test_claude_settings.py @@ -1,9 +1,6 @@ import json import os -import pathlib -import shlex import stat -import subprocess import sys import time from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError @@ -56,7 +53,6 @@ def _owners(*backup_paths): CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings" -AUTH_MODULE = "litellm.proxy.client.cli.commands.auth" @pytest.fixture def paths(tmp_path): return tmp_path / "claude" / "settings.json", tmp_path / "backup.json" @@ -174,35 +170,6 @@ class TestConfigureClaudeSettings: class TestLiteApiKeyHelper: - @pytest.mark.skipif(sys.platform.startswith("win"), reason="runs the command through sh, as Claude Code does on POSIX") - def test_sh_runs_the_lite_on_path_against_this_proxy(self, tmp_path): - lite = tmp_path / "my tools" / "lite" - lite.parent.mkdir() - lite.write_text('#!/bin/sh\nprintf "%s\\n" "$0" "$@"\n') - lite.chmod(0o700) - - helper = lite_api_key_helper("https://proxy.example.com/", platform="linux", which=lambda name: str(lite)) - - assert helper is not None - ran = subprocess.run(["sh", "-c", helper.command], capture_output=True, text=True, check=True) - assert ran.stdout.splitlines() == [str(lite), "--base-url", "https://proxy.example.com", "auth", "print-token"] - - @pytest.mark.skipif(sys.platform.startswith("win"), reason="runs the command through sh, as Claude Code does on POSIX") - def test_a_lite_found_through_a_relative_path_entry_keeps_running_from_any_directory(self, tmp_path, monkeypatch): - for repo, name in ((tmp_path / "login-dir", "mine"), (tmp_path / "other-repo", "planted")): - lite = repo / "bin" / "lite" - lite.parent.mkdir(parents=True) - lite.write_text(f"#!/bin/sh\necho {name}\n") - lite.chmod(0o700) - monkeypatch.chdir(tmp_path / "login-dir") - - helper = lite_api_key_helper("https://proxy.example.com", platform="linux", which=lambda name: "bin/lite") - - assert helper is not None - ran = subprocess.run( - ["sh", "-c", helper.command], cwd=tmp_path / "other-repo", capture_output=True, text=True, check=True - ) - assert ran.stdout == "mine\n" def test_quotes_every_token_for_cmd_on_windows(self): helper = lite_api_key_helper( @@ -262,23 +229,6 @@ class TestInstalledExecutable: path.write_text(f"#!/bin/sh\n{body}\n") path.chmod(0o700) - def test_the_helper_runs_the_lite_installed_beside_this_python_over_a_shim_earlier_on_path( - self, tmp_path, monkeypatch - ): - self._script(tmp_path / "shims" / "lite", "echo 'pyenv: lite: command not found' >&2; exit 127") - self._script(tmp_path / "venv" / "bin" / "lite", "echo real") - monkeypatch.setenv("PATH", f"{tmp_path / 'shims'}{os.pathsep}{os.environ['PATH']}") - - helper = lite_api_key_helper( - "https://proxy.example.com", - platform="linux", - which=lambda name: installed_executable(name, scripts_dir=str(tmp_path / "venv" / "bin")), - ) - - assert helper is not None - ran = subprocess.run(["sh", "-c", helper.command], capture_output=True, text=True) - assert (ran.returncode, ran.stdout) == (0, "real\n") - def test_falls_back_to_path_when_this_python_has_no_lite(self, tmp_path, monkeypatch): self._script(tmp_path / "pipx" / "lite", "echo pipx") (tmp_path / "venv" / "bin").mkdir(parents=True) @@ -931,18 +881,6 @@ class TestStatusLine: merged = merge_claude_settings({}, PROXY, StaticToken("tok"), status_line=self.COMMAND) assert merged["statusLine"] == {"type": "command", "command": self.COMMAND} - def test_the_installed_script_is_the_bundled_one_and_the_command_runs_this_interpreter(self, tmp_path): - from litellm.proxy.client.cli.commands import statusline_script - - script = tmp_path / "lite" / "statusline.py" - command = install_statusline_script(script) - assert script.read_bytes().split(b"\n", 1)[1] == pathlib.Path(statusline_script.__file__).read_bytes() - assert shlex.split(command) == [sys.executable, str(script)] - assert command == statusline_command(script) - assert stat.S_IMODE(script.stat().st_mode) == 0o600 - assert stat.S_IMODE(script.parent.stat().st_mode) == 0o700 - assert install_statusline_script(script) == command - def test_a_reinstall_replaces_the_script_in_one_step_and_a_refused_one_leaves_the_old_script_whole(self, tmp_path): # Claude Code may be running the script at the moment `lite` reinstalls it; the file it has open # must stay complete, and a reinstall that cannot land must not leave a truncated script behind. diff --git a/tests/unit/proxy/conftest.py b/tests/unit/proxy/conftest.py index 1aec7831b8e..05463bd3555 100644 --- a/tests/unit/proxy/conftest.py +++ b/tests/unit/proxy/conftest.py @@ -2,15 +2,12 @@ import _hashlib import asyncio -import copy import hashlib import importlib -import inspect import logging import os import sys import tempfile -import warnings from collections.abc import Awaitable, Callable, Iterator from contextlib import contextmanager from types import ModuleType @@ -36,6 +33,7 @@ from litellm.integrations.otel.model.payloads import ServiceSpanData from litellm.integrations.otel.model.spans import service_span_name from tests.unit.litellm_core_utils.fake_secret_vault import FakeSecretVault from tests._master_key import MASTER_KEY +from tests._module_state import collect_flushable_caches, flush_caches, restore_mutable_state, snapshot_mutable_state class StubClientNotConnectedError(ClientNotConnectedError): @@ -56,104 +54,9 @@ def disconnected_prisma() -> DisconnectedPrisma: return DisconnectedPrisma() -# Top-level assignments of these types are the ones importlib.reload(litellm) -# would have effectively reset. We snapshot them at conftest import time and -# deep-copy the snapshot back before every test. -_SNAPSHOT_TYPES = (list, dict, set, tuple, str, int, float, bool, bytes) - - -def _snapshot_mutable_state(module): - """Capture a per-module snapshot of primitive and collection attributes.""" - snapshot = {} - for attr in list(vars(module)): - if attr.startswith("_"): - continue - try: - value = getattr(module, attr) - except Exception as exc: - warnings.warn( - f"conftest: could not read {module.__name__}.{attr} during snapshot: {exc}", - stacklevel=2, - ) - continue - if value is None or isinstance(value, _SNAPSHOT_TYPES): - try: - snapshot[attr] = _restored_value(value) - except Exception as exc: - warnings.warn( - f"conftest: could not snapshot {module.__name__}.{attr}: {exc}", - stacklevel=2, - ) - return snapshot - - -_MUTABLE_CONTAINERS = (list, dict, set, bytearray) - - -def _holds_mutable_container(value) -> bool: - if isinstance(value, _MUTABLE_CONTAINERS): - return True - if isinstance(value, tuple): - return any(_holds_mutable_container(element) for element in value) - return False - - -def _restored_value(value): - return copy.deepcopy(value) if _holds_mutable_container(value) else value - - -def _restore_mutable_state(module, snapshot): - for attr, default in snapshot.items(): - try: - setattr(module, attr, _restored_value(default)) - except Exception as exc: - warnings.warn( - f"conftest: could not restore {module.__name__}.{attr}: {exc}", - stacklevel=2, - ) - - -def _collect_flushable_caches(): - """Return (module, attr) pairs whose values expose flush_cache().""" - targets = [] - for module in (litellm, litellm.proxy.proxy_server): - for attr in list(vars(module)): - if attr.startswith("_"): - continue - try: - value = getattr(module, attr) - except Exception: - continue - # Only instances — a class reference has an unbound flush_cache - # that can't be called without a self argument. - if inspect.isclass(value) or inspect.ismodule(value): - continue - if callable(getattr(value, "flush_cache", None)): - targets.append((module, attr)) - return targets - - -def _flush_caches(targets): - for module, attr in targets: - try: - value = getattr(module, attr) - except Exception: - continue - flush = getattr(value, "flush_cache", None) - if callable(flush): - try: - flush() - except Exception as exc: - warnings.warn( - f"conftest: flush_cache failed on {module.__name__}.{attr}: {exc}", - stacklevel=2, - ) - - -# Snapshot once at conftest import — these are the "clean" module states. -_LITELLM_STATE = _snapshot_mutable_state(litellm) -_PROXY_SERVER_STATE = _snapshot_mutable_state(litellm.proxy.proxy_server) -_FLUSHABLE_CACHES = _collect_flushable_caches() +_LITELLM_STATE: Final = snapshot_mutable_state(litellm) +_PROXY_SERVER_STATE: Final = snapshot_mutable_state(litellm.proxy.proxy_server) +_FLUSHABLE_CACHES: Final = collect_flushable_caches((litellm, litellm.proxy.proxy_server)) _PROXY_APP_LIFESPAN: Final = litellm.proxy.proxy_server.app.router.lifespan_context @@ -182,9 +85,9 @@ def setup_and_teardown(): Use pytest's monkeypatch.setattr() or a local fixture for those cases — don't rely on this autouse fixture to undo them. """ - _restore_mutable_state(litellm, _LITELLM_STATE) - _restore_mutable_state(litellm.proxy.proxy_server, _PROXY_SERVER_STATE) - _flush_caches(_FLUSHABLE_CACHES) + restore_mutable_state(litellm, _LITELLM_STATE) + restore_mutable_state(litellm.proxy.proxy_server, _PROXY_SERVER_STATE) + flush_caches(_FLUSHABLE_CACHES) loop = asyncio.get_event_loop_policy().new_event_loop() asyncio.set_event_loop(loop) diff --git a/tests/unit/proxy/hooks/litellm_skills/test_main.py b/tests/unit/proxy/hooks/litellm_skills/test_main.py index dc7ebf6e2ec..f716a8533d8 100644 --- a/tests/unit/proxy/hooks/litellm_skills/test_main.py +++ b/tests/unit/proxy/hooks/litellm_skills/test_main.py @@ -1,7 +1,3 @@ -import os -import subprocess -import sys -from typing import Final from unittest.mock import AsyncMock, patch import pytest @@ -69,19 +65,3 @@ async def test_execute_code_loop_dispatches_litellm_skill_tool(): mock_exec.assert_awaited_once() assert mock_exec.await_args.args[0] == SKILL_TOOL_NAME assert result is final_response - - -def test_importing_proxy_hooks_does_not_mutate_litellm_callbacks(): - script: Final = ( - "import litellm; " - "litellm.callbacks = []; " - "import litellm.proxy.hooks; " - "assert len(litellm.callbacks) == 0, f'mutated: {litellm.callbacks}'" - ) - result: Final = subprocess.run( - [sys.executable, "-c", script], - capture_output=True, - text=True, - env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, - ) - assert result.returncode == 0, f"stdout: {result.stdout}\nstderr: {result.stderr}" diff --git a/tests/unit/proxy/mcp_server/conftest.py b/tests/unit/proxy/mcp_server/conftest.py index da5b7d8d9a3..605d910e0b7 100644 --- a/tests/unit/proxy/mcp_server/conftest.py +++ b/tests/unit/proxy/mcp_server/conftest.py @@ -1,11 +1,15 @@ import asyncio -import importlib import os +from typing import Final import pytest import litellm from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from tests._module_state import collect_flushable_caches, flush_caches, restore_mutable_state, snapshot_mutable_state + +_LITELLM_STATE: Final = snapshot_mutable_state(litellm) +_FLUSHABLE_CACHES: Final = collect_flushable_caches((litellm,)) @pytest.fixture(scope="session") @@ -20,11 +24,8 @@ def event_loop(): @pytest.fixture(scope="function", autouse=True) def setup_and_teardown(): - """ - This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. - """ - importlib.reload(litellm) - import asyncio + restore_mutable_state(litellm, _LITELLM_STATE) + flush_caches(_FLUSHABLE_CACHES) loop = asyncio.get_event_loop_policy().new_event_loop() asyncio.set_event_loop(loop) diff --git a/tests/unit/proxy/mcp_server/test_openapi_to_mcp_generator.py b/tests/unit/proxy/mcp_server/test_openapi_to_mcp_generator.py index adb185e4242..a658433b00c 100644 --- a/tests/unit/proxy/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/unit/proxy/mcp_server/test_openapi_to_mcp_generator.py @@ -1636,21 +1636,3 @@ class TestBoundedOpenAPISpecLoading: else: assert await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=100) == {"paths": {}} assert destination.call_count == 1 - - -def test_openapi_generator_import_does_not_require_mcp_sdk() -> None: - import subprocess - import sys - - script = """ -import builtins -original_import = builtins.__import__ -def without_mcp(name, *args, **kwargs): - if name == 'mcp' or name.startswith('mcp.'): - raise ModuleNotFoundError('MCP SDK unavailable') - return original_import(name, *args, **kwargs) -builtins.__import__ = without_mcp -import litellm.proxy.mcp_server.openapi_to_mcp_generator -""" - result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True) - assert result.returncode == 0, result.stderr diff --git a/tests/unit/proxy/rust_forward/supervisor/test_runtime.py b/tests/unit/proxy/rust_forward/supervisor/test_runtime.py index c5d1bbc7a36..1dda28b7e1f 100644 --- a/tests/unit/proxy/rust_forward/supervisor/test_runtime.py +++ b/tests/unit/proxy/rust_forward/supervisor/test_runtime.py @@ -1,5 +1,4 @@ import signal -import subprocess import sys from dataclasses import dataclass, field from pathlib import Path @@ -9,9 +8,7 @@ from typing import Final import anyio import pytest -from litellm.proxy.rust_forward.platform import ChildProcess, Platform, current_platform -from litellm.proxy.rust_forward.platform.windows import WindowsPlatform -from litellm.proxy.rust_forward.supervisor import supervise +from litellm.proxy.rust_forward.platform import ChildProcess, current_platform from litellm.proxy.rust_forward.supervisor.children import ChildSpec from litellm.proxy.rust_forward.supervisor.lifecycle import ( Action, @@ -24,42 +21,6 @@ from litellm.proxy.rust_forward.supervisor.lifecycle import ( Spawn, ) from litellm.proxy.rust_forward.supervisor.runtime import SpawnFailed, drive, exit_like, open_children, spawn_subprocess -from tests.unit.proxy.rust_forward.supervisor._support import python_server - - -PYTHON_CHILD: Final = """ -import json, os, pathlib, socket, time -time.sleep(0.5) -pathlib.Path(os.environ["MARKERS"], "listening").touch() -private = json.loads(os.environ["_LITELLM_GATEWAY_SUPERVISED"])["listener"] -listener = socket.socket(fileno=private["fd"]) if "fd" in private else socket.create_server(("127.0.0.1", private["port"])) -listener.listen() -time.sleep(10) -""" -RUST_CHILD: Final = """ -import os, pathlib -markers = pathlib.Path(os.environ["MARKERS"]) -(markers / "rust_saw_python_listening").write_text(str((markers / "listening").exists())) -""" - - -def _fake_gateway(scripts: Path) -> Platform: - if sys.platform != "win32": - gateway: Final = scripts / "litellm-gateway" - gateway.write_text(f"#!{sys.executable} -I\n{RUST_CHILD}") - gateway.chmod(0o755) - return current_platform() - script: Final = scripts / "rust_child.py" - script.write_text(RUST_CHILD) - (scripts / "litellm-gateway.cmd").write_text(f'@"{sys.executable}" -I "{script}"\r\n') - return WindowsPlatform(gateway_name="litellm-gateway.cmd") - - -def test_rust_opens_the_public_port_only_once_python_listens(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("MARKERS", str(tmp_path)) - command: Final = (sys.executable, "-I", "-c", PYTHON_CHILD) - supervise(python_server(), command, scripts_dir=str(tmp_path), platform=_fake_gateway(tmp_path)) - assert (tmp_path / "rust_saw_python_listening").read_text() == "True" TERM: Final = signal.SIGTERM @@ -190,15 +151,6 @@ async def test_a_cancelled_supervisor_still_waits_for_the_child_it_stops() -> No assert (while_child_runs, once_child_exits) == (((child,), 1, False), True) -async def test_exiting_children_are_reported_as_events() -> None: - async with open_children( - current_platform(), spawn_subprocess, _children((sys.executable, "-c", "exit(7)")), 1 - ) as children: - assert await children.spawn(Child.RUST) is None - with anyio.fail_after(5): - assert await children.next_event() == ChildExited(Child.RUST, 7) - - async def test_a_child_that_cannot_start_is_returned_as_a_failure(tmp_path: Path) -> None: async with open_children( current_platform(), spawn_subprocess, _children((str(tmp_path / "missing"),)), 1 @@ -223,13 +175,3 @@ def _dropped(_signum: signal.Signals) -> None: ) def test_a_supervisor_the_kernel_will_not_signal_still_exits_as_python_alone_would(status: int, exits: int) -> None: assert exit_like(status, STOPS, die_by=_dropped) == exits - - -@POSIX_ONLY -@pytest.mark.parametrize("signum", [signal.SIGTERM, signal.SIGKILL]) -def test_the_supervisor_dies_by_the_signal_that_killed_python(signum: signal.Signals) -> None: - script: Final = ( - "from litellm.proxy.rust_forward.supervisor.runtime import exit_like\n" - f"raise SystemExit(exit_like({-signum}, frozenset()))" - ) - assert subprocess.run((sys.executable, "-I", "-c", script), check=False, timeout=60).returncode == -signum diff --git a/tests/unit/proxy/test_prometheus_metrics_server.py b/tests/unit/proxy/test_prometheus_metrics_server.py index e52110e3b75..b080493a63b 100644 --- a/tests/unit/proxy/test_prometheus_metrics_server.py +++ b/tests/unit/proxy/test_prometheus_metrics_server.py @@ -151,37 +151,6 @@ def test_main_rejects_missing_multiproc_dir(monkeypatch: pytest.MonkeyPatch): run.assert_not_called() -def test_start_metrics_server_process_returns_only_once_child_serves(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) - _write_worker_sample(pid=3001, value=4) - port: Final = _free_port() - with patch("atexit.register") as register: - process: Final = start_metrics_server_process(host="127.0.0.1", port=port, multiproc_dir=str(tmp_path)) - try: - register.assert_called_once_with(process.terminate) - assert process.poll() is None - startup_metrics: Final = httpx.get(f"http://127.0.0.1:{port}/metrics", follow_redirects=True, timeout=5.0) - assert startup_metrics.status_code == 200 - assert startup_metrics.headers[PID_HEADER] == str(process.pid) - metrics: Final = httpx.get(f"http://127.0.0.1:{port}/metrics", follow_redirects=True, timeout=10.0) - assert 'litellm_requests_metric_total{model="gpt-5"} 4.0' in metrics.text - finally: - process.kill() - process.wait(timeout=10) - - -def test_start_metrics_server_process_fails_when_port_is_taken(tmp_path: Path): - with socket.socket() as occupied: - occupied.bind(("127.0.0.1", 0)) - occupied.listen() - port: Final = occupied.getsockname()[1] - with ( - patch("atexit.register"), - pytest.raises(MetricsServerStartupError), - ): - start_metrics_server_process(host="127.0.0.1", port=port, multiproc_dir=str(tmp_path)) - - class _ImpostorMetrics(BaseHTTPRequestHandler): """An unrelated service already on the port that answers /metrics with 200 and plausible metrics.""" @@ -196,22 +165,6 @@ class _ImpostorMetrics(BaseHTTPRequestHandler): return -def test_start_metrics_server_process_rejects_metrics_from_another_service_on_the_port(tmp_path: Path): - with ThreadingHTTPServer(("127.0.0.1", 0), _ImpostorMetrics) as impostor: - threading.Thread(target=impostor.serve_forever, daemon=True).start() - port: Final = impostor.server_address[1] - impostor_response: Final = httpx.get(f"http://127.0.0.1:{port}/metrics") - assert impostor_response.status_code == 200 - assert "# HELP impostor_metric" in impostor_response.text - assert PID_HEADER not in impostor_response.headers - with ( - patch("atexit.register"), - pytest.raises(MetricsServerStartupError), - ): - start_metrics_server_process(host="127.0.0.1", port=port, multiproc_dir=str(tmp_path)) - impostor.shutdown() - - def test_metrics_server_process_serves_and_exits_with_parent(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) _write_worker_sample(pid=2001, value=7) diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 7fb034022d0..fdc38e3434a 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -3201,33 +3201,6 @@ async def test_user_config_router_enforces_its_limit_across_proxy_requests( ) -@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) -@pytest.mark.asyncio -async def test_concurrent_user_config_requests_each_count_once_against_the_tpm_limit( - fake_redis_port: int, arrives_with_logging_obj: bool -) -> None: - tpm_for_three_requests: Final = await _tpm_for_three_requests() - - def request() -> Coroutine[None, None, str]: - return _unsettled_user_config_request_outcome( - fake_redis_port, - {"routing_strategy": "usage-based-routing-v2"}, - {"tpm": tpm_for_three_requests}, - _proxy_logging_obj() if arrives_with_logging_obj else None, - ) - - async def settled(*requests: Coroutine[None, None, str]) -> list[str]: - outcomes: Final = await asyncio.gather(*requests) - await _settle_router_accounting() - return outcomes - - await _wait_out_a_rate_limit_minute_rollover() - concurrent_outcomes: Final = await settled(request(), request()) - later_outcomes: Final = [*(await settled(request())), *(await settled(request()))] - - assert [*concurrent_outcomes, *later_outcomes] == ["200", "200", "200", "429"] - - @pytest.mark.asyncio async def test_overlapping_user_config_requests_share_one_rpm_limit(fake_redis_port: int) -> None: def request() -> Coroutine[None, None, str]: diff --git a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py index 627f00b96b9..9835e0e2017 100644 --- a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -1,8 +1,6 @@ import asyncio import importlib -import subprocess import sys -import textwrap import types from types import MappingProxyType from collections.abc import Mapping, Sequence @@ -1032,52 +1030,6 @@ async def test_execute_tool_calls_propagates_request_tags_to_function_setup(monk assert captured["metadata"]["tags"] == ["team-a", "prod"] -def test_completion_with_function_tools_works_without_fastapi_installed(): - script = textwrap.dedent( - """ - import sys - - class _FastapiBlocker: - def find_spec(self, fullname, path=None, target=None): - if fullname == "fastapi" or fullname.startswith("fastapi."): - raise ModuleNotFoundError("No module named 'fastapi'") - return None - - sys.meta_path.insert(0, _FastapiBlocker()) - - import litellm - - response = litellm.completion( - model="openai/gpt-5.5", - messages=[{"role": "user", "content": "What is the weather in SF?"}], - tools=[ - { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get the current weather for a location", - "parameters": { - "type": "object", - "properties": {"location": {"type": "string"}}, - "required": ["location"], - }, - }, - } - ], - mock_response="sunny", - ) - assert response.choices[0].message.content == "sunny" - """ - ) - result = subprocess.run( - [sys.executable, "-c", script], - capture_output=True, - text=True, - timeout=120, - ) - assert result.returncode == 0, result.stderr - - def test_extract_tool_call_details_reads_anthropic_tool_use_input(): """ Regression test (LIT-4517): an Anthropic tool_use block carries its arguments diff --git a/tests/unit/rust_bridge/test_configuration.py b/tests/unit/rust_bridge/test_configuration.py index 70223421471..714718b9afc 100644 --- a/tests/unit/rust_bridge/test_configuration.py +++ b/tests/unit/rust_bridge/test_configuration.py @@ -1,8 +1,4 @@ from __future__ import annotations - -import os -import subprocess -import sys from collections.abc import Generator from concurrent.futures import ThreadPoolExecutor from typing import Final @@ -132,21 +128,3 @@ def test_process_override_and_reset_apply_to_existing_threads() -> None: assert executor.submit(configuration.rust_enabled).result() is True configuration.reset_rust_configuration() assert executor.submit(configuration.rust_enabled).result() is False - - -@pytest.mark.parametrize(("value", "expected"), (("1", "True"), ("0", "False"))) -def test_environment_controls_startup(value: str, expected: str) -> None: - environment: Final = {**os.environ, "LITELLM_RUST": value} - result: Final = subprocess.run( - ( - sys.executable, - "-c", - "from litellm.rust_bridge.configuration import rust_enabled; print(rust_enabled())", - ), - check=True, - capture_output=True, - text=True, - env=environment, - ) - - assert result.stdout.strip() == expected diff --git a/tests/unit/test_lazy_imports.py b/tests/unit/test_lazy_imports.py index 1f39bdede5a..850eac03011 100644 --- a/tests/unit/test_lazy_imports.py +++ b/tests/unit/test_lazy_imports.py @@ -1,7 +1,4 @@ import importlib -import os -import subprocess -import sys from typing import Final import pytest @@ -12,22 +9,6 @@ from litellm._lazy_imports_registry import LITELLM_ATTRIBUTES, UTILS_ATTRIBUTES from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter -def test_import_litellm_does_not_load_fastapi_or_bpe_table(): - result: Final = subprocess.run( - [ - sys.executable, - "-c", - "import sys, litellm; print(','.join(m for m in ('fastapi','starlette','litellm.proxy.proxy_cli','litellm.litellm_core_utils.default_encoding','openai.types.realtime') if m in sys.modules))", - ], - check=True, - capture_output=True, - text=True, - env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, - ) - - assert result.stdout.strip() == "" - - def _declared(location: tuple[str, str]) -> object: module_path, attr_name = location return getattr(importlib.import_module(module_path, package="litellm"), attr_name) diff --git a/tests/unit/test_litellm_cli.py b/tests/unit/test_litellm_cli.py index 6015bd1458b..fcd3ca6bf05 100644 --- a/tests/unit/test_litellm_cli.py +++ b/tests/unit/test_litellm_cli.py @@ -1,8 +1,4 @@ -import json -import os -import subprocess import sys -from collections.abc import Mapping from pathlib import Path from typing import Final @@ -16,373 +12,8 @@ if sys.version_info >= (3, 11): else: import tomli as tomllib -_HELP_ARGV: Final = ( - ("--help",), - ("logout", "--help"), - ("login", "--help"), - ("models", "list", "--help"), - ("claude", "--help"), - ("opencode", "--help"), - ("codex", "--help"), - ("pi", "--help"), - ("keys", "list", "--help"), -) - -_COMMAND_ARGV: Final = (("logout",), ("whoami",)) - -_RUN_EACH_ARGV: Final = """ -import contextlib, io, json, sys -from {module} import cli -runs = {{}} -for argv in json.loads(sys.argv[1]): - sys.argv = ["lite", *argv] - output = io.StringIO() - with contextlib.redirect_stdout(output): - try: - cli() - except SystemExit as stopped: - code = stopped.code - else: - code = 0 - runs[" ".join(argv)] = {{"exit": code, "output": output.getvalue()}} -print(json.dumps({{"runs": runs, "sdk_initialized": "completion" in vars(sys.modules["litellm"])}})) -""" - -_USE_SDK_AFTER_HELP: Final = """ -import builtins, importlib, json, sys -interpreter_imports = (builtins.__import__, importlib.import_module) -from {module} import cli -sys.argv = ["lite", "--help"] -try: - cli() -except SystemExit: - pass -before = "completion" in vars(sys.modules["litellm"]) -{first_use} -from litellm.router import Router -import litellm -import litellm.responses.main -from litellm import acompletion -print(json.dumps({{ - "before": before, - "router": Router.__name__, - "completion": callable(litellm.completion), - "acompletion": acompletion is litellm.acompletion, - "responses_is_function": callable(litellm.responses) and not isinstance(litellm.responses, type(litellm)), - "namespace": sorted(vars(litellm)), - "module_type": type(litellm).__name__, - "imports_restored": (builtins.__import__, importlib.import_module) == interpreter_imports, -}})) -""" - -_FIRST_USE_FROM_THREADS: Final = """ -import json, sys, threading -from litellm_cli.main import cli -sys.argv = ["lite", "--help"] -try: - cli() -except SystemExit: - pass -import litellm -start = threading.Barrier(4) -seen = [] -def use_sdk(): - start.wait() - try: - seen.append(callable(litellm.completion)) - except AttributeError as error: - seen.append(str(error)) -threads = [threading.Thread(target=use_sdk) for _ in range(4)] -for thread in threads: - thread.start() -for thread in threads: - thread.join() -print(json.dumps({"seen": seen})) -""" - -_STAR_IMPORT_AFTER_HELP: Final = """ -import json, sys -from litellm_cli.main import cli -sys.argv = ["lite", "--help"] -try: - cli() -except SystemExit: - pass -before = "completion" in vars(sys.modules["litellm"]) -from litellm import * -print(json.dumps({"before": before, "completion": callable(globals().get("completion")), "acompletion": callable(globals().get("acompletion"))})) -""" - -_SDK_WHOSE_SUBMODULE_READS_IT_WHILE_LOADING: Final = { - "sdk_sync.py": """ -import threading -initializing = threading.Event() -submodule_loading = threading.Event() -""", - "litellm/__init__.py": """ -import sdk_sync -sdk_sync.initializing.set() -sdk_sync.submodule_loading.wait(2) -from litellm import submodule -ready = True -""", - "litellm/submodule.py": """ -import litellm -import sdk_sync -sdk_sync.submodule_loading.set() -ready_when_loaded = getattr(litellm, "ready", False) -""", - "litellm/_process_defaults.py": "def apply_process_defaults() -> None: ...\n", - "litellm/proxy/__init__.py": "", - "litellm/proxy/client/__init__.py": "", - "litellm/proxy/client/cli/__init__.py": "", - "litellm/proxy/client/cli/main.py": """ -import json, threading -import litellm -import sdk_sync - -def cli(): - def initialize_sdk(): - litellm.ready - - def import_submodule_once_initialization_starts(): - sdk_sync.initializing.wait() - import litellm.submodule - - threads = [threading.Thread(target=initialize_sdk, daemon=True), - threading.Thread(target=import_submodule_once_initialization_starts, daemon=True)] - for thread in threads: - thread.start() - for thread in threads: - thread.join(10) - print(json.dumps({"finished": [not thread.is_alive() for thread in threads]})) - -litellm_proxy_cli = cli -""", -} - -_SDK_WHOSE_INIT_FAILS: Final = { - "litellm/__init__.py": 'raise ImportError("a dependency of the SDK is broken")\n', - "litellm/utils.py": "", - "litellm/_process_defaults.py": "def apply_process_defaults() -> None: ...\n", - "litellm/proxy/__init__.py": "", - "litellm/proxy/client/__init__.py": "", - "litellm/proxy/client/cli/__init__.py": "", - "litellm/proxy/client/cli/main.py": """ -import json -import litellm - -def error_from(use): - try: - use() - except Exception as error: - return f"{type(error).__name__}: {error}" - return None - -def import_sdk_code(): - import litellm.utils - -def cli(): - print(json.dumps([error_from(lambda: litellm.completion), error_from(lambda: litellm.completion), error_from(import_sdk_code)])) - -litellm_proxy_cli = cli -""", -} - -_RUN_LITE_AGAINST_SDK_AT: Final = """ -import sys -sys.path.insert(0, sys.argv[1]) -from litellm_cli.main import cli -cli() -""" - _PYPROJECT: Final = Path(__file__).resolve().parents[2] / "pyproject.toml" -_RUN_ENTRY_POINT_HELP: Final = """ -import contextlib, io, json, sys -from {module} import {function} -imported_sdk = "litellm" in sys.modules -sys.argv = ["litellm", "--help"] -output = io.StringIO() -with contextlib.redirect_stdout(output): - try: - {function}() - except SystemExit as stopped: - code = stopped.code - else: - code = 0 -print(json.dumps({{"imported_sdk": imported_sdk, "exit": code, "help": output.getvalue()}})) -""" - -_WORKER_LOGGERS_UNDER_A_DEFAULT_LOG_CONFIG: Final = """ -import json, logging, os -os.environ["LITELLM_WORKER_LOG_CONFIG"] = json.dumps({{"version": 1, "root": {{"level": "INFO"}}}}) -from litellm_cli.main import proxy_worker_app -proxy_worker_app() -print(json.dumps({{name: logging.getLogger(name).disabled for name in {names!r}}})) -""" - -_COMMAND_MODULES_LOADED_BY: Final = """ -import json, sys -from litellm_cli.main import cli -sys.argv = ["lite", *sys.argv[1:]] -try: - cli() -except SystemExit: - pass -prefix = "litellm.proxy.client.cli.commands." -print(json.dumps({"modules": sorted(name[len(prefix):] for name in sys.modules if name.startswith(prefix))})) -""" - - -def _run(source: str, argv: tuple[str, ...], home: Path) -> dict[str, object]: - checkout: Final = str(_PYPROJECT.parent) - bootstrap: Final = ( - "import pathlib,sys\n" - f"sys.path.insert(0, {checkout!r})\n" - f"sys.path.insert(0, str(pathlib.Path({checkout!r}) / '.github/packaging'))\n" - "from cli_package import editable_support\n" - f"exec(editable_support(pathlib.Path({checkout!r}) / 'litellm/cli')[0][1])\n" - ) - result: Final = subprocess.run( - [sys.executable, "-I", "-c", bootstrap + source, *argv], - capture_output=True, - text=True, - timeout=120, - env={**os.environ, "HOME": str(home), "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, - ) - assert result.returncode == 0, result.stderr - return json.loads(result.stdout.rstrip("\n").rpartition("\n")[2]) - - -@pytest.fixture(scope="module") -def cli_runs(tmp_path_factory: pytest.TempPathFactory) -> Mapping[str, dict[str, object]]: - argv: Final = (json.dumps(_HELP_ARGV + _COMMAND_ARGV),) - return { - "thin": _run(_RUN_EACH_ARGV.format(module="litellm_cli.main"), argv, tmp_path_factory.mktemp("thin")), - "full": _run( - _RUN_EACH_ARGV.format(module="litellm.proxy.client.cli.main"), argv, tmp_path_factory.mktemp("full") - ), - } - - -def test_litellm_command_imports_the_proxy_only_when_it_runs(tmp_path: Path): - module, _, function = tomllib.loads(_PYPROJECT.read_text())["project"]["scripts"]["litellm"].partition(":") - run: Final = _run(_RUN_ENTRY_POINT_HELP.format(module=module, function=function), (), tmp_path) - assert run["imported_sdk"] is False - assert run["exit"] == 0 - assert "--num_workers" in str(run["help"]) - - -def test_a_worker_log_config_silences_the_sdk_loggers_but_not_the_ones_the_proxy_app_creates(tmp_path: Path): - disabled: Final = _run( - _WORKER_LOGGERS_UNDER_A_DEFAULT_LOG_CONFIG.format(names=("LiteLLM Proxy", "litellm.proxy.health_check")), - (), - tmp_path, - ) - assert disabled == {"LiteLLM Proxy": True, "litellm.proxy.health_check": False} - - -def test_lite_runs_commands_without_initializing_the_sdk(cli_runs: Mapping[str, dict[str, object]]): - assert cli_runs["thin"]["sdk_initialized"] is False - assert cli_runs["full"]["sdk_initialized"] is True - - -@pytest.mark.parametrize("argv", _HELP_ARGV + _COMMAND_ARGV, ids=" ".join) -def test_lite_prints_what_the_full_sdk_cli_prints(argv: tuple[str, ...], cli_runs: Mapping[str, dict[str, object]]): - thin: Final = cli_runs["thin"]["runs"] - full: Final = cli_runs["full"]["runs"] - assert isinstance(thin, dict) and isinstance(full, dict) - - assert thin[" ".join(argv)] == full[" ".join(argv)] - assert thin[" ".join(argv)]["output"] - - -@pytest.mark.parametrize( - ("argv", "loaded", "not_loaded"), - [ - (("logout", "--help"), {"auth"}, {"agents", "autoroute.commands", "models", "keys"}), - (("models", "list", "--help"), {"models"}, {"agents", "autoroute.commands", "keys"}), - (("claude", "--help"), {"agents"}, {"autoroute.commands", "models", "keys"}), - ], -) -def test_lite_imports_only_the_command_it_runs( - argv: tuple[str, ...], loaded: set[str], not_loaded: set[str], tmp_path: Path -): - modules: Final = _run(_COMMAND_MODULES_LOADED_BY, argv, tmp_path)["modules"] - assert isinstance(modules, list) - - assert loaded <= set(modules) - assert not not_loaded & set(modules) - - -@pytest.fixture(scope="module") -def sdk_used_after_full_cli(tmp_path_factory: pytest.TempPathFactory) -> dict[str, object]: - return _run( - _USE_SDK_AFTER_HELP.format(module="litellm.proxy.client.cli.main", first_use="pass"), - (), - tmp_path_factory.mktemp("full"), - ) - - -@pytest.mark.parametrize( - "first_use", - [ - "import litellm; litellm.completion", - "import litellm; litellm.types.utils.ModelResponse", - "from litellm import acompletion", - "from litellm.router import Router", - "from litellm.types import utils", - "import litellm.responses.main", - 'importlib.import_module("litellm.router")', - 'importlib.import_module(".utils", "litellm")', - 'exec("from .types import utils", {"__name__": "litellm.lazy_import", "__package__": "litellm"})', - ], -) -def test_sdk_initializes_in_place_on_first_use_after_lite_imports( - first_use: str, sdk_used_after_full_cli: dict[str, object], tmp_path: Path -): - thin: Final = _run(_USE_SDK_AFTER_HELP.format(module="litellm_cli.main", first_use=first_use), (), tmp_path) - - assert thin["before"] is False - assert sdk_used_after_full_cli["before"] is True - assert {**thin, "before": True} == sdk_used_after_full_cli - assert (thin["router"], thin["completion"], thin["acompletion"], thin["responses_is_function"]) == ( - "Router", - True, - True, - True, - ) - - -def test_star_import_after_lite_imports_gets_the_whole_sdk(tmp_path: Path): - assert _run(_STAR_IMPORT_AFTER_HELP, (), tmp_path) == {"before": False, "completion": True, "acompletion": True} - - -def test_threads_that_use_the_sdk_at_once_all_wait_for_one_initialization(tmp_path: Path): - used: Final = _run(_FIRST_USE_FROM_THREADS, (), tmp_path) - - assert used == {"seen": [True, True, True, True]} - - -def _write_sdk(root: Path, files: Mapping[str, str]) -> Path: - for relative_path, source in files.items(): - (root / relative_path).parent.mkdir(parents=True, exist_ok=True) - (root / relative_path).write_text(source) - return root - - -def test_importing_sdk_code_while_another_thread_initializes_the_sdk_does_not_hang(tmp_path: Path): - sdk_root: Final = _write_sdk(tmp_path / "sdk", _SDK_WHOSE_SUBMODULE_READS_IT_WHILE_LOADING) - - assert _run(_RUN_LITE_AGAINST_SDK_AT, (str(sdk_root),), tmp_path) == {"finished": [True, True]} - - -def test_every_use_after_a_failed_sdk_init_raises_the_init_error(tmp_path: Path): - sdk_root: Final = _write_sdk(tmp_path / "sdk", _SDK_WHOSE_INIT_FAILS) - - assert _run(_RUN_LITE_AGAINST_SDK_AT, (str(sdk_root),), tmp_path) == ["ImportError: a dependency of the SDK is broken"] * 3 - def test_lite_keeps_an_sdk_that_was_already_imported( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] diff --git a/tests/unit/test_router/test_router_fuzz_regressions_rust_only.py b/tests/unit/test_router/test_router_fuzz_regressions_rust_only.py index 8551e6ea884..3c55c2a2c15 100644 --- a/tests/unit/test_router/test_router_fuzz_regressions_rust_only.py +++ b/tests/unit/test_router/test_router_fuzz_regressions_rust_only.py @@ -3,15 +3,12 @@ Router behavior the router-code differential fuzz run found changed once the Rus only core: count settings below zero, an allowed_fails_policy write that is no policy, a model response object as a deployment's mock_response, and endpoints read or patched on the Router class. """ - -import asyncio import gc import json import socket import threading import weakref from collections.abc import Iterator, Mapping -from datetime import datetime from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Final from unittest import mock @@ -245,95 +242,6 @@ def upstream() -> Iterator[str]: thread.join() -def _served(name: str, deployment_id: str, model: str, api_base: str) -> dict[str, object]: - return { - "model_name": name, - "litellm_params": {"model": model, "api_key": "k", "api_base": api_base}, - "model_info": {"id": deployment_id}, - } - - -def _request_logging(call_type: str, *, stream: bool) -> object: - from litellm.litellm_core_utils.litellm_logging import Logging - - return Logging( - model="g", - messages=MESSAGES, - stream=stream, - call_type=call_type, - start_time=datetime.now(), # noqa: DTZ005 # litellm's own naive start time - litellm_call_id="call-1", - function_id="f", - ) - - -async def _recorded() -> None: - """Let the background recordings of a stream's outcome run.""" - for _ in range(50): - await asyncio.sleep(0.02) - - -@pytest.mark.asyncio -async def test_a_responses_stream_cut_after_it_opened_counts_one_fail_with_a_shared_logging_object( - upstream: str, -) -> None: - router: Final = Router( - model_list=[_served("g", "d1", "openai/gpt-4o", f"{upstream}/cut/v1")], num_retries=0, disable_cooldowns=True - ) - - stream: Final = await router.aresponses( - model="g", input="hi", stream=True, litellm_logging_obj=_request_logging("aresponses", stream=True) - ) - with pytest.raises(Exception): # noqa: B017, PT011 # the transport's read error, whatever class maps it - async for _ in stream: # pyright: ignore[reportGeneralTypeIssues] # a stream when stream=True - pass - await _recorded() - - assert router.health().store.get_cache(key="d1:fails") == 1 - - -@pytest.mark.asyncio -async def test_a_blocked_response_raised_for_a_fallback_counts_no_success_with_a_shared_logging_object( - upstream: str, -) -> None: - router: Final = Router( - model_list=[ - _served("a", "a1", "openai/gpt-4o", f"{upstream}/filtered/v1"), - _served("safe", "s1", "openai/gpt-4o", f"{upstream}/ok/v1"), - ], - content_policy_fallbacks=[{"a": ["safe"]}], - num_retries=0, - disable_cooldowns=True, - ) - - response: Final = await router.acompletion( - model="a", messages=MESSAGES, litellm_logging_obj=_request_logging("acompletion", stream=False) - ) - await _recorded() - - assert response.choices[0].message.content == "ok" - assert router.health().store.get_cache(key="a1:successes") is None - assert router.health().store.get_cache(key="s1:successes") == 1 - - -@pytest.mark.asyncio -async def test_a_text_completion_stream_that_fails_before_its_first_chunk_counts_nothing(upstream: str) -> None: - router: Final = Router( - model_list=[_served("g", "d1", "openai/gpt-3.5-turbo-instruct", f"{upstream}/down/v1")], - num_retries=0, - disable_cooldowns=True, - ) - - stream: Final = await router.atext_completion(model="g", prompt="hi", stream=True) - with pytest.raises(litellm.BadGatewayError): - async for _ in stream: # pyright: ignore[reportGeneralTypeIssues] # a stream when stream=True - pass - await _recorded() - - assert router.health().store.get_cache(key="d1:successes") is None - assert router.health().store.get_cache(key="d1:fails") is None - - @pytest.mark.asyncio async def test_a_text_completion_without_a_prompt_raises_type_error_before_routing() -> None: router: Final = Router(model_list=[_deployment("g", "ok")], num_retries=0) @@ -389,20 +297,6 @@ async def test_a_refused_pick_by_deployment_id_names_the_deployments_group() -> await router.async_get_available_deployment(model="p2", messages=MESSAGES, request_kwargs={}) -@pytest.mark.asyncio -async def test_a_request_refused_before_the_provider_counts_against_no_deployment(upstream: str) -> None: - router: Final = Router( - model_list=[_served("g", "d1", "openai/gpt-4o", f"{upstream}/ok/v1")], num_retries=0, disable_cooldowns=True - ) - - with pytest.raises(litellm.UnsupportedParamsError): - await router.acompletion(model="g", messages=MESSAGES, reasoning_effort="low") - await _recorded() - - assert router.health().store.get_cache(key="d1:fails") is None - assert not [key for key in router.health().store.in_memory_cache.cache_dict if ":rpm" in key] - - @pytest.mark.asyncio async def test_the_async_pick_read_on_the_class_is_a_function_taking_the_router() -> None: router: Final = Router(model_list=[_deployment("g", "ok")]) diff --git a/tests/unit/test_router/test_router_host.py b/tests/unit/test_router/test_router_host.py index 1da96e71d7b..7be9d6b3767 100644 --- a/tests/unit/test_router/test_router_host.py +++ b/tests/unit/test_router/test_router_host.py @@ -1,4 +1,3 @@ -import subprocess import sys import pytest @@ -7,19 +6,6 @@ import litellm.router_host as router_host_module from litellm.router_host import RouterHost, router_host -def test_router_uses_the_proxy_host_without_importing_the_proxy_first(): - # A fresh interpreter, so no earlier test has imported litellm.proxy. - script = ( - "import sys\n" - "from litellm.router import Router\n" - "from litellm.router_host import router_host\n" - "assert 'litellm.proxy' not in sys.modules, 'importing the router imported the proxy'\n" - "print(type(router_host()).__name__)\n" - ) - result = subprocess.run([sys.executable, "-P", "-c", script], capture_output=True, text=True, check=True) - assert result.stdout.strip() == "ProxyRouterHost" - - @pytest.fixture(autouse=True) def nothing_installed(): saved = router_host_module._installed[0] diff --git a/tests/windows_tests/AGENTS.md b/tests/windows_tests/AGENTS.md index 565dba91ce6..36e16a99b5e 100644 --- a/tests/windows_tests/AGENTS.md +++ b/tests/windows_tests/AGENTS.md @@ -1,6 +1,6 @@ # Windows tests -The `windows-sdk` job in `.github/workflows/test-merge-smoke.yml` is the only place these run, on a Windows runner with no proxy dependencies installed. They prove the SDK works on Windows without them, so a test here must not import uvicorn, fastapi, psutil or anything else outside the `dev` group, and must not need the `litellm-gateway.exe` binary +The `windows-sdk` job in `.github/workflows/test-merge-smoke.yml` is the only place these run, on pushes to main and manual runs but not on pull requests, on a Windows runner with no proxy dependencies installed. They prove the SDK works on Windows without them, so a test here must not import uvicorn, fastapi, psutil or anything else outside the `dev` group, and must not need the `litellm-gateway.exe` binary A test that cannot run off Windows carries a `skipif` on `sys.platform`, so collecting this folder on another OS never fails