add localhost retry

litellm/a2a_protocol/main.py

@@ -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)
 
@@ -403,15 +448,15 @@ 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"
+
+    # Track for logging
+    start_time = datetime.datetime.now()
     model = f"a2a_agent/{agent_name}"
 
     logging_obj = Logging(
@@ -443,15 +488,56 @@ async def asend_message_streaming(
     logging_obj.model_call_details["litellm_params"] = _litellm_params
     logging_obj.model_call_details["metadata"] = metadata or {}
 
-    iterator = A2AStreamingIterator(
-        stream=stream,
-        request=request,
-        logging_obj=logging_obj,
-        agent_name=agent_name,
-    )
+    # 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,
+        )
 
-    async for chunk in iterator:
-        yield chunk
+        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(