litellm/litellm/a2a_protocol/exception_mapping_utils.py

"""
A2A Protocol Exception Mapping Utils.

Maps A2A SDK exceptions to LiteLLM A2A exception types.
"""

from typing import TYPE_CHECKING, Any, Final

from litellm._logging import verbose_logger
from litellm.a2a_protocol.card_resolver import (
    is_localhost_or_internal_url,
    set_agent_card_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 Client as A2AClientType


try:
    from a2a.client import Client, ClientConfig, create_client

    A2A_SDK_AVAILABLE = True
except ImportError:
    A2A_SDK_AVAILABLE = False
    Client = None
    ClientConfig = None
    create_client = None


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: Final = error_str.lower()
        return any(pattern in error_str_lower for pattern in CONNECTION_ERROR_PATTERNS)

    @staticmethod
    def is_localhost_url(url: str | None) -> 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: Final = error_str.lower()
        agent_card_patterns: Final = [
            "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: str | None = None,
    api_base: str | None = None,
    model: str | None = 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: Final = 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,
    )


async 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:
        raise ImportError("A2A SDK is required for localhost retry handling. Install it with: pip install a2a-sdk")

    if agent_card is None:
        raise RuntimeError(
            "Cannot retry A2A localhost URL fix: no agent card is available to "
            "rewrite, so the upstream URL cannot be corrected."
        )

    request_type: Final = "streaming " if is_streaming else ""
    verbose_logger.warning(
        "A2A %srequest to '%s' failed: %s. Agent card contains localhost/internal URL. Retrying with base_url '%s'.",
        request_type,
        error.localhost_url,
        error.original_error,
        error.base_url,
    )

    # Fix the agent card URL
    set_agent_card_url(agent_card, error.base_url)

    # Reuse the httpx client and call context LiteLLM attached at creation, since the
    # context carries this agent's trace-id/auth headers. Only clients built by
    # ``create_a2a_client`` have them; an externally-supplied client cannot be retried.
    httpx_client: Final = getattr(a2a_client, "_litellm_httpx_client", None)
    if httpx_client is None:
        raise RuntimeError(
            "Cannot retry A2A localhost URL fix: the client was not created by "
            "create_a2a_client, so no LiteLLM httpx client is attached."
        )

    new_client: Final = await create_client(  # pyright: ignore[reportOptionalCall]
        agent_card,
        client_config=ClientConfig(  # pyright: ignore[reportOptionalCall]
            httpx_client=httpx_client,
            streaming=is_streaming,
        ),
    )
    new_client._litellm_httpx_client = httpx_client
    new_client._litellm_call_context = getattr(  # pyright: ignore[reportAttributeAccessIssue]  # LiteLLM-owned stash
        a2a_client, "_litellm_call_context", None
    )
    new_client._litellm_agent_card = agent_card
    return new_client