test agnts

This commit is contained in:
Ishaan Jaffer 2026-02-06 13:50:41 -08:00
parent 5c7f4636b1
commit 3621e512f6
3 changed files with 55 additions and 77 deletions

View file

@ -1,8 +1,7 @@
"""
Custom A2A Card Resolver for LiteLLM.
Extends the A2A SDK's card resolver to support multiple well-known paths
and fixes agent card URL issues where the card contains internal/localhost URLs.
Extends the A2A SDK's card resolver to support multiple well-known paths.
"""
from typing import TYPE_CHECKING, Any, Dict, Optional
@ -30,7 +29,7 @@ except ImportError:
def is_localhost_or_internal_url(url: Optional[str]) -> bool:
"""
Check if a URL is a localhost or internal URL that should be replaced.
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.
@ -39,17 +38,7 @@ def is_localhost_or_internal_url(url: Optional[str]) -> bool:
url: The URL to check
Returns:
True if the URL is localhost/internal and should be replaced
Examples:
>>> is_localhost_or_internal_url("http://localhost:8000/")
True
>>> is_localhost_or_internal_url("http://0.0.0.0:8001/")
True
>>> is_localhost_or_internal_url("https://my-agent.example.com/")
False
>>> is_localhost_or_internal_url(None)
False
True if the URL is localhost/internal
"""
if not url:
return False
@ -59,50 +48,40 @@ def is_localhost_or_internal_url(url: Optional[str]) -> bool:
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) and /.well-known/agent.json (previous/alternative)
- Fix agent card URLs that contain localhost/internal addresses by replacing them with the original base_url
This fixes a common issue where agent cards are deployed with internal URLs
(e.g., "http://0.0.0.0:8001/") that don't work when accessed from external clients.
Extends the base A2ACardResolver to try both:
- /.well-known/agent-card.json (standard)
- /.well-known/agent.json (previous/alternative)
"""
def _fix_agent_card_url(self, agent_card: "AgentCard") -> "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 method
replaces such URLs with the original base_url used to fetch the card.
Args:
agent_card: The agent card to fix
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 = self.base_url.rstrip("/") + "/"
verbose_logger.warning(
f"Agent card contains localhost/internal URL '{card_url}'. "
f"Replacing with base_url '{fixed_url}'. "
f"Consider updating the agent's configuration to use the correct public URL."
)
# Create a new agent card with the fixed URL
# We need to handle this carefully to preserve all other fields
agent_card.url = fixed_url
return agent_card
async def get_agent_card(
self,
relative_card_path: Optional[str] = None,
@ -110,57 +89,54 @@ 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.
Also fixes agent card URLs that contain localhost/internal addresses.
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 (with URL fixed if necessary)
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:
agent_card = await super().get_agent_card(
return await super().get_agent_card(
relative_card_path=relative_card_path,
http_kwargs=http_kwargs,
)
return self._fix_agent_card_url(agent_card)
# 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}"
)
agent_card = await super().get_agent_card(
return await super().get_agent_card(
relative_card_path=path,
http_kwargs=http_kwargs,
)
return self._fix_agent_card_url(agent_card)
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}. "

View file

@ -1,6 +1,11 @@
"""
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:

View file

@ -9,7 +9,10 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.a2a_protocol.card_resolver import is_localhost_or_internal_url
from litellm.a2a_protocol.card_resolver import (
fix_agent_card_url,
is_localhost_or_internal_url,
)
@pytest.mark.asyncio
@ -83,19 +86,13 @@ def test_is_localhost_or_internal_url():
def test_fix_agent_card_url_replaces_localhost():
"""Test that _fix_agent_card_url replaces localhost URLs with base_url."""
from litellm.a2a_protocol.card_resolver import LiteLLMA2ACardResolver
"""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/"
# Create resolver instance without calling __init__
resolver = LiteLLMA2ACardResolver.__new__(LiteLLMA2ACardResolver)
resolver.base_url = "https://my-public-agent.example.com"
# Fix the URL
result = resolver._fix_agent_card_url(mock_card)
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/"