From 1c5b28a25400a0ff43b6267483b88e012c628687 Mon Sep 17 00:00:00 2001 From: shin-bot-litellm Date: Sun, 1 Feb 2026 04:27:38 +0000 Subject: [PATCH] fix(tests): fix flaky MCP streaming test patch context - Move stream consumption inside patch context in test_acompletion_with_mcp_streaming_metadata_in_correct_chunks The test was consuming the stream outside the patch context, causing follow-up acompletion calls to hit the real API instead of mocks - Handle missing a2a dependency gracefully in card_resolver.py Previously would raise confusing 'TypeError: NoneType takes no arguments' Now provides a placeholder class with clear ImportError message - Add pytest.importorskip for a2a tests to skip when dependency unavailable --- litellm/a2a_protocol/card_resolver.py | 139 ++++++++++-------- .../a2a_protocol/test_cost_calculator.py | 3 + .../mcp/test_chat_completions_handler.py | 17 ++- 3 files changed, 91 insertions(+), 68 deletions(-) diff --git a/litellm/a2a_protocol/card_resolver.py b/litellm/a2a_protocol/card_resolver.py index 7c4c5af149d..8df869716de 100644 --- a/litellm/a2a_protocol/card_resolver.py +++ b/litellm/a2a_protocol/card_resolver.py @@ -13,6 +13,7 @@ if TYPE_CHECKING: # Runtime imports with availability check _A2ACardResolver: Any = None +_A2A_AVAILABLE: bool = False AGENT_CARD_WELL_KNOWN_PATH: str = "/.well-known/agent-card.json" PREV_AGENT_CARD_WELL_KNOWN_PATH: str = "/.well-known/agent.json" @@ -22,76 +23,90 @@ try: AGENT_CARD_WELL_KNOWN_PATH, PREV_AGENT_CARD_WELL_KNOWN_PATH, ) + _A2A_AVAILABLE = True except ImportError: pass -class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc] - """ - Custom A2A card resolver that supports multiple well-known paths. - - Extends the base A2ACardResolver to try both: - - /.well-known/agent-card.json (standard) - - /.well-known/agent.json (previous/alternative) - """ - - async def get_agent_card( - self, - relative_card_path: Optional[str] = None, - http_kwargs: Optional[Dict[str, Any]] = None, - ) -> "AgentCard": +# Only define the class if a2a is available, otherwise provide a placeholder +if _A2A_AVAILABLE and _A2ACardResolver is not None: + class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc] """ - Fetch the agent card, trying multiple well-known paths. + Custom A2A card resolver that supports multiple well-known paths. - First tries the standard path, then falls back to the previous path. - - Args: - relative_card_path: Optional path to the agent card endpoint. - If None, tries both well-known paths. - http_kwargs: Optional dictionary of keyword arguments to pass to httpx.get - - Returns: - AgentCard from the A2A agent - - Raises: - A2AClientHTTPError or A2AClientJSONError if both paths fail + Extends the base A2ACardResolver to try both: + - /.well-known/agent-card.json (standard) + - /.well-known/agent.json (previous/alternative) """ - # If a specific path is provided, use the parent implementation - if relative_card_path is not None: - return await super().get_agent_card( - relative_card_path=relative_card_path, - http_kwargs=http_kwargs, - ) - # Try both well-known paths - paths = [ - AGENT_CARD_WELL_KNOWN_PATH, - PREV_AGENT_CARD_WELL_KNOWN_PATH, - ] - - last_error = None - for path in paths: - try: - verbose_logger.debug( - f"Attempting to fetch agent card from {self.base_url}{path}" - ) + async def get_agent_card( + self, + relative_card_path: Optional[str] = None, + http_kwargs: Optional[Dict[str, Any]] = None, + ) -> "AgentCard": + """ + Fetch the agent card, trying multiple well-known paths. + + First tries the standard path, then falls back to the previous path. + + Args: + relative_card_path: Optional path to the agent card endpoint. + If None, tries both well-known paths. + http_kwargs: Optional dictionary of keyword arguments to pass to httpx.get + + Returns: + AgentCard from the A2A agent + + Raises: + A2AClientHTTPError or A2AClientJSONError if both paths fail + """ + # If a specific path is provided, use the parent implementation + if relative_card_path is not None: return await super().get_agent_card( - relative_card_path=path, + relative_card_path=relative_card_path, http_kwargs=http_kwargs, ) - except Exception as e: - verbose_logger.debug( - f"Failed to fetch agent card from {self.base_url}{path}: {e}" - ) - last_error = e - continue - - # If we get here, all paths failed - re-raise the last error - if last_error is not None: - raise last_error - - # This shouldn't happen, but just in case - raise Exception( - f"Failed to fetch agent card from {self.base_url}. " - f"Tried paths: {', '.join(paths)}" - ) + + # Try both well-known paths + paths = [ + AGENT_CARD_WELL_KNOWN_PATH, + PREV_AGENT_CARD_WELL_KNOWN_PATH, + ] + + last_error = None + for path in paths: + try: + verbose_logger.debug( + f"Attempting to fetch agent card from {self.base_url}{path}" + ) + return await super().get_agent_card( + relative_card_path=path, + http_kwargs=http_kwargs, + ) + except Exception as e: + verbose_logger.debug( + f"Failed to fetch agent card from {self.base_url}{path}: {e}" + ) + last_error = e + continue + + # If we get here, all paths failed - re-raise the last error + if last_error is not None: + raise last_error + + # This shouldn't happen, but just in case + raise Exception( + f"Failed to fetch agent card from {self.base_url}. " + f"Tried paths: {', '.join(paths)}" + ) +else: + # Provide a placeholder class that raises ImportError when instantiated + class LiteLLMA2ACardResolver: # type: ignore[no-redef] + """ + Placeholder class when a2a SDK is not installed. + """ + def __init__(self, *args, **kwargs): + raise ImportError( + "The 'a2a' package is required to use LiteLLMA2ACardResolver. " + "Install it with: pip install a2a-sdk" + ) diff --git a/tests/test_litellm/a2a_protocol/test_cost_calculator.py b/tests/test_litellm/a2a_protocol/test_cost_calculator.py index 0a472c089b1..368a3004bca 100644 --- a/tests/test_litellm/a2a_protocol/test_cost_calculator.py +++ b/tests/test_litellm/a2a_protocol/test_cost_calculator.py @@ -8,6 +8,9 @@ from unittest.mock import AsyncMock, MagicMock import pytest +# Skip entire module if a2a is not available +a2a = pytest.importorskip("a2a", reason="a2a SDK not installed") + import litellm from litellm.integrations.custom_logger import CustomLogger diff --git a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py index 3cca61092ba..b661e588dbe 100644 --- a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py +++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py @@ -751,6 +751,8 @@ async def test_acompletion_with_mcp_streaming_metadata_in_correct_chunks(monkeyp ) # Patch litellm.acompletion at module level to catch function-level imports + # NOTE: The stream consumption must be inside the patch context because + # the _MCPAutoExecStreamWrapper makes follow-up acompletion calls when consuming with patch("litellm.acompletion", mock_acompletion_func), \ patch.object(chat_completions_handler, "litellm_acompletion", side_effect=mock_acompletion, create=True): result = await acompletion_with_mcp( @@ -760,13 +762,16 @@ async def test_acompletion_with_mcp_streaming_metadata_in_correct_chunks(monkeyp stream=True, ) - # Verify result is CustomStreamWrapper - assert isinstance(result, CustomStreamWrapper) + # Verify result is CustomStreamWrapper + assert isinstance(result, CustomStreamWrapper) + + # Consume the stream and verify metadata placement + # This MUST be inside the patch context because consuming the stream + # triggers follow-up acompletion calls + all_chunks = [] + async for chunk in result: + all_chunks.append(chunk) - # Consume the stream and verify metadata placement - all_chunks = [] - async for chunk in result: - all_chunks.append(chunk) assert len(all_chunks) > 0 # Find first chunk and final chunk from initial response