From ce4a25721f9f560350614fb294db69e2a9cb1c44 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:23:55 +0000 Subject: [PATCH] fix(proxy): route registered a2a agents before wildcard models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/route_llm_request.py | 30 +++-- .../basic/test_chat_completions_basic_a2a.py | 104 +++++++++++++++++- tests/unit/proxy/test_route_llm_request.py | 58 ++++++++++ 3 files changed, 182 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index e16e14c8d54..f64be103d32 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -888,6 +888,22 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr return getattr(llm_router, f"{route_type}")(**data) elif data["model"] not in router_model_names: + a2a_model_name: Final[object] = JSON_OBJECT_ADAPTER.validate_python(data).get("model", "") + if isinstance(a2a_model_name, str) and _is_a2a_agent_model(a2a_model_name): + from litellm.proxy.common_utils.registry_read_through import ( + get_agent_with_read_through, + ) + + agent_name: Final = a2a_model_name[4:] + if await get_agent_with_read_through(agent_name) is not None: + from litellm.proxy.agent_endpoints.a2a_routing import ( + route_a2a_agent_request, + ) + + result: Final = await route_a2a_agent_request(data, route_type, user_api_key_dict=user_api_key_dict) + if result is not None: + return result + # Check wildcards before checking deployment_names # Priority: 1. Exact model_name match, 2. Wildcard match, 3. deployment_names match if llm_router.router_general_settings.pass_through_all_models: @@ -944,16 +960,12 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr except Exception: # If router fails (e.g., model not found in router), fall back to direct call return getattr(litellm, f"{route_type}")(**direct_litellm_kwargs()) - elif _is_a2a_agent_model(data.get("model", "")): - from litellm.proxy.agent_endpoints.a2a_routing import ( - route_a2a_agent_request, + elif isinstance(a2a_model_name, str) and _is_a2a_agent_model(a2a_model_name): + raise ProxyModelNotFoundError( + route=ROUTE_ENDPOINT_MAPPING.get(route_type, route_type), + model_name=a2a_model_name, + retryable_with_model_read_through=False, ) - - result: Final = await route_a2a_agent_request(data, route_type, user_api_key_dict=user_api_key_dict) - if result is not None: - return result - # Fall through to raise exception below if result is None - elif user_model is not None or route_type == "allm_passthrough_route": return getattr(litellm, f"{route_type}")(**direct_litellm_kwargs()) diff --git a/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_a2a.py b/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_a2a.py index 707608890f8..d53f442a205 100644 --- a/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_a2a.py +++ b/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_a2a.py @@ -5,7 +5,7 @@ from typing import Final import httpx import pytest -from integration._support.client import Gateway +from integration._support.client import Gateway, object_value, string_value from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue @@ -164,3 +164,105 @@ def test_a2a_completion_through_proxy( requests: Final = wire.drain() assert len(tuple(item for item in requests if item.method == "POST")) == 1 + + +def test_registered_a2a_agent_routes_before_wildcard_model(gateway: Gateway) -> None: + agent_name: Final = "translation-a2a-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + if request.method == "GET": + assert request.target in ("/.well-known/agent-card.json", "/.well-known/agent.json") + return Reply( + body=json.dumps( + { + "protocolVersion": "0.3", + "name": agent_name, + "description": "Scripted A2A peer", + "version": "1.0.0", + "url": wire.url + "/", + "capabilities": {"streaming": True}, + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "skills": [], + } + ).encode() + ) + + assert request.method == "POST" and request.target == "/" + outbound: Final = json.loads(request.body) + request_id: Final = outbound["id"] + return Reply( + body=json.dumps( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "kind": "message", + "role": "agent", + "messageId": agent_name, + "parts": [{"kind": "text", "text": "A2A wildcard response"}], + }, + } + ).encode() + ) + + with wire_server(upstream) as wire, gateway.scenario() as scenario: + wildcard_prefix: Final = "integration-" + uuid.uuid4().hex + wildcard_created: Final = gateway.request( + "POST", + "/model/new", + { + "model_name": f"{wildcard_prefix}/*", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "integration-provider-key", + "api_base": gateway.upstream_url, + }, + }, + ) + assert wildcard_created.status_code == 200, wildcard_created.text + wildcard_id: Final = string_value(object_value(wildcard_created.json()["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, wildcard_id) + + missing_agent: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": f"a2a/{agent_name}-missing", + "messages": [{"role": "user", "content": "Hello"}], + }, + ) + assert missing_agent.status_code == 400, missing_agent.text + + card: Final = { + "protocolVersion": "0.3", + "name": agent_name, + "description": "Scripted A2A peer", + "version": "1.0.0", + "url": wire.url + "/", + "capabilities": {"streaming": True}, + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "skills": [], + } + created: Final = gateway.request( + "POST", + "/v1/agents", + {"agent_name": agent_name, "agent_card_params": card}, + ) + assert created.status_code == 200, created.text + agent_id: Final = string_value(created.json()["agent_id"]) + scenario.cleanups.callback( + lambda: gateway.request("DELETE", f"/v1/agents/{agent_id}") + ) + + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": f"a2a/{agent_name}", + "messages": [{"role": "user", "content": "Hello"}], + }, + ) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "A2A wildcard response" diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 8928b79bcd4..e57bee0b743 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -2117,6 +2117,64 @@ async def test_route_request_routing_group_name_passes_model_gate(): spy.assert_called_once_with(**data) +@pytest.mark.asyncio +@pytest.mark.parametrize("agent_registered", [True, False]) +async def test_route_request_a2a_agent_precedes_pattern_route( + agent_registered: bool, fresh_agent_read_through: object, monkeypatch: pytest.MonkeyPatch +) -> None: + from types import SimpleNamespace + from unittest.mock import AsyncMock + import uuid + + import litellm + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.types.agents import AgentResponse + + agent_name: Final = "pattern-route-agent-" + uuid.uuid4().hex + model_name: Final = f"a2a/{agent_name}" + router: Final = litellm.Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}, + } + ] + ) + pattern_route: Final = AsyncMock(return_value="pattern response") + monkeypatch.setattr(router, "acompletion", pattern_route) + agent: Final = AgentResponse( + agent_id=f"{agent_name}-id", + agent_name=agent_name, + agent_card_params={"name": agent_name, "url": "http://agent.example"}, + ) + agents_find_unique: Final = AsyncMock(side_effect=[None, agent] if agent_registered else [None, None]) + fake_prisma, _ = _fake_prisma_client_with_models([]) + fake_prisma.db.litellm_agentstable = SimpleNamespace(find_unique=agents_find_unique) + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", router) + + data: Final = {"model": model_name, "messages": [{"role": "user", "content": "hi"}]} + admin_auth: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + if agent_registered: + agent_route: Final = AsyncMock(return_value="agent response") + monkeypatch.setattr(litellm, "acompletion", agent_route) + agent_response: Final[object] = await (await route_request(data, router, None, "acompletion", admin_auth)) + assert agent_response == "agent response" + agent_route.assert_awaited_once() + agent_call: Final = agent_route.await_args + assert agent_call is not None + assert agent_call.kwargs["api_base"] == "http://agent.example" + pattern_route.assert_not_awaited() + return + + pattern_response: Final[object] = await (await route_request(data, router, None, "acompletion")) + assert pattern_response == "pattern response" + pattern_route.assert_awaited_once_with(**data) + assert agents_find_unique.await_count == 2 + + @pytest.mark.asyncio async def test_route_request_a2a_agent_miss_does_not_consume_model_read_through(fresh_agent_read_through, monkeypatch): from types import SimpleNamespace