ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Fix] A2a Agent Gateway Fixes - A2A agents deployed with localhost/internal URLs in their agent cards (e.g., http://0.0.0.0:8001/) (#20604)

Browse Source 
* v1 card resolver fix

* fix: is_localhost_or_internal_url

* fix code

* test_fix_agent_card_url_replaces_localhost

* test restruct

* test_a2a_non_streaming

* test agnts

* add exception handling

* init errors

* add localhost retry

* add agent_testing

* test_a2a_non_streaming

* _build_streaming_logging_obj

* code qa fixes

* test_card_resolver_fallback_from_new_to_old_path

* fix linting
This commit is contained in:
Ishaan Jaff 2026-02-06 15:02:34 -08:00 committed by GitHub
parent 6a213fc3bc
commit b78f4c924c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
12 changed files with 773 additions and 82 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

53
.circleci/config.yml
View File

@ -1372,6 +1372,51 @@ jobs:
paths:
- mcp_coverage.xml
- mcp_coverage
agent_testing:
docker:
- image: cimg/python:3.11
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
working_directory: ~/project
steps:
- checkout
- setup_google_dns
- run:
name: Install Dependencies
command: |
python -m pip install --upgrade pip
python -m pip install -r requirements.txt
pip install "pytest==7.3.1"
pip install "pytest-retry==1.6.3"
pip install "pytest-cov==5.0.0"
pip install "pytest-asyncio==0.21.1"
pip install "respx==0.22.0"
pip install "pydantic==2.11.0"
pip install "a2a-sdk"
# Run pytest and generate JUnit XML report
- run:
name: Run tests
command: |
pwd
ls
python -m pytest -vv tests/agent_tests --ignore=tests/agent_tests/local_only_agent_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5
no_output_timeout: 120m
- run:
name: Rename the coverage files
command: |
mv coverage.xml agent_coverage.xml
mv .coverage agent_coverage
# Store test results
- store_test_results:
path: test-results
- persist_to_workspace:
root: .
paths:
- agent_coverage.xml
- agent_coverage
guardrails_testing:
docker:
- image: cimg/python:3.11
@ -4264,6 +4309,12 @@ workflows:
only:
- main
- /litellm_.*/
- agent_testing:
filters:
branches:
only:
- main
- /litellm_.*/
- guardrails_testing:
filters:
branches:
@ -4371,6 +4422,7 @@ workflows:
- llm_translation_testing
- realtime_translation_testing
- mcp_testing
- agent_testing
- google_generate_content_endpoint_testing
- guardrails_testing
- llm_responses_api_testing
@ -4449,6 +4501,7 @@ workflows:
- llm_translation_testing
- realtime_translation_testing
- mcp_testing
- agent_testing
- google_generate_content_endpoint_testing
- llm_responses_api_testing
- ocr_testing

14
litellm/a2a_protocol/__init__.py
View File

@ -39,6 +39,12 @@ Example usage (class-based):
"""
from litellm.a2a_protocol.client import A2AClient
from litellm.a2a_protocol.exceptions import (
A2AAgentCardError,
A2AConnectionError,
A2AError,
A2ALocalhostURLError,
)
from litellm.a2a_protocol.main import (
aget_agent_card,
asend_message,
@ -49,11 +55,19 @@ from litellm.a2a_protocol.main import (
from litellm.types.agents import LiteLLMSendMessageResponse
__all__ = [
# Client
"A2AClient",
# Functions
"asend_message",
"send_message",
"asend_message_streaming",
"aget_agent_card",
"create_a2a_client",
# Response types
"LiteLLMSendMessageResponse",
# Exceptions
"A2AError",
"A2AConnectionError",
"A2AAgentCardError",
"A2ALocalhostURLError",
]

67
litellm/a2a_protocol/card_resolver.py
View File

@ -7,6 +7,7 @@ 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
from litellm.constants import LOCALHOST_URL_PATTERNS
if TYPE_CHECKING:
from a2a.types import AgentCard
@ -26,15 +27,61 @@ except ImportError:
pass
def is_localhost_or_internal_url(url: Optional[str]) -> bool:
"""
Check if a URL is a localhost or internal URL.
This detects common development URLs that are accidentally left in
agent cards when deploying to production.
Args:
url: The URL to check
Returns:
True if the URL is localhost/internal
"""
if not url:
return False
url_lower = url.lower()
return any(pattern in url_lower for pattern in LOCALHOST_URL_PATTERNS)
def fix_agent_card_url(agent_card: "AgentCard", base_url: str) -> "AgentCard":
"""
Fix the agent card URL if it contains a localhost/internal address.
Many A2A agents are deployed with agent cards that contain internal URLs
like "http://0.0.0.0:8001/" or "http://localhost:8000/". This function
replaces such URLs with the provided base_url.
Args:
agent_card: The agent card to fix
base_url: The base URL to use as replacement
Returns:
The agent card with the URL fixed if necessary
"""
card_url = getattr(agent_card, "url", None)
if card_url and is_localhost_or_internal_url(card_url):
# Normalize base_url to ensure it ends with /
fixed_url = base_url.rstrip("/") + "/"
agent_card.url = fixed_url
return agent_card
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,
@ -42,17 +89,17 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc]
) -> "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
"""
@ -62,13 +109,13 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc]
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:
@ -85,11 +132,11 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc]
)
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}. "

203
litellm/a2a_protocol/exception_mapping_utils.py Normal file
View File

@ -0,0 +1,203 @@
"""
A2A Protocol Exception Mapping Utils.
Maps A2A SDK exceptions to LiteLLM A2A exception types.
"""
from typing import TYPE_CHECKING, Any, Optional
from litellm._logging import verbose_logger
from litellm.a2a_protocol.card_resolver import (
fix_agent_card_url,
is_localhost_or_internal_url,
)
from litellm.a2a_protocol.exceptions import (
A2AAgentCardError,
A2AConnectionError,
A2AError,
A2ALocalhostURLError,
)
from litellm.constants import CONNECTION_ERROR_PATTERNS
if TYPE_CHECKING:
from a2a.client import A2AClient as A2AClientType
# Runtime import
A2A_SDK_AVAILABLE = False
try:
from a2a.client import A2AClient as _A2AClient # type: ignore[no-redef]
A2A_SDK_AVAILABLE = True
except ImportError:
_A2AClient = None # type: ignore[misc]
class A2AExceptionCheckers:
"""
Helper class for checking various A2A error conditions.
"""
@staticmethod
def is_connection_error(error_str: str) -> bool:
"""
Check if an error string indicates a connection error.
Args:
error_str: The error string to check
Returns:
True if the error indicates a connection issue
"""
if not isinstance(error_str, str):
return False
error_str_lower = error_str.lower()
return any(pattern in error_str_lower for pattern in CONNECTION_ERROR_PATTERNS)
@staticmethod
def is_localhost_url(url: Optional[str]) -> bool:
"""
Check if a URL is a localhost/internal URL.
Args:
url: The URL to check
Returns:
True if the URL is localhost/internal
"""
return is_localhost_or_internal_url(url)
@staticmethod
def is_agent_card_error(error_str: str) -> bool:
"""
Check if an error string indicates an agent card error.
Args:
error_str: The error string to check
Returns:
True if the error is related to agent card fetching/parsing
"""
if not isinstance(error_str, str):
return False
error_str_lower = error_str.lower()
agent_card_patterns = [
"agent card",
"agent-card",
".well-known",
"card not found",
"invalid agent",
]
return any(pattern in error_str_lower for pattern in agent_card_patterns)
def map_a2a_exception(
original_exception: Exception,
card_url: Optional[str] = None,
api_base: Optional[str] = None,
model: Optional[str] = None,
) -> Exception:
"""
Map an A2A SDK exception to a LiteLLM A2A exception type.
Args:
original_exception: The original exception from the A2A SDK
card_url: The URL from the agent card (if available)
api_base: The original API base URL
model: The model/agent name
Returns:
A mapped LiteLLM A2A exception
Raises:
A2ALocalhostURLError: If the error is a connection error to a localhost URL
A2AConnectionError: If the error is a general connection error
A2AAgentCardError: If the error is related to agent card issues
A2AError: For other A2A-related errors
"""
error_str = str(original_exception)
# Check for localhost URL connection error (special case - retryable)
if (
card_url
and api_base
and A2AExceptionCheckers.is_localhost_url(card_url)
and A2AExceptionCheckers.is_connection_error(error_str)
):
raise A2ALocalhostURLError(
localhost_url=card_url,
base_url=api_base,
original_error=original_exception,
model=model,
)
# Check for agent card errors
if A2AExceptionCheckers.is_agent_card_error(error_str):
raise A2AAgentCardError(
message=error_str,
url=api_base,
model=model,
)
# Check for general connection errors
if A2AExceptionCheckers.is_connection_error(error_str):
raise A2AConnectionError(
message=error_str,
url=card_url or api_base,
model=model,
)
# Default: wrap in generic A2AError
raise A2AError(
message=error_str,
model=model,
)
def handle_a2a_localhost_retry(
error: A2ALocalhostURLError,
agent_card: Any,
a2a_client: "A2AClientType",
is_streaming: bool = False,
) -> "A2AClientType":
"""
Handle A2ALocalhostURLError by fixing the URL and creating a new client.
This is called when we catch an A2ALocalhostURLError and want to retry
with the corrected URL.
Args:
error: The localhost URL error
agent_card: The agent card object to fix
a2a_client: The current A2A client
is_streaming: Whether this is a streaming request (for logging)
Returns:
A new A2A client with the fixed URL
Raises:
ImportError: If the A2A SDK is not installed
"""
if not A2A_SDK_AVAILABLE or _A2AClient is None:
raise ImportError(
"A2A SDK is required for localhost retry handling. "
"Install it with: pip install a2a"
)
request_type = "streaming " if is_streaming else ""
verbose_logger.warning(
f"A2A {request_type}request to '{error.localhost_url}' failed: {error.original_error}. "
f"Agent card contains localhost/internal URL. "
f"Retrying with base_url '{error.base_url}'."
)
# Fix the agent card URL
fix_agent_card_url(agent_card, error.base_url)
# Create a new client with the fixed agent card (transport caches URL)
return _A2AClient(
httpx_client=a2a_client._transport.httpx_client, # type: ignore[union-attr]
agent_card=agent_card,
)

150
litellm/a2a_protocol/exceptions.py Normal file
View File

@ -0,0 +1,150 @@
"""
A2A Protocol Exceptions.
Custom exception types for A2A protocol operations, following LiteLLM's exception pattern.
"""
from typing import Optional
import httpx
class A2AError(Exception):
"""
Base exception for A2A protocol errors.
Follows the same pattern as LiteLLM's main exceptions.
"""
def __init__(
self,
message: str,
status_code: int = 500,
llm_provider: str = "a2a_agent",
model: Optional[str] = None,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
):
self.status_code = status_code
self.message = f"litellm.A2AError: {message}"
self.llm_provider = llm_provider
self.model = model
self.litellm_debug_info = litellm_debug_info
self.max_retries = max_retries
self.num_retries = num_retries
self.response = response or httpx.Response(
status_code=self.status_code,
request=httpx.Request(method="POST", url="https://litellm.ai"),
)
super().__init__(self.message)
def __str__(self) -> str:
_message = self.message
if self.num_retries:
_message += f" LiteLLM Retried: {self.num_retries} times"
if self.max_retries:
_message += f", LiteLLM Max Retries: {self.max_retries}"
return _message
def __repr__(self) -> str:
return self.__str__()
class A2AConnectionError(A2AError):
"""
Raised when connection to an A2A agent fails.
This typically occurs when:
- The agent is unreachable
- The agent card contains a localhost/internal URL
- Network issues prevent connection
"""
def __init__(
self,
message: str,
url: Optional[str] = None,
model: Optional[str] = None,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
):
self.url = url
super().__init__(
message=message,
status_code=503,
llm_provider="a2a_agent",
model=model,
response=response,
litellm_debug_info=litellm_debug_info,
max_retries=max_retries,
num_retries=num_retries,
)
class A2AAgentCardError(A2AError):
"""
Raised when there's an issue with the agent card.
This includes:
- Failed to fetch agent card
- Invalid agent card format
- Missing required fields
"""
def __init__(
self,
message: str,
url: Optional[str] = None,
model: Optional[str] = None,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
):
self.url = url
super().__init__(
message=message,
status_code=404,
llm_provider="a2a_agent",
model=model,
response=response,
litellm_debug_info=litellm_debug_info,
)
class A2ALocalhostURLError(A2AConnectionError):
"""
Raised when an agent card contains a localhost/internal URL.
Many A2A agents are deployed with agent cards that contain internal URLs
like "http://0.0.0.0:8001/" or "http://localhost:8000/". This error
indicates that the URL needs to be corrected and the request should be retried.
Attributes:
localhost_url: The localhost/internal URL found in the agent card
base_url: The public base URL that should be used instead
original_error: The original connection error that was raised
"""
def __init__(
self,
localhost_url: str,
base_url: str,
original_error: Optional[Exception] = None,
model: Optional[str] = None,
):
self.localhost_url = localhost_url
self.base_url = base_url
self.original_error = original_error
message = (
f"Agent card contains localhost/internal URL '{localhost_url}'. "
f"Retrying with base URL '{base_url}'."
)
super().__init__(
message=message,
url=localhost_url,
model=model,
)

184
litellm/a2a_protocol/main.py
View File

@ -44,6 +44,11 @@ except ImportError:
# Import our custom card resolver that supports multiple well-known paths
from litellm.a2a_protocol.card_resolver import LiteLLMA2ACardResolver
from litellm.a2a_protocol.exception_mapping_utils import (
handle_a2a_localhost_retry,
map_a2a_exception,
)
from litellm.a2a_protocol.exceptions import A2ALocalhostURLError
# Use our custom resolver instead of the default A2A SDK resolver
A2ACardResolver = LiteLLMA2ACardResolver
@ -244,10 +249,50 @@ async def asend_message(
verbose_logger.info(f"A2A send_message request_id={request.id}, agent={agent_name}")
a2a_response = await a2a_client.send_message(request)
# Get agent card URL for localhost retry logic
agent_card = getattr(a2a_client, "_litellm_agent_card", None) or getattr(
a2a_client, "agent_card", None
)
card_url = getattr(agent_card, "url", None) if agent_card else None
# Retry loop: if connection fails due to localhost URL in agent card, retry with fixed URL
a2a_response = None
for _ in range(2): # max 2 attempts: original + 1 retry
try:
a2a_response = await a2a_client.send_message(request)
break # success, exit retry loop
except A2ALocalhostURLError as e:
# Localhost URL error - fix and retry
a2a_client = handle_a2a_localhost_retry(
error=e,
agent_card=agent_card,
a2a_client=a2a_client,
is_streaming=False,
)
card_url = agent_card.url if agent_card else None
except Exception as e:
# Map exception - will raise A2ALocalhostURLError if applicable
try:
map_a2a_exception(e, card_url, api_base, model=agent_name)
except A2ALocalhostURLError as localhost_err:
# Localhost URL error - fix and retry
a2a_client = handle_a2a_localhost_retry(
error=localhost_err,
agent_card=agent_card,
a2a_client=a2a_client,
is_streaming=False,
)
card_url = agent_card.url if agent_card else None
continue
except Exception:
# Re-raise the mapped exception
raise
verbose_logger.info(f"A2A send_message completed, request_id={request.id}")
# a2a_response is guaranteed to be set if we reach here (loop breaks on success or raises)
assert a2a_response is not None
# Wrap in LiteLLM response type for _hidden_params support
response = LiteLLMSendMessageResponse.from_a2a_response(a2a_response)
@ -307,6 +352,48 @@ def send_message(
)
def _build_streaming_logging_obj(
request: "SendStreamingMessageRequest",
agent_name: str,
agent_id: Optional[str],
litellm_params: Optional[Dict[str, Any]],
metadata: Optional[Dict[str, Any]],
proxy_server_request: Optional[Dict[str, Any]],
) -> Logging:
"""Build logging object for streaming A2A requests."""
start_time = datetime.datetime.now()
model = f"a2a_agent/{agent_name}"
logging_obj = Logging(
model=model,
messages=[{"role": "user", "content": "streaming-request"}],
stream=False,
call_type="asend_message_streaming",
start_time=start_time,
litellm_call_id=str(request.id),
function_id=str(request.id),
)
logging_obj.model = model
logging_obj.custom_llm_provider = "a2a_agent"
logging_obj.model_call_details["model"] = model
logging_obj.model_call_details["custom_llm_provider"] = "a2a_agent"
if agent_id:
logging_obj.model_call_details["agent_id"] = agent_id
_litellm_params = litellm_params.copy() if litellm_params else {}
if metadata:
_litellm_params["metadata"] = metadata
if proxy_server_request:
_litellm_params["proxy_server_request"] = proxy_server_request
logging_obj.litellm_params = _litellm_params
logging_obj.optional_params = _litellm_params
logging_obj.model_call_details["litellm_params"] = _litellm_params
logging_obj.model_call_details["metadata"] = metadata or {}
return logging_obj
async def asend_message_streaming(
a2a_client: Optional["A2AClientType"] = None,
request: Optional["SendStreamingMessageRequest"] = None,
@ -403,55 +490,72 @@ async def asend_message_streaming(
verbose_logger.info(f"A2A send_message_streaming request_id={request.id}")
# Track for logging
start_time = datetime.datetime.now()
stream = a2a_client.send_message_streaming(request)
# Build logging object for streaming completion callbacks
agent_card = getattr(a2a_client, "_litellm_agent_card", None) or getattr(
a2a_client, "agent_card", None
)
card_url = getattr(agent_card, "url", None) if agent_card else None
agent_name = getattr(agent_card, "name", "unknown") if agent_card else "unknown"
model = f"a2a_agent/{agent_name}"
logging_obj = Logging(
model=model,
messages=[{"role": "user", "content": "streaming-request"}],
stream=False, # complete response logging after stream ends
call_type="asend_message_streaming",
start_time=start_time,
litellm_call_id=str(request.id),
function_id=str(request.id),
)
logging_obj.model = model
logging_obj.custom_llm_provider = "a2a_agent"
logging_obj.model_call_details["model"] = model
logging_obj.model_call_details["custom_llm_provider"] = "a2a_agent"
if agent_id:
logging_obj.model_call_details["agent_id"] = agent_id
# Propagate litellm_params for spend logging (includes cost_per_query, etc.)
_litellm_params = litellm_params.copy() if litellm_params else {}
# Merge metadata into litellm_params.metadata (required for proxy cost tracking)
if metadata:
_litellm_params["metadata"] = metadata
if proxy_server_request:
_litellm_params["proxy_server_request"] = proxy_server_request
logging_obj.litellm_params = _litellm_params
logging_obj.optional_params = _litellm_params # used by cost calc
logging_obj.model_call_details["litellm_params"] = _litellm_params
logging_obj.model_call_details["metadata"] = metadata or {}
iterator = A2AStreamingIterator(
stream=stream,
logging_obj = _build_streaming_logging_obj(
request=request,
logging_obj=logging_obj,
agent_name=agent_name,
agent_id=agent_id,
litellm_params=litellm_params,
metadata=metadata,
proxy_server_request=proxy_server_request,
)
async for chunk in iterator:
yield chunk
# Retry loop: if connection fails due to localhost URL in agent card, retry with fixed URL
# Connection errors in streaming typically occur on first chunk iteration
first_chunk = True
for attempt in range(2): # max 2 attempts: original + 1 retry
stream = a2a_client.send_message_streaming(request)
iterator = A2AStreamingIterator(
stream=stream,
request=request,
logging_obj=logging_obj,
agent_name=agent_name,
)
try:
first_chunk = True
async for chunk in iterator:
if first_chunk:
first_chunk = False # connection succeeded
yield chunk
return # stream completed successfully
except A2ALocalhostURLError as e:
# Only retry on first chunk, not mid-stream
if first_chunk and attempt == 0:
a2a_client = handle_a2a_localhost_retry(
error=e,
agent_card=agent_card,
a2a_client=a2a_client,
is_streaming=True,
)
card_url = agent_card.url if agent_card else None
else:
raise
except Exception as e:
# Only map exception on first chunk
if first_chunk and attempt == 0:
try:
map_a2a_exception(e, card_url, api_base, model=agent_name)
except A2ALocalhostURLError as localhost_err:
# Localhost URL error - fix and retry
a2a_client = handle_a2a_localhost_retry(
error=localhost_err,
agent_card=agent_card,
a2a_client=a2a_client,
is_streaming=True,
)
card_url = agent_card.url if agent_card else None
continue
except Exception:
# Re-raise the mapped exception
raise
raise
async def create_a2a_client(

16
litellm/constants.py
View File

@ -306,6 +306,22 @@ DEFAULT_MAX_TOKENS_FOR_TRITON = int(os.getenv("DEFAULT_MAX_TOKENS_FOR_TRITON", 2
#### Networking settings ####
request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", 6000)) # time in seconds
DEFAULT_A2A_AGENT_TIMEOUT: float = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)) # 10 minutes
# Patterns that indicate a localhost/internal URL in A2A agent cards that should be
# replaced with the original base_url. This is a common misconfiguration where
# developers deploy agents with development URLs in their agent cards.
LOCALHOST_URL_PATTERNS: List[str] = [
"localhost",
"127.0.0.1",
"0.0.0.0",
"[::1]", # IPv6 localhost
]
# Patterns in error messages that indicate a connection failure
CONNECTION_ERROR_PATTERNS: List[str] = [
"connect",
"connection",
"network",
"refused",
]
STREAM_SSE_DONE_STRING: str = "[DONE]"
STREAM_SSE_DATA_PREFIX: str = "data: "
### SPEND TRACKING ###

0
tests/agent_tests/local_vertex_agent.py → tests/agent_tests/local_only_agent_tests/local_vertex_agent.py
View File

0
tests/agent_tests/test_a2a.py → tests/agent_tests/local_only_agent_tests/test_a2a.py
View File

0
tests/agent_tests/test_a2a_completion_bridge.py → tests/agent_tests/local_only_agent_tests/test_a2a_completion_bridge.py
View File

83
tests/agent_tests/test_a2a_agent.py Normal file
View File

@ -0,0 +1,83 @@
"""
Simple A2A agent tests - non-streaming and streaming.
These tests validate the localhost URL retry logic: if an A2A agent's card
contains a localhost/internal URL (e.g., http://0.0.0.0:8001/), the request
will fail with a connection error. LiteLLM detects this and automatically
retries using the original api_base URL instead.
Requires A2A_AGENT_URL environment variable to be set.
Run with:
A2A_AGENT_URL=https://your-agent.example.com pytest tests/agent_tests/test_a2a_agent.py -v -s
"""
import os
import pytest
from uuid import uuid4
def get_a2a_agent_url():
"""Get A2A agent URL from environment, skip test if not set."""
url = os.environ.get("A2A_AGENT_URL")
return url
@pytest.mark.asyncio
async def test_a2a_non_streaming():
"""Test non-streaming A2A request."""
from a2a.types import MessageSendParams, SendMessageRequest
from litellm.a2a_protocol import asend_message
api_base = get_a2a_agent_url()
request = SendMessageRequest(
id=str(uuid4()),
params=MessageSendParams(
message={
"role": "user",
"parts": [{"kind": "text", "text": "Say hello in one word"}],
"messageId": uuid4().hex,
}
),
)
response = await asend_message(
request=request,
api_base=api_base,
)
assert response is not None
print(f"\nNon-streaming response: {response}")
@pytest.mark.asyncio
async def test_a2a_streaming():
"""Test streaming A2A request."""
from a2a.types import MessageSendParams, SendStreamingMessageRequest
from litellm.a2a_protocol import asend_message_streaming
api_base = get_a2a_agent_url()
request = SendStreamingMessageRequest(
id=str(uuid4()),
params=MessageSendParams(
message={
"role": "user",
"parts": [{"kind": "text", "text": "Say hello in one word"}],
"messageId": uuid4().hex,
}
),
)
chunks = []
async for chunk in asend_message_streaming(
request=request,
api_base=api_base,
):
chunks.append(chunk)
print(f"\nStreaming chunk: {chunk}")
assert len(chunks) > 0, "Should receive at least one chunk"
print(f"\nTotal chunks received: {len(chunks)}")

85
tests/test_litellm/a2a_protocol/test_card_resolver.py
View File

@ -4,11 +4,16 @@ 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
from litellm.a2a_protocol.card_resolver import (
LiteLLMA2ACardResolver,
fix_agent_card_url,
is_localhost_or_internal_url,
)
@pytest.mark.asyncio
async def test_card_resolver_fallback_from_new_to_old_path():
@ -24,39 +29,31 @@ async def test_card_resolver_fallback_from_new_to_old_path():
# 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,
},
# Create a mock for the parent's get_agent_card method
async def mock_parent_get_agent_card(
self, relative_card_path=None, http_kwargs=None
):
# Import after patching
from litellm.a2a_protocol.card_resolver import LiteLLMA2ACardResolver
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
resolver = LiteLLMA2ACardResolver(base_url="http://test-agent:8000")
# Create a mock httpx client
mock_httpx_client = MagicMock()
# Patch the parent class's get_agent_card method
# We need to patch the actual parent class method that super() calls
with patch.object(
LiteLLMA2ACardResolver.__bases__[0],
"get_agent_card",
mock_parent_get_agent_card,
):
resolver = LiteLLMA2ACardResolver(
httpx_client=mock_httpx_client, base_url="http://test-agent:8000"
)
result = await resolver.get_agent_card()
# Verify both paths were tried in correct order
@ -67,3 +64,27 @@ async def test_card_resolver_fallback_from_new_to_old_path():
# Verify the result
assert result == mock_agent_card
assert result.name == "Test Agent"
def test_is_localhost_or_internal_url():
"""Test that localhost/internal URLs are correctly detected."""
# Should return True for localhost variants
assert is_localhost_or_internal_url("http://localhost:8000/") is True
assert is_localhost_or_internal_url("http://0.0.0.0:8001/") is True
# Should return False for public URLs
assert is_localhost_or_internal_url("https://my-agent.example.com/") is False
assert is_localhost_or_internal_url(None) is False
def test_fix_agent_card_url_replaces_localhost():
"""Test that fix_agent_card_url replaces localhost URLs with base_url."""
# Create mock agent card with localhost URL
mock_card = MagicMock()
mock_card.url = "http://0.0.0.0:8001/"
# Fix the URL
result = fix_agent_card_url(mock_card, "https://my-public-agent.example.com")
# Verify localhost URL was replaced with base_url
assert result.url == "https://my-public-agent.example.com/"