scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-02 07:13:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 148

fix(proxy): serve /a2a/{agent_id}/message/stream routes (#49471)

Browse source 
* test(a2a): cover message stream route aliases

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): serve /a2a/{agent_id}/message/stream routes

Co-authored-by: Rustaman Delvianova <96533738+Rustaman1280@users.noreply.github.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): document A2A stream route schemas

* fix(proxy): use immutable A2A schema mappings

* fix(proxy): map v1 a2a stream route to call types and keep stream alias unbilled

* test(a2a): use real guardrail registry and typed locals in stream route tests

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Rustaman Delvianova <96533738+Rustaman1280@users.noreply.github.com>
This commit is contained in:
moe-berri 2026-10-31 16:11:51 -07:00 • committed by GitHub
parent ddab201d68
commit bba63c1934
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 590 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/_lazy_features.py
View file

@ -105,10 +105,10 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
LazyFeature(
name="a2a",
module_path="litellm.proxy.agent_endpoints.a2a_endpoints",
# ``/v1/a2a/{agent_id}/message/send`` is caught via the suffix so the
# ``/v1/a2a/{agent_id}/message/send`` and ``/message/stream`` are caught via suffixes so the
# ``/v1/a2a`` prefix doesn't subsume the discover prefix below.
path_prefixes=("/a2a",),
path_suffixes=("/message/send",),
path_suffixes=("/message/send", "/message/stream"),
),
LazyFeature(
name="a2a_registration",

277
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -54,6 +54,57 @@
],
"title": "ValidationError",
"type": "object"
},
"_A2AJsonRpcResponse": {
"additionalProperties": true,
"properties": {
"error": {
"anyOf": [
{
"additionalProperties": true,
"type": "object"
},
{
"type": "null"
}
],
"title": "Error"
},
"id": {
"anyOf": [
{
"type": "string"
},
{
"type": "integer"
},
{
"type": "null"
}
],
"title": "Id"
},
"jsonrpc": {
"const": "2.0",
"title": "Jsonrpc",
"type": "string"
},
"result": {
"anyOf": [
{},
{
"type": "null"
}
],
"title": "Result"
}
},
"required": [
"jsonrpc",
"id"
],
"title": "_A2AJsonRpcResponse",
"type": "object"
}
}
},
@ -242,6 +293,119 @@
]
}
},
"/a2a/{agent_id}/message/stream": {
"post": {
"description": "Invoke an agent using the A2A protocol (JSON-RPC 2.0).\n\nSupported methods:\n- message/send: Send a message and get a response\n- message/stream: Send a message and stream the response",
"operationId": "invoke_agent_a2a_a2a__agent_id__message_stream_post",
"parameters": [
{
"in": "path",
"name": "agent_id",
"required": true,
"schema": {
"title": "Agent Id",
"type": "string"
}
}
],
"requestBody": {
"content": {
"application/json": {
"example": {
"id": "stream-1",
"jsonrpc": "2.0",
"method": "message/stream",
"params": {
"message": {
"kind": "message",
"messageId": "message-1",
"parts": [
{
"kind": "text",
"text": "Hello"
}
],
"role": "user"
}
}
},
"schema": {
"additionalProperties": true,
"properties": {
"id": {
"anyOf": [
{
"type": "string"
},
{
"type": "integer"
},
{
"type": "null"
}
],
"title": "Id"
},
"jsonrpc": {
"const": "2.0",
"title": "Jsonrpc",
"type": "string"
},
"method": {
"title": "Method",
"type": "string"
},
"params": {
"additionalProperties": true,
"title": "Params",
"type": "object"
}
},
"required": [
"jsonrpc",
"method",
"params"
],
"title": "_A2AJsonRpcRequest",
"type": "object"
}
}
},
"required": true
},
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/_A2AJsonRpcResponse"
}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Invoke Agent A2A",
"tags": [
"a2a"
]
}
},
"/v1/a2a/{agent_id}/message/send": {
"post": {
"description": "Invoke an agent using the A2A protocol (JSON-RPC 2.0).\n\nSupported methods:\n- message/send: Send a message and get a response\n- message/stream: Send a message and stream the response",
@ -287,6 +451,119 @@
"a2a"
]
}
},
"/v1/a2a/{agent_id}/message/stream": {
"post": {
"description": "Invoke an agent using the A2A protocol (JSON-RPC 2.0).\n\nSupported methods:\n- message/send: Send a message and get a response\n- message/stream: Send a message and stream the response",
"operationId": "invoke_agent_a2a_v1_a2a__agent_id__message_stream_post",
"parameters": [
{
"in": "path",
"name": "agent_id",
"required": true,
"schema": {
"title": "Agent Id",
"type": "string"
}
}
],
"requestBody": {
"content": {
"application/json": {
"example": {
"id": "stream-1",
"jsonrpc": "2.0",
"method": "message/stream",
"params": {
"message": {
"kind": "message",
"messageId": "message-1",
"parts": [
{
"kind": "text",
"text": "Hello"
}
],
"role": "user"
}
}
},
"schema": {
"additionalProperties": true,
"properties": {
"id": {
"anyOf": [
{
"type": "string"
},
{
"type": "integer"
},
{
"type": "null"
}
],
"title": "Id"
},
"jsonrpc": {
"const": "2.0",
"title": "Jsonrpc",
"type": "string"
},
"method": {
"title": "Method",
"type": "string"
},
"params": {
"additionalProperties": true,
"title": "Params",
"type": "object"
}
},
"required": [
"jsonrpc",
"method",
"params"
],
"title": "_A2AJsonRpcRequest",
"type": "object"
}
}
},
"required": true
},
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/_A2AJsonRpcResponse"
}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Invoke Agent A2A",
"tags": [
"a2a"
]
}
}
}
},

1
litellm/proxy/_types.py
View file

@ -501,6 +501,7 @@ class LiteLLMRoutes(enum.Enum):
"/a2a/{agent_id}/message/send",
"/v1/a2a/{agent_id}/message/send",
"/a2a/{agent_id}/message/stream",
"/v1/a2a/{agent_id}/message/stream",
"/a2a/{agent_id}/.well-known/agent-card.json",
)

53
litellm/proxy/agent_endpoints/a2a_endpoints.py
View file

@ -14,7 +14,7 @@ import json
from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping
from copy import deepcopy
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
from urllib.parse import urlparse
from fastapi import APIRouter, Depends, HTTPException, Request, Response
@ -43,6 +43,7 @@ from litellm.proxy.agent_endpoints.databricks_oauth import (
)
from litellm.proxy.agent_endpoints.utils import merge_agent_headers
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.custom_openapi_spec import inline_request_body
from litellm.proxy.common_utils.openai_error_payload import status_typed_proxy_exception
from litellm.proxy.common_utils.sse_keepalive import (
SSE_COMMENT_PING,
@ -59,6 +60,42 @@ if TYPE_CHECKING:
router: Final = APIRouter()
class _A2AJsonRpcRequest(BaseModel):
model_config = ConfigDict(extra="allow", frozen=True)
jsonrpc: Literal["2.0"]
id: str | int | None = None
method: str
params: Mapping[str, object]
class _A2AJsonRpcResponse(BaseModel):
model_config = ConfigDict(extra="allow", frozen=True)
jsonrpc: Literal["2.0"]
id: str | int | None
result: object | None = None
error: Mapping[str, object] | None = None
_A2A_MESSAGE_STREAM_REQUEST_BODY: Final = inline_request_body(
_A2AJsonRpcRequest,
{
"jsonrpc": "2.0",
"id": "stream-1",
"method": "message/stream",
"params": {
"message": {
"messageId": "message-1",
"role": "user",
"parts": [{"kind": "text", "text": "Hello"}],
"kind": "message",
}
},
},
)
# Mirrors the native seam's own headers: a reverse proxy that batches the whole
# stream would swallow the keepalives this route sends to defeat idle timeouts.
_SSE_KEEPALIVE_HEADERS: Final[Mapping[str, str]] = MappingProxyType(
@ -746,6 +783,20 @@ async def get_agent_card(
tags=["[beta] A2A Agents"],
dependencies=[Depends(user_api_key_auth)],
)
@router.post(
"/a2a/{agent_id}/message/stream",
tags=["[beta] A2A Agents"],
dependencies=[Depends(user_api_key_auth)],
openapi_extra={"requestBody": _A2A_MESSAGE_STREAM_REQUEST_BODY},
responses={200: {"model": _A2AJsonRpcResponse}},
)
@router.post(
"/v1/a2a/{agent_id}/message/stream",
tags=["[beta] A2A Agents"],
dependencies=[Depends(user_api_key_auth)],
openapi_extra={"requestBody": _A2A_MESSAGE_STREAM_REQUEST_BODY},
responses={200: {"model": _A2AJsonRpcResponse}},
)
async def invoke_agent_a2a(
agent_id: str,
request: Request,

2
litellm/types/utils.py
View file

@ -1151,6 +1151,8 @@ API_ROUTE_TO_CALL_TYPES: Final[Mapping[str, Sequence[CallTypes]]] = {
# A2A (Agent-to-Agent)
"/a2a/{agent_id}": [CallTypes.asend_message, CallTypes.send_message],
"/a2a/{agent_id}/message/send": [CallTypes.asend_message, CallTypes.send_message],
"/a2a/{agent_id}/message/stream": [CallTypes.asend_message, CallTypes.send_message],
"/v1/a2a/{agent_id}/message/stream": [CallTypes.asend_message, CallTypes.send_message],
# Passthrough endpoints
"/llm_passthrough": [
CallTypes.allm_passthrough_route,

9
tests/integration/compatibility/test_a2a_wire_versions.py
View file

@ -452,8 +452,13 @@ def test_guard_no_agent_key_or_proxy_caller_key_is_sent(gateway: Gateway) -> Non
(("0.3", "message/stream", _MESSAGE_V03), ("1.0", "SendStreamingMessage", _MESSAGE_V10)),
ids=("v0.3", "v1.0"),
)
@pytest.mark.parametrize(
"path",
("/a2a/{identity}", "/a2a/{identity}/message/stream", "/v1/a2a/{identity}/message/stream"),
ids=("bare", "message_stream", "v1_message_stream"),
)
def test_a2a_stream_relays_each_upstream_event_as_its_own_sse_frame_in_the_pinned_dialect(
gateway: Gateway, version: str, method: str, message: dict[str, JsonValue]
gateway: Gateway, version: str, method: str, message: dict[str, JsonValue], path: str
) -> None:
marker: Final = "a2astream" + uuid.uuid4().hex
first_frame_seen: Final = threading.Event()
@ -482,7 +487,7 @@ def test_a2a_stream_relays_each_upstream_event_as_its_own_sse_frame_in_the_pinne
)
with gateway.client.stream(
"POST",
f"/a2a/{identity}",
path.format(identity=identity),
headers={
"Authorization": f"Bearer {key}",
"a2a-version": version,

12
tests/unit/litellm_core_utils/test_api_route_to_call_types.py
View file

@ -8,6 +8,8 @@ Regression coverage for the guardrail route table bugs:
take call_types[0] to a handler-less type
"""
from typing import Final
from litellm.litellm_core_utils.api_route_to_call_types import (
get_call_types_for_route,
)
@ -103,6 +105,16 @@ class TestExistingRouteResolutionUnchanged:
assert call_types is not None
assert list(call_types) == [CallTypes.asend_message, CallTypes.send_message]
def test_single_segment_placeholder_message_stream_resolves(self):
call_types: Final = get_call_types_for_route("/a2a/my-agent/message/stream")
assert call_types is not None
assert list(call_types) == [CallTypes.asend_message, CallTypes.send_message]
def test_v1_single_segment_placeholder_message_stream_resolves(self):
call_types: Final = get_call_types_for_route("/v1/a2a/my-agent/message/stream")
assert call_types is not None
assert list(call_types) == [CallTypes.asend_message, CallTypes.send_message]
def test_longer_route_does_not_collapse_into_bare_placeholder_pattern(self):
call_types = get_call_types_for_route("/responses/resp_123/input_items")
assert call_types is not None

1
tests/unit/proxy/auth/test_route_checks.py
View file

@ -3925,6 +3925,7 @@ AGENT_INFERENCE_ROUTES = [
"/a2a/abc-123",
"/a2a/abc-123/message/send",
"/a2a/abc-123/message/stream",
"/v1/a2a/abc-123/message/stream",
"/a2a/abc-123/.well-known/agent-card.json",
]

22
tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
View file

@ -2025,6 +2025,28 @@ class TestStreamingHttpErrorFrames:
def _use_real_mappings(self, monkeypatch):
_patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings())
@pytest.mark.asyncio
async def test_v1_a2a_stream_route_scans_response_chunks(self):
guardrail: Final = RecordingGuardrail()
guardrail.streaming_sampling_rate = 1
chunk: Final = {
"jsonrpc": "2.0",
"id": "stream-1",
"result": {
"kind": "task",
"artifacts": [{"parts": [{"kind": "text", "text": "streamed text"}]}],
},
}
await _drive_stream(
UnifiedLLMGuardrails(),
guardrail,
[chunk],
request_route="/v1/a2a/my-agent/message/stream",
)
assert guardrail.apply_calls[0]["inputs"]["texts"] == ["streamed text"]
@pytest.mark.asyncio
async def test_chat_eos_block_emits_data_error_frame(self):
import json as _json

6
tests/unit/proxy/test__lazy_features.py
View file

@ -10,6 +10,7 @@ from fastapi.testclient import TestClient
from pydantic import BaseModel
from litellm.proxy._lazy_features import (
LAZY_FEATURES,
LazyFeature,
LazyFeatureMiddleware,
attach_lazy_features,
@ -66,6 +67,11 @@ def _has_lazy_middleware(app: FastAPI) -> bool:
return any(middleware.cls is LazyFeatureMiddleware for middleware in app.user_middleware)
def test_a2a_feature_matches_v1_message_stream() -> None:
feature: Final = next(feature for feature in LAZY_FEATURES if feature.name == "a2a")
assert feature.matches("/v1/a2a/x/message/stream")
@pytest.mark.parametrize("disable_lazy_routes", ("false", "true"))
def test_retired_roi_calculator_cannot_be_loaded_or_called(
monkeypatch: pytest.MonkeyPatch, disable_lazy_routes: str

208
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -452,6 +452,30 @@ export interface paths {
patch?: never;
trace?: never;
};
"/a2a/{agent_id}/message/stream": {
parameters: {
query?: never;
header?: never;
path?: never;
cookie?: never;
};
get?: never;
put?: never;
/**
* Invoke Agent A2A
* @description Invoke an agent using the A2A protocol (JSON-RPC 2.0).
*
* Supported methods:
* - message/send: Send a message and get a response
* - message/stream: Send a message and stream the response
*/
post: operations["invoke_agent_a2a_a2a__agent_id__message_stream_post"];
delete?: never;
options?: never;
head?: never;
patch?: never;
trace?: never;
};
"/access_group/list": {
parameters: {
query?: never;
@ -19307,6 +19331,30 @@ export interface paths {
patch?: never;
trace?: never;
};
"/v1/a2a/{agent_id}/message/stream": {
parameters: {
query?: never;
header?: never;
path?: never;
cookie?: never;
};
get?: never;
put?: never;
/**
* Invoke Agent A2A
* @description Invoke an agent using the A2A protocol (JSON-RPC 2.0).
*
* Supported methods:
* - message/send: Send a message and get a response
* - message/stream: Send a message and stream the response
*/
post: operations["invoke_agent_a2a_v1_a2a__agent_id__message_stream_post"];
delete?: never;
options?: never;
head?: never;
patch?: never;
trace?: never;
};
"/v1/access_group": {
parameters: {
query?: never;
@ -50790,6 +50838,24 @@ export interface components {
/** Status */
status?: ("pending" | "running" | "paused" | "completed" | "failed") | null;
};
/** _A2AJsonRpcResponse */
_A2AJsonRpcResponse: {
/** Error */
error?: {
[key: string]: unknown;
} | null;
/** Id */
id: string | number | null;
/**
* Jsonrpc
* @constant
*/
jsonrpc: "2.0";
/** Result */
result?: unknown | null;
} & {
[key: string]: unknown;
};
/** ModelInfo */
litellm__proxy___types__ModelInfo: {
/** Base Model */
@ -52149,6 +52215,77 @@ export interface operations {
};
};
};
invoke_agent_a2a_a2a__agent_id__message_stream_post: {
parameters: {
query?: never;
header?: never;
path: {
agent_id: string;
};
cookie?: never;
};
requestBody: {
content: {
/**
* @example {
* "id": "stream-1",
* "jsonrpc": "2.0",
* "method": "message/stream",
* "params": {
* "message": {
* "kind": "message",
* "messageId": "message-1",
* "parts": [
* {
* "kind": "text",
* "text": "Hello"
* }
* ],
* "role": "user"
* }
* }
* }
*/
"application/json": {
/** Id */
id?: string | number | null;
/**
* Jsonrpc
* @constant
*/
jsonrpc: "2.0";
/** Method */
method: string;
/** Params */
params: {
[key: string]: unknown;
};
} & {
[key: string]: unknown;
};
};
};
responses: {
/** @description Successful Response */
200: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["_A2AJsonRpcResponse"];
};
};
/** @description Validation Error */
422: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["HTTPValidationError"];
};
};
};
};
list_access_groups_access_group_list_get: {
parameters: {
query?: never;
@ -77782,6 +77919,77 @@ export interface operations {
};
};
};
invoke_agent_a2a_v1_a2a__agent_id__message_stream_post: {
parameters: {
query?: never;
header?: never;
path: {
agent_id: string;
};
cookie?: never;
};
requestBody: {
content: {
/**
* @example {
* "id": "stream-1",
* "jsonrpc": "2.0",
* "method": "message/stream",
* "params": {
* "message": {
* "kind": "message",
* "messageId": "message-1",
* "parts": [
* {
* "kind": "text",
* "text": "Hello"
* }
* ],
* "role": "user"
* }
* }
* }
*/
"application/json": {
/** Id */
id?: string | number | null;
/**
* Jsonrpc
* @constant
*/
jsonrpc: "2.0";
/** Method */
method: string;
/** Params */
params: {
[key: string]: unknown;
};
} & {
[key: string]: unknown;
};
};
};
responses: {
/** @description Successful Response */
200: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["_A2AJsonRpcResponse"];
};
};
/** @description Validation Error */
422: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["HTTPValidationError"];
};
};
};
};
list_access_groups_v1_access_group_get: {
parameters: {
query?: never;