diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 16ab77c2c2d..5eacc404913 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -2,6 +2,7 @@ from __future__ import annotations import hashlib import json +import os import re from collections import deque from collections.abc import Iterator, Mapping, Sequence @@ -55,6 +56,7 @@ _MAX_VALIDATION_DEPTH: Final = 64 _MAX_VALIDATION_NODES: Final = 10_000 _MAX_VALIDATION_CHARACTERS: Final = 1_048_576 _VALIDATION_TIMEOUT_SECONDS: Final = 30 +_MAX_VALIDATION_WORKERS: Final = 4 def _validation_nodes(value: JsonValue | Mapping[str, JsonValue]) -> Iterator[tuple[int, int]]: @@ -101,6 +103,11 @@ def _validate_tool_arguments(schema: Mapping[str, JsonValue], arguments: Mapping return None +def _schema_validation_worker_count() -> int: + available: Final = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count() or 1 + return max(1, min(_MAX_VALIDATION_WORKERS, available)) + + async def _tool_argument_validation_error( schema: Mapping[str, JsonValue], arguments: Mapping[str, JsonValue], @@ -110,7 +117,7 @@ async def _tool_argument_validation_error( if limit_error is not None: return limit_error existing: Final = _SCHEMA_VALIDATION_LIMITER.get(None) - limiter: Final = existing if existing is not None else anyio.CapacityLimiter(4) + limiter: Final = existing if existing is not None else anyio.CapacityLimiter(_schema_validation_worker_count()) if existing is None: _SCHEMA_VALIDATION_LIMITER.set(limiter) try: diff --git a/tests/unit/caching/test_request_redis_batch_pre_call.py b/tests/unit/caching/test_request_redis_batch_pre_call.py index 69ef67e595b..b1da2a413b6 100644 --- a/tests/unit/caching/test_request_redis_batch_pre_call.py +++ b/tests/unit/caching/test_request_redis_batch_pre_call.py @@ -6,6 +6,7 @@ from __future__ import annotations import asyncio import hashlib import json +from datetime import datetime, timezone from itertools import chain from typing import Any, Final from unittest.mock import AsyncMock, MagicMock @@ -837,7 +838,12 @@ async def test_the_prefetch_reads_cooldowns_again_once_the_read_interval_elapses @pytest.mark.asyncio async def test_the_prefetch_mget_carries_only_the_keys_whose_read_is_due(monkeypatch): first_time: Final = 1_000_000.0 + pinned: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc) monkeypatch.setattr(dual_cache_module.time, "time", lambda: first_time) + monkeypatch.setattr( + "litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime", + lambda: pinned, + ) client = FakeClient(_lua_ok_replies) redis_cache = FakeRedisCache(client) router = _router(redis_cache, routing_strategy="usage-based-routing-v2") diff --git a/tests/unit/integrations/generic_api/test_generic_api_callback.py b/tests/unit/integrations/generic_api/test_generic_api_callback.py index 7a712786a84..afaae02321d 100644 --- a/tests/unit/integrations/generic_api/test_generic_api_callback.py +++ b/tests/unit/integrations/generic_api/test_generic_api_callback.py @@ -443,6 +443,42 @@ async def test_generic_api_callback_sumologic_uses_ndjson(): assert len(mine) == 2, f"Expected this test's 2 calls as NDJSON lines, got {len(mine)} of {len(records)}" +@pytest.mark.asyncio +async def test_generic_api_callback_sends_queued_events_on_its_own(monkeypatch: pytest.MonkeyPatch): + mock_post = AsyncMock() + mock_post.return_value.status_code = 200 + mock_post.return_value.text = "OK" + + test_endpoint = "https://example.com/api/logs" + test_headers = {"Authorization": "Bearer test_token"} + monkeypatch.setenv("GENERIC_LOGGER_ENDPOINT", test_endpoint) + + generic_logger = GenericAPILogger(endpoint=test_endpoint, headers=test_headers, flush_interval=0.05) + generic_logger.async_httpx_client.post = mock_post + monkeypatch.setattr(litellm, "callbacks", [generic_logger]) + + for i in range(3): + await litellm.acompletion( + model="gpt-5.5", + messages=[{"role": "user", "content": f"Hello, world! {i}"}], + mock_response="hi", + user="test_user", + ) + + async def _poll() -> None: + while sum(len(json.loads(call.kwargs["data"])) for call in mock_post.call_args_list) < 3: + await asyncio.sleep(0.01) + + await asyncio.wait_for(_poll(), timeout=10) + + posted_bodies: Final = [json.loads(call.kwargs["data"]) for call in mock_post.call_args_list] + for body in posted_bodies: + assert isinstance(body, list), "Each posted body should be a JSON array" + posted_items: Final = sum(posted_bodies, []) + contents: Final = sorted(item["messages"][0]["content"] for item in posted_items) + assert contents == [f"Hello, world! {i}" for i in range(3)] + + @pytest.mark.asyncio async def test_generic_api_callback_invalid_log_format(): """ diff --git a/tests/unit/integrations/langfuse/test_langfuse_sdk.py b/tests/unit/integrations/langfuse/test_langfuse_sdk.py index cf13667608b..cef2b13f0fd 100644 --- a/tests/unit/integrations/langfuse/test_langfuse_sdk.py +++ b/tests/unit/integrations/langfuse/test_langfuse_sdk.py @@ -1080,11 +1080,19 @@ def test_auth_check_and_project_id_make_one_round_trip_when_langfuse_is_down( [(500, 2), (503, 2), (429, 1), (404, 1)], ids=["http-500", "http-503", "http-429", "http-404"], ) -def test_cold_prompt_miss_never_sleeps_when_langfuse_is_down(status: int, round_trips: int): +def test_cold_prompt_miss_never_sleeps_when_langfuse_is_down( + status: int, round_trips: int, monkeypatch: pytest.MonkeyPatch +) -> None: """A cold ``get_prompt`` fetches inline on the event loop; with the generated client's default retries a 429 carrying ``Retry-After: 30`` used to hold the loop for a minute. A 5xx gets the v2 client's one quick retry, a 429 or 4xx none.""" requests: list[httpx.Request] = [] + sleeps: Final[list[float]] = [] + monkeypatch.setattr( + langfuse_http_client, + "time", + SimpleNamespace(sleep=sleeps.append, time=time.time), + ) def fail(request: httpx.Request) -> httpx.Response: requests.append(request) @@ -1097,18 +1105,25 @@ def test_cold_prompt_miss_never_sleeps_when_langfuse_is_down(status: int, round_ httpx_client=httpx.Client(transport=httpx.MockTransport(fail)), ) - started = monotonic() with pytest.raises(LangfusePromptError) as caught: client.get_prompt("greeting") assert len(requests) == round_trips - assert monotonic() - started < 0.5 + assert sleeps == [], "a cold prompt miss should not sleep for REST retries" assert caught.value.status_code == status @pytest.mark.parametrize("first_failure", [503, "connect-error"], ids=["http-503", "connect-error"]) -def test_one_transient_failure_on_a_cold_prompt_miss_does_not_fail_the_call(first_failure: int | str): +def test_one_transient_failure_on_a_cold_prompt_miss_does_not_fail_the_call( + first_failure: int | str, monkeypatch: pytest.MonkeyPatch +) -> None: """The v2 client retried a cold fetch once; a single Langfuse blip must not fail the LLM call.""" requests: list[httpx.Request] = [] + sleeps: Final[list[float]] = [] + monkeypatch.setattr( + langfuse_http_client, + "time", + SimpleNamespace(sleep=sleeps.append, time=time.time), + ) def flaky(request: httpx.Request) -> httpx.Response: requests.append(request) @@ -1125,10 +1140,9 @@ def test_one_transient_failure_on_a_cold_prompt_miss_does_not_fail_the_call(firs httpx_client=httpx.Client(transport=httpx.MockTransport(flaky)), ) - started = monotonic() assert client.get_prompt("greeting").compile() == "hello" assert len(requests) == 2 - assert monotonic() - started < 0.5 + assert sleeps == [], "one transient failure should not sleep for REST retries" assert client.get_prompt("greeting").compile() == "hello", "the retried prompt is cached like any other" assert len(requests) == 2 diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py index ab6ab882095..c109d3fdbd0 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -4121,6 +4121,18 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock( ) mcp_server._stateful_session_owners[session_id] = mcp_server._owner_fingerprint_for(owner_auth) + class _RecordingLock(asyncio.Lock): + def __init__(self) -> None: + super().__init__() + self.contended = asyncio.Event() + + async def acquire(self) -> bool: + if self.locked(): + self.contended.set() + return await super().acquire() + + session_lock: Final = _RecordingLock() + mcp_server._stateful_session_locks[session_id] = session_lock gate = asyncio.Event() request_in_handle = asyncio.Event() response_handled = asyncio.Event() @@ -4186,12 +4198,19 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock( ), ): req_task = asyncio.create_task(call(request_body)) - await asyncio.wait_for(request_in_handle.wait(), timeout=1.0) + await request_in_handle.wait() resp_task = asyncio.create_task(call(response_body)) - # Under a flat substring heuristic the response would acquire the - # lock held by req_task and this wait would time out (deadlock). - await asyncio.wait_for(response_handled.wait(), timeout=1.0) + response_handled_task: Final = asyncio.create_task(response_handled.wait()) + contended_task: Final = asyncio.create_task(session_lock.contended.wait()) + _, pending = await asyncio.wait( + (response_handled_task, contended_task), + return_when=asyncio.FIRST_COMPLETED, + ) + for event_task in pending: + event_task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + assert response_handled.is_set(), "response POST contended for the request's session lock" await resp_task assert not req_task.done() @@ -4207,6 +4226,7 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock( finally: gate.set() mcp_server._stateful_session_auth_contexts.pop(session_id, None) + mcp_server._stateful_session_auth_context_last_seen.pop(session_id, None) mcp_server._stateful_session_owners.pop(session_id, None) mcp_server._stateful_session_locks.pop(session_id, None) mcp_server._stateful_session_active_request_counts.pop(session_id, None) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_tool_search.py index fd94fccf8fd..9cb3e230aa6 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -1427,7 +1427,12 @@ async def test_mcp_tool_search_leaves_the_listed_tools_slot_empty(monkeypatch: p @pytest.mark.asyncio async def test_cold_concurrent_schema_validation_accepts_valid_arguments() -> None: - from litellm.proxy._experimental.mcp_server.tool_search import _tool_argument_validation_error + import psutil + + from litellm.proxy._experimental.mcp_server.tool_search import ( + _MAX_VALIDATION_WORKERS, + _tool_argument_validation_error, + ) schema: Final = { "type": "object", @@ -1435,6 +1440,7 @@ async def test_cold_concurrent_schema_validation_accepts_valid_arguments() -> No "properties": {"amount": {"$ref": "#/$defs/amount"}}, } original_affinity: Final = os.sched_getaffinity(0) if sys.platform == "linux" else None + existing_children: Final = frozenset(child.pid for child in psutil.Process().children()) try: if original_affinity is not None: os.sched_setaffinity(0, {min(original_affinity)}) @@ -1443,6 +1449,14 @@ async def test_cold_concurrent_schema_validation_accepts_valid_arguments() -> No *(_tool_argument_validation_error(schema, {"amount": 0.75}, timeout_seconds=90.0) for _ in range(8)) ) assert results == [None] * 8, "Cold and warm workers must accept valid concurrent tool arguments" + workers: Final = tuple( + child + for child in psutil.Process().children() + if child.pid not in existing_children and "anyio.to_process" in child.cmdline() + ) + assert 1 <= len(workers) <= (1 if original_affinity is not None else _MAX_VALIDATION_WORKERS), ( + "Validation must not start more cold workers than the CPUs it can run on" + ) finally: if original_affinity is not None: os.sched_setaffinity(0, original_affinity) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_operations.py b/tests/unit/proxy/_experimental/mcp_server/test_operations.py index af547ebf6c7..b037d3b5d25 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_operations.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_operations.py @@ -1287,17 +1287,17 @@ async def test_discovery_concurrent_listings_drain_on_failure_and_cancellation(o ): task = asyncio.create_task(GatewayOperations().execute(DiscoverRequest(), prepare_context(UserAPIKeyAuth(user_id="scoped")))) try: - await asyncio.wait_for(asyncio.gather(*(event.wait() for event in ready)), 1) + await asyncio.gather(*(event.wait() for event in ready)) if outcome == "cancel": task.cancel() else: release.set() if outcome == "success": - result = await asyncio.wait_for(task, 1) + result = await task assert result.capabilities.model_dump(exclude_none=True) == {} else: with pytest.raises(asyncio.CancelledError if outcome == "cancel" else ValueError): - await asyncio.wait_for(task, 1) + await task assert all(event.is_set() for event in closed) assert tools.call_args.kwargs["log_list_tools_to_spendlogs"] is False finally: diff --git a/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py b/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py index 1ad320f5f19..fe4d3c9bc4b 100644 --- a/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py +++ b/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py @@ -499,12 +499,14 @@ async def test_native_count_finishing_after_quarter_worker_budget_keeps_plan_and from litellm.litellm_core_utils import logging_worker from litellm.litellm_core_utils.logging_worker import LoggingWorker - release: Final = asyncio.Event() + counted: Final = asyncio.Event() + loop: Final = asyncio.get_running_loop() + real_time: Final = loop.time + monkeypatch.setattr(loop, "time", lambda: real_time() + (2.05 if counted.is_set() else 0.0)) async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int: - if not release.is_set(): - asyncio.get_running_loop().call_later(2.05, release.set) - await release.wait() + counted.set() + await asyncio.sleep(0.001) return await _count(model, api_key, body) worker: Final = LoggingWorker(timeout=8.0) @@ -519,7 +521,6 @@ async def test_native_count_finishing_after_quarter_worker_budget_keeps_plan_and assert observed.plan is not None and observed.plan.breakpoints[0].prefix_tokens == 5000 assert payload["response_cost"] is not None and worker._timeout_total == 0 finally: - release.set() await worker.stop() diff --git a/tests/unit/proxy/roi_calculator/test_gitlab.py b/tests/unit/proxy/roi_calculator/test_gitlab.py index e3b237a1d41..38ff412b049 100644 --- a/tests/unit/proxy/roi_calculator/test_gitlab.py +++ b/tests/unit/proxy/roi_calculator/test_gitlab.py @@ -48,15 +48,15 @@ async def test_fork_lookups_overlap_with_a_bounded_number_of_requests() -> None: source: Final = GitLab(ROISettings(source_provider="gitlab"), httpx.MockTransport(respond)) pending: Final = asyncio.create_task(source.pulls("group/repo", date(2026, 9, 1), date(2026, 9, 30))) try: - first_wave: Final = tuple([await asyncio.wait_for(started.get(), timeout=1) for _ in range(8)]) + first_wave: Final = tuple([await started.get() for _ in range(8)]) assert len(set(first_wave)) == 8 assert started.empty() release[first_wave[0] - 2].set() - next_id: Final = await asyncio.wait_for(started.get(), timeout=1) + next_id: Final = await started.get() assert next_id not in first_wave for event in release: event.set() - pulls: Final = await asyncio.wait_for(pending, timeout=1) + pulls: Final = await pending assert tuple(pull.head.repo.full_name if pull.head and pull.head.repo else None for pull in pulls) == tuple( None if source_id == 3 else f"fork-{source_id}/repo" for source_id in source_ids ) diff --git a/tests/unit/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/unit/proxy/spend_tracking/test_key_metadata_recovery.py index 2a5f03bcdf3..7a0767257dc 100644 --- a/tests/unit/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/unit/proxy/spend_tracking/test_key_metadata_recovery.py @@ -27,6 +27,17 @@ from litellm.proxy.utils import hash_token from litellm.proxy.db.log_db_metrics import record_db_io +class _FakeClock: + def __init__(self) -> None: + self.now = 1000.0 + + def advance(self, seconds: float) -> None: + self.now += seconds + + def __call__(self) -> float: + return self.now + + def _digest_row(digest: str, key_alias: str | None, team_id: str | None, user_id: str | None) -> dict[str, str | None]: return {"digest": digest, "key_alias": key_alias, "team_id": team_id, "user_id": user_id} @@ -422,34 +433,44 @@ async def test_recover_key_metadata_from_spend_logs_rescans_when_the_window_chan @pytest.mark.asyncio -async def test_recover_key_metadata_from_spend_logs_retries_a_failed_query_only_after_the_miss_ttl(): - digest = hash_token("cli-session-retry") - window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) - cache = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL) - mock_prisma = MagicMock() - query_raw = _spend_log_transaction(mock_prisma, AsyncMock(side_effect=PrismaError("statement timeout"))) - started = time.time() - assert await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) == {} - query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(digest, "back-online", None, None)])) +async def test_recover_key_metadata_from_spend_logs_retries_a_failed_query_only_after_the_miss_ttl() -> None: + digest: Final = hash_token("cli-session-retry") + window: Final = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + clock: Final = _FakeClock() + cache: Final = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL, clock=clock) + mock_prisma: Final = MagicMock() + query_raw: Final = _spend_log_transaction( + mock_prisma, + AsyncMock( + side_effect=[ + PrismaError("statement timeout"), + [_spend_log_row(digest, "back-online", None, None)], + ] + ), + ) assert await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) == {} - query_raw.assert_not_awaited() - miss_key = next(key for key in cache.ttl_dict if digest in key and not key.endswith(":missed-before")) - assert cache.ttl_dict[miss_key] - started <= SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1 - cache.ttl_dict[miss_key] = time.time() - 1 + miss_key: Final = next(key for key in cache.ttl_dict if digest in key and not key.endswith(":missed-before")) + assert cache.ttl_dict[miss_key] == clock.now + SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL - result = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) + assert await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) == {} + assert query_raw.await_count == 1 + clock.advance(SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1) + + result: Final = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) assert result[digest]["key_alias"] == "back-online" + assert query_raw.await_count == 2 @pytest.mark.asyncio -async def test_recover_key_metadata_from_spend_logs_never_caches_repeated_query_failures_as_long_as_a_hit(): +async def test_recover_key_metadata_from_spend_logs_never_caches_repeated_query_failures_as_long_as_a_hit() -> None: digest: Final = hash_token("cli-session-repeated-failure") window: Final = (datetime(2026, 9, 7), datetime(2026, 9, 10)) - cache: Final = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL) + clock: Final = _FakeClock() + cache: Final = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL, clock=clock) mock_prisma: Final = MagicMock() - _spend_log_transaction( + query_raw: Final = _spend_log_transaction( mock_prisma, AsyncMock( side_effect=[ @@ -461,17 +482,19 @@ async def test_recover_key_metadata_from_spend_logs_never_caches_repeated_query_ ) await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) miss_key: Final = next(key for key in cache.ttl_dict if digest in key and not key.endswith(":missed-before")) - cache.ttl_dict[miss_key] = time.time() - 1 - second_query_started: Final = time.time() + assert cache.ttl_dict[miss_key] == clock.now + SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + clock.advance(SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1) await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) - assert cache.ttl_dict[miss_key] - second_query_started <= SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1 - cache.ttl_dict[miss_key] = time.time() - 1 + assert query_raw.await_count == 2 + assert cache.ttl_dict[miss_key] == clock.now + SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + clock.advance(SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1) result: Final = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) assert result[digest]["key_alias"] == "back-online" + assert query_raw.await_count == 3 @pytest.mark.asyncio diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index f3dbdbcc556..4ac917d706d 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -1071,6 +1071,7 @@ def _shared_redis_stub(store: dict) -> MagicMock: @pytest.mark.asyncio +@pytest.mark.usefixtures("router_minute_pinned") async def test_headers_on_fresh_worker_reflect_shared_redis_usage(): store: dict = {} worker_a = _rpm_tpm_router("lit-3058-workers")