Fix three bugs in A2A discovery flow

cursoragent · yassin-berriai · cursoragent · commit a65357f9bb79 · 2026-05-28T03:36:43.000Z
1. UI: Stabilize discoveryRequest deps to avoid redundant /v1/a2a/discover
   API calls. The parent rebuilds the discoveryRequest object on every form
   keystroke, so depend on primitive proxies (discovery_mode + serialized
   params) rather than the object identity. Read the actual object via a
   ref inside handleDiscover.

2. Backend: Route the well-known card fetch through async_safe_get so the
   admin /v1/a2a/discover endpoint can't be used to probe private/loopback
   addresses or cloud metadata endpoints. SSRFError is a separate handled
   case so it surfaces a clear AgentCardDiscoveryError.

3. Streaming: Make openai_chunk_to_a2a_chunk emit the same flat result
   shape as the non-streaming response (kind/role/parts/messageId at the
   result level), with envelope-level 'final' added. Matches the existing
   create_artifact_update_event pattern and lets consumers read a uniform
   result shape across streaming and non-streaming.

Co-authored-by: Yassin Kortam &lt;yassin@berri.ai&gt;
diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py
@@ -335,19 +335,22 @@ def openai_chunk_to_a2a_chunk(
         if not content and not is_final:
             return None
 
-        # Build A2A streaming chunk (legacy format). ``final`` is an
-        # envelope-level streaming property per the A2A spec and must live
-        # alongside ``message`` in ``result``, not inside the message object.
+        # Build A2A streaming chunk. Mirrors the non-streaming response
+        # shape (``result`` is the message itself, with ``kind: "message"``
+        # as the result-level event discriminator — matching how
+        # ``create_artifact_update_event`` uses ``kind: "artifact-update"``
+        # at the result level). ``final`` is an envelope-level streaming
+        # property per the A2A spec and is appended alongside the message
+        # fields so consumers can read a uniform ``result`` shape across
+        # streaming and non-streaming.
         a2a_chunk = {
             "jsonrpc": "2.0",
             "id": request_id,
             "result": {
-                "message": {
-                    "kind": "message",
-                    "role": "agent",
-                    "parts": [{"kind": "text", "text": content}],
-                    "messageId": uuid4().hex,
-                },
+                "kind": "message",
+                "role": "agent",
+                "parts": [{"kind": "text", "text": content}],
+                "messageId": uuid4().hex,
                 "final": is_final,
             },
         }
diff --git a/litellm/proxy/a2a/discovery.py b/litellm/proxy/a2a/discovery.py
@@ -19,6 +19,7 @@
 from urllib.parse import urlencode
 
 from litellm._logging import verbose_proxy_logger
+from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get
 from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
 from litellm.types.llms.custom_http import httpxSpecialProvider
 
@@ -113,7 +114,18 @@ async def fetch_well_known_card(
     for path in paths:
         url = f"{normalized}{path}"
         try:
-            response = await client.get(url, headers=headers)
+            # ``async_safe_get`` validates the URL against the SSRF blocklist
+            # (private/loopback IPs, cloud metadata endpoints, etc.) on every
+            # redirect hop. Even though the discovery endpoint is admin-only,
+            # we don't want a compromised admin key to be able to probe
+            # internal infrastructure through this fetcher.
+            response = await async_safe_get(client, url, headers=headers)
+        except SSRFError as exc:
+            last_error = f"{url}: {exc!s}"
+            verbose_proxy_logger.debug(
+                "A2A discovery blocked by SSRF guard for %s: %s", url, exc
+            )
+            continue
         except Exception as exc:
             last_error = f"{url}: {exc!s}"
             verbose_proxy_logger.debug("A2A discovery failed for %s: %s", url, exc)
diff --git a/tests/test_litellm/proxy/a2a/test_discovery.py b/tests/test_litellm/proxy/a2a/test_discovery.py
@@ -6,6 +6,7 @@
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 
+import litellm
 from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
 from litellm.proxy.a2a.discovery import (
     AGENT_CARD_WELL_KNOWN_PATHS,
@@ -16,6 +17,17 @@
 from litellm.proxy.a2a.endpoints import router as a2a_router
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 
+
+@pytest.fixture(autouse=True)
+def _disable_url_validation_for_mocks(monkeypatch):
+    """The fetch tests use placeholder hostnames (``upstream.example``,
+    ``localhost:2024``) with mocked HTTP clients. ``async_safe_get`` would
+    otherwise resolve those hostnames and either fail DNS or block on the
+    SSRF guard. Disabling validation here lets the unit tests focus on
+    fallback / parsing logic; SSRF behavior is covered in its own test."""
+    monkeypatch.setattr(litellm, "user_url_validation", False)
+
+
 # ---------------------------------------------------------------------------
 # fetch_well_known_card
 # ---------------------------------------------------------------------------
diff --git a/ui/litellm-dashboard/src/components/agents/agent_card_discovery.tsx b/ui/litellm-dashboard/src/components/agents/agent_card_discovery.tsx
@@ -1,6 +1,6 @@
 "use client";
 
-import React, { useCallback, useEffect, useRef, useState } from "react";
+import React, { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import {
   Alert,
   Button,
@@ -110,6 +110,20 @@ const AgentCardDiscovery: React.FC<AgentCardDiscoveryProps> = ({
   onApplyRef.current = onApply;
   const discoverRequestIdRef = useRef(0);
   const lastSyncedSelectionRef = useRef<string | null>(null);
+  // Hold the latest ``discoveryRequest`` in a ref so ``handleDiscover`` can
+  // read its ``discovery_mode``/``params`` without depending on the object
+  // identity itself — the parent recreates the object on every form keystroke
+  // even when the underlying values are unchanged. We use stable primitive
+  // keys (``discoveryMode`` + ``discoveryParamsKey``) as the actual deps so
+  // the callback / effect only re-run when content actually changes.
+  const discoveryRequestRef = useRef(discoveryRequest);
+  discoveryRequestRef.current = discoveryRequest;
+
+  const discoveryMode = discoveryRequest?.discovery_mode;
+  const discoveryParamsKey = useMemo(
+    () => JSON.stringify(discoveryRequest?.params ?? null),
+    [discoveryRequest?.params],
+  );
 
   const handleDiscover = useCallback(async () => {
     if (!accessToken) {
@@ -129,17 +143,18 @@ const AgentCardDiscovery: React.FC<AgentCardDiscoveryProps> = ({
       return;
     }
 
+    const currentDiscoveryRequest = discoveryRequestRef.current;
     const requestId = ++discoverRequestIdRef.current;
     setLoading(true);
     setError(null);
     try {
       const response = await discoverAgentCardCall(
         accessToken,
         trimmed,
-        isParentDriven
+        isParentDriven && currentDiscoveryRequest
           ? {
-              discovery_mode: discoveryRequest!.discovery_mode,
-              params: discoveryRequest!.params,
+              discovery_mode: currentDiscoveryRequest.discovery_mode,
+              params: currentDiscoveryRequest.params,
             }
           : undefined,
       );
@@ -158,7 +173,11 @@ const AgentCardDiscovery: React.FC<AgentCardDiscoveryProps> = ({
         setLoading(false);
       }
     }
-  }, [accessToken, discoveryRequest, effectiveUrl, isParentDriven]);
+    // ``discoveryMode`` / ``discoveryParamsKey`` are primitive proxies for
+    // ``discoveryRequest`` content; the actual object is read via the ref
+    // above so identity churn from the parent doesn't recreate this callback.
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [accessToken, effectiveUrl, isParentDriven, discoveryMode, discoveryParamsKey]);
 
   // Auto-discover when the URL (or parent plan) becomes available.
   useEffect(() => {
@@ -177,15 +196,7 @@ const AgentCardDiscovery: React.FC<AgentCardDiscoveryProps> = ({
       void handleDiscover();
     }, debounceMs);
     return () => window.clearTimeout(timer);
-  }, [
-    accessToken,
-    effectiveUrl,
-    isParentDriven,
-    discoveryRequest?.url,
-    discoveryRequest?.discovery_mode,
-    discoveryRequest?.params,
-    handleDiscover,
-  ]);
+  }, [accessToken, effectiveUrl, isParentDriven, handleDiscover]);
 
   const toggleSkill = (id: string, checked: boolean) => {
     setSelectedSkillIds((prev) => {
