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: