From a65357f9bb79c1bc1172cea5d71e92e043f98826 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 28 May 2026 03:36:08 +0000 Subject: [PATCH] Fix three bugs in A2A discovery flow 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 --- .../transformation.py | 21 +++++----- litellm/proxy/a2a/discovery.py | 14 ++++++- .../test_litellm/proxy/a2a/test_discovery.py | 12 ++++++ .../agents/agent_card_discovery.tsx | 39 ++++++++++++------- 4 files changed, 62 insertions(+), 24 deletions(-) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py index 3b9af87757d..69723f5f917 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py @@ -335,19 +335,22 @@ class A2ACompletionBridgeTransformation: 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 index 9a7d6e711f5..a28e980818d 100644 --- a/litellm/proxy/a2a/discovery.py +++ b/litellm/proxy/a2a/discovery.py @@ -19,6 +19,7 @@ from typing import Any, Dict, Optional, Tuple 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 index cd27309e481..ac1e7dfbb56 100644 --- a/tests/test_litellm/proxy/a2a/test_discovery.py +++ b/tests/test_litellm/proxy/a2a/test_discovery.py @@ -6,6 +6,7 @@ import pytest 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.discovery import ( 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 index 5505c9035ca..a5ca5f4d562 100644 --- 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 = ({ onApplyRef.current = onApply; const discoverRequestIdRef = useRef(0); const lastSyncedSelectionRef = useRef(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,6 +143,7 @@ const AgentCardDiscovery: React.FC = ({ return; } + const currentDiscoveryRequest = discoveryRequestRef.current; const requestId = ++discoverRequestIdRef.current; setLoading(true); setError(null); @@ -136,10 +151,10 @@ const AgentCardDiscovery: React.FC = ({ 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 = ({ 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 = ({ 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) => {