From 22c9e383f3131e18eb9cf29ca7521d1cbac62bae Mon Sep 17 00:00:00 2001 From: VimarshDwivedi <167737761+VimarshDwivedi@users.noreply.github.com> Date: Tue, 13 Oct 2026 22:45:22 +0530 Subject: [PATCH] fix(mcp): let tool-argument validation calls set their own timeout budget (#46184) test_cold_concurrent_schema_validation_accepts_valid_arguments pins the process to a single CPU core and fans out 8 concurrent calls through anyio.to_process's subprocess worker pool, each needing to cold-spawn a fresh interpreter and import jsonschema. Under a slow or contended CI runner that routinely exceeds the production 30-second budget before any actual validation work runs, failing the test on a genuine timeout rather than a validation bug. _tool_argument_validation_error now takes an optional timeout_seconds, defaulting to the existing constant, so production behavior is unchanged. The test passes a timeout sized for its own single-core, cold-start stress scenario, matching the margin the companion integration test (test_schema_validation_stops_expensive_work_and_recovers) already grants a similar concurrency=8 case. --- litellm/proxy/_experimental/mcp_server/tool_search.py | 6 ++++-- .../proxy/_experimental/mcp_server/test_mcp_tool_search.py | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index a8aff2b2e4b..23a63c00bd0 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -102,7 +102,9 @@ def _validate_tool_arguments(schema: Mapping[str, JsonValue], arguments: Mapping async def _tool_argument_validation_error( - schema: Mapping[str, JsonValue], arguments: Mapping[str, JsonValue] + schema: Mapping[str, JsonValue], + arguments: Mapping[str, JsonValue], + timeout_seconds: float = _VALIDATION_TIMEOUT_SECONDS, ) -> str | None: limit_error: Final = _validation_limit_error(schema, arguments) if limit_error is not None: @@ -112,7 +114,7 @@ async def _tool_argument_validation_error( if existing is None: _SCHEMA_VALIDATION_LIMITER.set(limiter) try: - with anyio.fail_after(_VALIDATION_TIMEOUT_SECONDS): + with anyio.fail_after(timeout_seconds): return await to_process.run_sync( _validate_tool_arguments, schema, arguments, cancellable=True, limiter=limiter ) 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 bd4962829c2..0d6b83b3275 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 @@ -1406,7 +1406,7 @@ async def test_cold_concurrent_schema_validation_accepts_valid_arguments() -> No os.sched_setaffinity(0, {min(original_affinity)}) for _ in range(2): results: Final = await asyncio.gather( - *(_tool_argument_validation_error(schema, {"amount": 0.75}) for _ in range(8)) + *(_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" finally: