diff --git a/litellm/a2a_protocol/card_resolver.py b/litellm/a2a_protocol/card_resolver.py new file mode 100644 index 00000000000..7c4c5af149d --- /dev/null +++ b/litellm/a2a_protocol/card_resolver.py @@ -0,0 +1,97 @@ +""" +Custom A2A Card Resolver for LiteLLM. + +Extends the A2A SDK's card resolver to support multiple well-known paths. +""" + +from typing import TYPE_CHECKING, Any, Dict, Optional + +from litellm._logging import verbose_logger + +if TYPE_CHECKING: + from a2a.types import AgentCard + +# Runtime imports with availability check +_A2ACardResolver: Any = None +AGENT_CARD_WELL_KNOWN_PATH: str = "/.well-known/agent-card.json" +PREV_AGENT_CARD_WELL_KNOWN_PATH: str = "/.well-known/agent.json" + +try: + from a2a.client import A2ACardResolver as _A2ACardResolver # type: ignore[no-redef] + from a2a.utils.constants import ( # type: ignore[no-redef] + AGENT_CARD_WELL_KNOWN_PATH, + PREV_AGENT_CARD_WELL_KNOWN_PATH, + ) +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": + """ + 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=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}" + ) + 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)}" + ) diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index ae95faede22..b326f9e7ed5 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -6,6 +6,7 @@ import asyncio import datetime +import uuid from typing import TYPE_CHECKING, Any, AsyncIterator, Coroutine, Dict, Optional, Union import litellm @@ -20,7 +21,6 @@ ) from litellm.types.agents import LiteLLMSendMessageResponse from litellm.utils import client -import uuid if TYPE_CHECKING: from a2a.client import A2AClient as A2AClientType @@ -36,13 +36,18 @@ _A2AClient: Any = None try: - from a2a.client import A2ACardResolver # type: ignore[no-redef] from a2a.client import A2AClient as _A2AClient # type: ignore[no-redef] A2A_SDK_AVAILABLE = True except ImportError: pass +# Import our custom card resolver that supports multiple well-known paths +from litellm.a2a_protocol.card_resolver import LiteLLMA2ACardResolver + +# Use our custom resolver instead of the default A2A SDK resolver +A2ACardResolver = LiteLLMA2ACardResolver + def _set_usage_on_logging_obj( kwargs: Dict[str, Any], diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index a21f3291d31..24727aacd75 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -119,6 +119,11 @@ async def stream_response(): tags=["[beta] A2A Agents"], dependencies=[Depends(user_api_key_auth)], ) +@router.get( + "/a2a/{agent_id}/.well-known/agent.json", + tags=["[beta] A2A Agents"], + dependencies=[Depends(user_api_key_auth)], +) async def get_agent_card( agent_id: str, request: Request, @@ -127,6 +132,10 @@ async def get_agent_card( """ Get the agent card for an agent (A2A discovery endpoint). + Supports both standard paths: + - /.well-known/agent-card.json + - /.well-known/agent.json + The URL in the agent card is rewritten to point to the LiteLLM proxy, so all subsequent A2A calls go through LiteLLM for logging and cost tracking. """ diff --git a/tests/test_litellm/a2a_protocol/test_card_resolver.py b/tests/test_litellm/a2a_protocol/test_card_resolver.py new file mode 100644 index 00000000000..a1bd33107b6 --- /dev/null +++ b/tests/test_litellm/a2a_protocol/test_card_resolver.py @@ -0,0 +1,69 @@ +""" +Mock tests for LiteLLMA2ACardResolver. + +Tests that the card resolver tries both old and new well-known paths. +""" + +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +@pytest.mark.asyncio +async def test_card_resolver_fallback_from_new_to_old_path(): + """ + Test that the card resolver tries the new path (/.well-known/agent-card.json) first, + and falls back to the old path (/.well-known/agent.json) if the new path fails. + """ + # Mock the AgentCard + mock_agent_card = MagicMock() + mock_agent_card.name = "Test Agent" + mock_agent_card.description = "A test agent" + + # Track which paths were called + paths_called = [] + + # Create a mock base class + class MockA2ACardResolver: + def __init__(self, base_url): + self.base_url = base_url + + async def get_agent_card(self, relative_card_path=None, http_kwargs=None): + paths_called.append(relative_card_path) + if relative_card_path == "/.well-known/agent-card.json": + # First call (new path) fails + raise Exception("404 Not Found") + else: + # Second call (old path) succeeds + return mock_agent_card + + # Create mock A2A module + mock_a2a_module = MagicMock() + mock_a2a_client = MagicMock() + mock_a2a_constants = MagicMock() + mock_a2a_constants.AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent-card.json" + mock_a2a_constants.PREV_AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent.json" + + with patch.dict( + sys.modules, + { + "a2a": mock_a2a_module, + "a2a.client": MagicMock(A2ACardResolver=MockA2ACardResolver), + "a2a.utils.constants": mock_a2a_constants, + }, + ): + # Import after patching + from litellm.a2a_protocol.card_resolver import LiteLLMA2ACardResolver + + resolver = LiteLLMA2ACardResolver(base_url="http://test-agent:8000") + result = await resolver.get_agent_card() + + # Verify both paths were tried in correct order + assert len(paths_called) == 2 + assert paths_called[0] == "/.well-known/agent-card.json" # New path tried first + assert paths_called[1] == "/.well-known/agent.json" # Old path tried second + + # Verify the result + assert result == mock_agent_card + assert result.name == "Test Agent"