From cf934d2f33ba72882c13debb6e9c7e7c97337a84 Mon Sep 17 00:00:00 2001 From: "A.Kebriaei" Date: Tue, 23 Jun 2026 09:41:08 +0100 Subject: [PATCH 1/3] feat(a2a): completion-bridge agents can use MCP tools and delegate to other agents A2A completion-bridge agents were limited to a single model call. This adds an agentic loop so an agent can call the MCP tools it is entitled to and delegate to other A2A agents it can access, driven by the model until it produces a final answer. The inter-agent hop is a real A2A message/send, so a delegated agent behaves exactly as if it were called directly. The loop converts MCP tool results (_execute_tool_calls returns {tool_call_id, result, name}) into proper {role: "tool", tool_call_id, content} messages; without this the model never sees the result, re-calls the tool until the loop is exhausted, and returns an empty answer. Routing also drops any tools/tool_choice carried in from the agent's litellm_params before the loop or acompletion_with_mcp run, since those paths pass their own tools= explicitly. Delegation safety. The set of agents exposed for delegation is the caller's allowed agents intersected with the source agent's own allowed agents (its object_permission.agents plus access-group agents), so a prompt-injected delegation cannot exceed what either the caller or the source agent was granted. _delegate_to_agent then re-enforces that whitelist before dispatch (the target id is model-supplied and asend_message bypasses the endpoint auth gate), refuses targets that carry proxy-side guardrails (this direct path cannot run the guardrail / post-call hooks a direct invocation would), and re-scopes auth to the target. Auth is also bound to the invoked agent via _scope_auth_to_agent before MCP tools run, so tool access is the caller's perms intersected with the agent's own object_permission. The endpoint resolves a bare proxy model-list alias to a provider-prefixed model and caps delegation at one level. The dashboard agent create and Settings forms gain a prompt dropdown (populated from the prompts list) plus enable_mcp_tools and enable_agent_calls toggles. The prompt editor's model can now be cleared, so a prompt with no model leaves the calling agent's model in place while a prompt that pins a model overrides it. --- .../litellm_completion_bridge/agentic_loop.py | 199 +++++++++++ .../litellm_completion_bridge/handler.py | 107 ++++-- .../proxy/agent_endpoints/a2a_endpoints.py | 212 +++++++++++- .../a2a_protocol/test_agentic_loop.py | 312 ++++++++++++++++++ .../agent_endpoints/test_a2a_endpoints.py | 184 +++++++++++ .../prompt_editor_view/ModelConfigCard.tsx | 9 +- .../prompt_editor_view/utils.test.ts | 37 +++ .../components/prompt_editor_view/utils.ts | 6 +- .../src/components/agents/add_agent_form.tsx | 13 +- .../src/components/agents/agent_config.ts | 27 +- .../components/agents/agent_form_fields.tsx | 22 +- .../src/components/agents/agent_info.tsx | 14 +- .../common_components/ModelSelector.tsx | 17 +- 13 files changed, 1121 insertions(+), 38 deletions(-) create mode 100644 litellm/a2a_protocol/litellm_completion_bridge/agentic_loop.py create mode 100644 tests/test_litellm/a2a_protocol/test_agentic_loop.py diff --git a/litellm/a2a_protocol/litellm_completion_bridge/agentic_loop.py b/litellm/a2a_protocol/litellm_completion_bridge/agentic_loop.py new file mode 100644 index 00000000000..b64648ad71e --- /dev/null +++ b/litellm/a2a_protocol/litellm_completion_bridge/agentic_loop.py @@ -0,0 +1,199 @@ +"""Agentic tool loop for completion-bridge A2A agents. + +Drives the model in a loop (model -> tool/agent calls -> execute -> feed results +back -> repeat) until it produces a final answer. An agent can call MCP tools +and delegate to other A2A agents; the inter-agent hop is a real A2A message/send +supplied by the proxy layer via ``call_agent``. +""" + +import json +from typing import Any, Awaitable, Callable, Iterable, Mapping, Optional + +import litellm +from litellm._logging import verbose_logger + +AGENT_TOOL_PREFIX = "agent__" +DEFAULT_MAX_ITERATIONS = 8 + + +def _mcp_results_to_tool_messages( + tool_results: Iterable[Mapping[str, object]], +) -> tuple[dict[str, object], ...]: + """``_execute_tool_calls`` returns ``{tool_call_id, result, name}``, but the + chat API needs ``{role: "tool", tool_call_id, content}``; without this + conversion the model never sees the result and re-calls the tool until the + loop is exhausted, yielding an empty final answer.""" + return tuple( + { + "role": "tool", + "tool_call_id": result.get("tool_call_id"), + "content": str(result.get("result", "")), + } + for result in tool_results + ) + + +def _build_agent_tools(callable_agents: list[dict[str, Any]]) -> list[dict[str, Any]]: + """One function tool per callable agent. The model picks an agent by name; + the tool name carries the target agent id so we can route the A2A hop.""" + tools: list[dict[str, Any]] = [] + for agent in callable_agents: + agent_id = agent.get("agent_id") or agent.get("id") + if not agent_id: + continue + name = agent.get("name") or agent_id + description = ( + agent.get("description") + or f"Delegate a self-contained task to the '{name}' specialist agent and use its reply." + ) + tools.append( + { + "type": "function", + "function": { + "name": f"{AGENT_TOOL_PREFIX}{agent_id}", + "description": description, + "parameters": { + "type": "object", + "properties": { + "message": { + "type": "string", + "description": "The complete instruction and all facts the specialist needs; it has no prior context.", + } + }, + "required": ["message"], + }, + }, + } + ) + return tools + + +async def _load_mcp_tools( + user_api_key_auth: Any, + litellm_trace_id: Optional[str], +): + """Returns (openai_tools, tool_server_map) for the agent's allowed MCP servers.""" + from litellm.responses.mcp.litellm_proxy_mcp_handler import ( + LiteLLM_Proxy_MCP_Handler, + ) + + mcp_tools_spec = [ + {"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"} + ] + ( + deduplicated, + tool_server_map, + ) = await LiteLLM_Proxy_MCP_Handler._process_mcp_tools_without_openai_transform( + user_api_key_auth=user_api_key_auth, + mcp_tools_with_litellm_proxy=mcp_tools_spec, + litellm_trace_id=litellm_trace_id, + ) + openai_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai( + deduplicated, target_format="chat" + ) + return openai_tools, tool_server_map + + +async def run_agentic_loop( + *, + model: str, + messages: list[dict[str, Any]], + completion_kwargs: dict[str, Any], + callable_agents: list[dict[str, Any]], + call_agent: Callable[[str, str], Awaitable[str]], + user_api_key_auth: Any = None, + enable_mcp_tools: bool = False, + max_iterations: int = DEFAULT_MAX_ITERATIONS, +): + """ + Drive the model until it stops requesting tools/agents. + + - Agent (``agent__``) tool calls -> real A2A hop via ``call_agent``. + - MCP tool calls -> executed against the proxy MCP manager (auth-scoped). + Returns the final ModelResponse. + """ + agent_tools = _build_agent_tools(callable_agents) + + mcp_openai_tools: list[dict[str, Any]] = [] + tool_server_map: dict[str, str] = {} + if enable_mcp_tools: + mcp_openai_tools, tool_server_map = await _load_mcp_tools( + user_api_key_auth=user_api_key_auth, + litellm_trace_id=completion_kwargs.get("litellm_trace_id"), + ) + + all_tools = (mcp_openai_tools + agent_tools) or None + messages = list(messages) + response = None + + for _ in range(max_iterations): + response = await litellm.acompletion( + model=model, + messages=messages, + tools=all_tools, + stream=False, + **completion_kwargs, + ) + message = response.choices[0].message + tool_calls = getattr(message, "tool_calls", None) or [] + if not tool_calls: + break + + messages.append( + message.model_dump() if hasattr(message, "model_dump") else message + ) + + agent_calls = [ + tc + for tc in tool_calls + if (tc.function.name or "").startswith(AGENT_TOOL_PREFIX) + ] + mcp_calls = [tc for tc in tool_calls if tc not in agent_calls] + + if mcp_calls: + from litellm.responses.mcp.litellm_proxy_mcp_handler import ( + LiteLLM_Proxy_MCP_Handler, + ) + + tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map=tool_server_map, + tool_calls=mcp_calls, + user_api_key_auth=user_api_key_auth, + ) + messages.extend(_mcp_results_to_tool_messages(tool_results)) + + for tool_call in agent_calls: + target_agent_id = (tool_call.function.name or "")[len(AGENT_TOOL_PREFIX) :] + try: + arguments = json.loads(tool_call.function.arguments or "{}") + except json.JSONDecodeError: + arguments = {} + sub_message = arguments.get("message", "") + try: + reply = await call_agent(target_agent_id, sub_message) + except Exception as e: + verbose_logger.warning( + f"A2A sub-agent call to '{target_agent_id}' failed: {e}" + ) + reply = f"Error calling agent '{target_agent_id}': {e}" + messages.append( + { + "role": "tool", + "tool_call_id": tool_call.id, + "content": reply, + } + ) + + # If the loop ended on a tool-calling turn (no prose), make one tool-free + # call so the agent summarises instead of returning an empty message. + final_message = response.choices[0].message if response is not None else None + final_text = getattr(final_message, "content", None) if final_message else None + if not (final_text or "").strip(): + response = await litellm.acompletion( + model=model, + messages=messages, + stream=False, + **completion_kwargs, + ) + + return response diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index a3502f21f95..fec01cd74de 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -33,11 +33,73 @@ _AGENT_ONLY_PARAMS = frozenset( "agent_name", "agent_id", "agent_card_params", + "enable_mcp_tools", + "user_api_key_auth", + "callable_agents", + "call_agent", A2A_USER_API_KEY_HASH_PARAM, } ) +def _mcp_tools_for_agent(litellm_params: dict[str, Any]) -> Optional[list]: + """Expose the proxy MCP tool for an MCP-entitled agent. Which servers the + agent can actually reach is enforced downstream by ``user_api_key_auth``.""" + if not litellm_params.get("enable_mcp_tools"): + return None + return [{"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"}] + + +async def _run_completion( + completion_params: dict[str, Any], + litellm_params: dict[str, Any], +): + """Run the model call, routing to the agentic loop when the agent can call + other agents, to ``acompletion_with_mcp`` when it only has MCP tools, or a + plain completion otherwise.""" + callable_agents = litellm_params.get("callable_agents") + call_agent = litellm_params.get("call_agent") + if callable_agents and call_agent is not None: + from litellm.a2a_protocol.litellm_completion_bridge.agentic_loop import ( + run_agentic_loop, + ) + + loop_kwargs = dict(completion_params) + model = loop_kwargs.pop("model") + messages = loop_kwargs.pop("messages") + loop_kwargs.pop("stream", None) + # The loop owns the tool list (MCP tools + agent-as-tools), so drop any + # tools/tool_choice carried in from the agent's litellm_params or they + # collide with the explicit tools= the loop passes. + loop_kwargs.pop("tools", None) + loop_kwargs.pop("tool_choice", None) + return await run_agentic_loop( + model=model, + messages=messages, + completion_kwargs=loop_kwargs, + callable_agents=callable_agents, + call_agent=call_agent, + user_api_key_auth=litellm_params.get("user_api_key_auth"), + enable_mcp_tools=bool(litellm_params.get("enable_mcp_tools")), + ) + + mcp_tools = _mcp_tools_for_agent(litellm_params) + if mcp_tools is None: + return await litellm.acompletion(**completion_params) + + from litellm.responses.mcp.chat_completions_handler import acompletion_with_mcp + + # Same collision guard: acompletion_with_mcp passes its own tools= explicitly. + mcp_kwargs = { + k: v for k, v in completion_params.items() if k not in ("tools", "tool_choice") + } + return await acompletion_with_mcp( + tools=mcp_tools, + user_api_key_auth=litellm_params.get("user_api_key_auth"), + **mcp_kwargs, + ) + + class A2ACompletionBridgeHandler: """ Static methods for handling A2A requests via LiteLLM completion. @@ -46,13 +108,13 @@ class A2ACompletionBridgeHandler: @staticmethod async def handle_non_streaming( request_id: str, - params: Dict[str, Any], - litellm_params: Dict[str, Any], + params: dict[str, Any], + litellm_params: dict[str, Any], api_base: Optional[str] = None, - agent_extra_headers: Optional[Dict[str, str]] = None, + agent_extra_headers: Optional[dict[str, str]] = None, *, _skip_a2a_provider_routing: bool = False, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """ Handle non-streaming A2A request via litellm.acompletion. @@ -139,8 +201,8 @@ class A2ACompletionBridgeHandler: static_headers=completion_params.get("extra_headers"), ) - # Call litellm.acompletion - response = await litellm.acompletion(**completion_params) + # Call the model (auto-executing MCP tools in a loop if the agent has any) + response = await _run_completion(completion_params, litellm_params) # Transform response to A2A format a2a_response = ( @@ -278,24 +340,31 @@ class A2ACompletionBridgeHandler: ) yield working_event - # Call litellm.acompletion with streaming - response = await litellm.acompletion(**completion_params) + # Call the model with streaming (auto-executing MCP tools if the agent has any) + response = await _run_completion(completion_params, litellm_params) - # 3. Accumulate content and emit artifact update + # 3. Accumulate content and emit artifact update. The agentic loop / MCP + # tool execution resolve to a single non-streamed response; only a plain + # completion yields an async stream, so handle both shapes. accumulated_text = "" chunk_count = 0 - async for chunk in response: # type: ignore[union-attr] - chunk_count += 1 + if hasattr(response, "__aiter__"): + async for chunk in response: # type: ignore[union-attr] + chunk_count += 1 - # Extract delta content - content = "" - if chunk is not None and hasattr(chunk, "choices") and chunk.choices: - choice = chunk.choices[0] - if hasattr(choice, "delta") and choice.delta: - content = choice.delta.content or "" + content = "" + if chunk is not None and hasattr(chunk, "choices") and chunk.choices: + choice = chunk.choices[0] + if hasattr(choice, "delta") and choice.delta: + content = choice.delta.content or "" - if content: - accumulated_text += content + if content: + accumulated_text += content + else: + chunk_count = 1 + choices = getattr(response, "choices", None) + if choices: + accumulated_text = getattr(choices[0].message, "content", None) or "" # Emit artifact update with accumulated content if accumulated_text: diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 7446f61ad1c..41310947b57 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -6,6 +6,7 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM """ import json +from collections.abc import Awaitable, Callable from typing import Any, AsyncGenerator, Dict, List, Optional from urllib.parse import urlparse @@ -494,6 +495,144 @@ async def get_agent_card( raise HTTPException(status_code=500, detail=str(e)) +def _resolve_model_and_provider( + litellm_params: dict[str, Any], + model_list: Optional[list[dict[str, Any]]], +) -> tuple[dict[str, Any], Optional[str]]: + """A completion-bridge agent needs a provider-prefixed model and a + custom_llm_provider. A UI-created agent may carry a bare proxy model-list + alias (e.g. "gemini-flash") with no provider; resolve the alias against the + model list, then derive the provider from the "provider/model" prefix so the + bridge can route without a URL.""" + custom_llm_provider = litellm_params.get("custom_llm_provider") + if custom_llm_provider: + return litellm_params, custom_llm_provider + + agent_model = litellm_params.get("model") + if isinstance(agent_model, str) and "/" not in agent_model: + for deployment in model_list or []: + if deployment.get("model_name") == agent_model: + resolved = (deployment.get("litellm_params") or {}).get("model") + if isinstance(resolved, str): + agent_model = resolved + break + litellm_params = {**litellm_params, "model": agent_model} + + if isinstance(agent_model, str) and "/" in agent_model: + custom_llm_provider = agent_model.split("/", 1)[0] + litellm_params = {**litellm_params, "custom_llm_provider": custom_llm_provider} + + return litellm_params, custom_llm_provider + + +def _scope_auth_to_agent(user_api_key_auth: Any, agent_id: str) -> Any: + """Bind the caller's auth to ``agent_id`` so the MCP permission handler applies + the agent's own object_permission (caller perms intersected with the agent's + allowlist). Without ``agent_id`` set only the caller's key/team perms gate MCP, + which would let an agent reach tools its own permissions exclude. Returns a + copy so the shared caller auth object is never mutated.""" + if user_api_key_auth is not None and hasattr(user_api_key_auth, "model_copy"): + return user_api_key_auth.model_copy(update={"agent_id": agent_id}) + return user_api_key_auth + + +async def _source_agent_allowed_targets(source_auth: Any) -> set | None: + """Agent ids the source agent itself may call, from its own + object_permission.agents plus agents resolved from its access groups. None + means the source agent has no agent restriction. Delegation targets are + intersected with this (not only the caller's perms), so a prompt-injected + delegation cannot exceed what the source agent was granted.""" + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + AgentRequestHandler, + ) + + obj_perm = await MCPRequestHandler._get_agent_object_permission( + user_api_key_auth=source_auth + ) + if obj_perm is None: + return None + + allowed = set(getattr(obj_perm, "agents", None) or []) + allowed |= set( + await AgentRequestHandler._get_agents_from_access_groups( + getattr(obj_perm, "agent_access_groups", None) or [] + ) + ) + return allowed or None + + +async def _delegate_to_agent( + *, + target_id: str, + message: str, + allowed_agent_ids: set, + get_agent: Callable[[str], Any], + user_api_key_auth: Any, + send_message: Callable[..., Awaitable[Any]], +) -> str: + """Dispatch a sub-agent call and return its text reply. + + The target id comes from a model tool call, so it is attacker-influenceable + via prompt injection, and ``send_message`` bypasses the endpoint's auth gate. + The auth-filtered ``allowed_agent_ids`` whitelist is therefore re-enforced + here so a caller can never reach an agent it was not granted.""" + if target_id not in allowed_agent_ids: + return f"Agent '{target_id}' is not available." + + target = get_agent(target_id) + if target is None: + return f"Agent '{target_id}' not found." + + # The target's proxy-side guardrails are enforced by the endpoint request + # pipeline, which this direct dispatch bypasses. Rather than return output a + # guardrail would have blocked on a direct call, refuse to delegate to a + # guarded agent (it can still be invoked directly). + if (target.litellm_params or {}).get("guardrails"): + return f"Agent '{target_id}' cannot be called indirectly." + + from uuid import uuid4 + + from a2a.types import MessageSendParams, SendMessageRequest + + target_card = target.agent_card_params or {} + # Strip delegation flags so a sub-agent cannot recurse into further + # agent-to-agent calls (one level of delegation). + target_params = { + k: v + for k, v in (target.litellm_params or {}).items() + if k not in ("enable_agent_calls", "callable_agents", "call_agent") + } + target_params["user_api_key_auth"] = _scope_auth_to_agent( + user_api_key_auth, target.agent_id + ) + sub_request = SendMessageRequest( + id="", + params=MessageSendParams( + message={ + "role": "user", + "parts": [{"kind": "text", "text": message}], + "messageId": uuid4().hex, + } + ), + ) + sub_response = await send_message( + request=sub_request, + api_base=target_card.get("url"), + litellm_params=target_params, + agent_id=target.agent_id, + ) + dumped = ( + sub_response.model_dump() + if hasattr(sub_response, "model_dump") + else sub_response + ) + parts = ((dumped or {}).get("result") or {}).get("parts") or [] + return "\n".join(p.get("text", "") for p in parts if p.get("kind") == "text") + + @router.post( "/a2a/{agent_id}", tags=["[beta] A2A Agents"], @@ -594,7 +733,12 @@ async def invoke_agent_a2a( # Get litellm_params (may include custom_llm_provider for completion bridge) litellm_params = agent.litellm_params or {} - custom_llm_provider = litellm_params.get("custom_llm_provider") + from litellm.proxy.proxy_server import llm_router + + litellm_params, custom_llm_provider = _resolve_model_and_provider( + litellm_params, + llm_router.model_list if llm_router is not None else None, + ) # Hand the authenticated key hash to the completion bridge so provider # configs can scope provider-side session state per key (e.g. LangFlow @@ -609,6 +753,72 @@ async def invoke_agent_a2a( A2A_USER_API_KEY_HASH_PARAM: user_api_key_dict.api_key, } + # MCP tools and/or agent-to-agent delegation both run inside the completion + # bridge; hand the caller's auth bound to this agent so MCP tools are scoped + # to the caller's perms intersected with the agent's own allowlist. + if custom_llm_provider and ( + litellm_params.get("enable_mcp_tools") + or litellm_params.get("enable_agent_calls") + ): + litellm_params = { + **litellm_params, + "user_api_key_auth": _scope_auth_to_agent( + user_api_key_dict, agent.agent_id + ), + } + + # Agents-as-tools: expose the other agents this caller may reach, and + # dispatch each call as a real A2A message/send to the target agent. + if custom_llm_provider and litellm_params.get("enable_agent_calls"): + from litellm.a2a_protocol import asend_message + from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry, + ) + + source_allowed_targets = await _source_agent_allowed_targets( + _scope_auth_to_agent(user_api_key_dict, agent.agent_id) + ) + + callable_agents: list[dict[str, Any]] = [] + for other in global_agent_registry.get_agent_list(): + if other.agent_id == agent.agent_id: + continue + if not await AgentRequestHandler.is_agent_allowed( + agent_id=other.agent_id, user_api_key_auth=user_api_key_dict + ): + continue + if ( + source_allowed_targets is not None + and other.agent_id not in source_allowed_targets + ): + continue + other_card = other.agent_card_params or {} + callable_agents.append( + { + "agent_id": other.agent_id, + "name": other_card.get("name", other.agent_id), + "description": other_card.get("description", ""), + } + ) + + allowed_agent_ids = {a["agent_id"] for a in callable_agents} + + async def _call_agent(target_id: str, message: str) -> str: + return await _delegate_to_agent( + target_id=target_id, + message=message, + allowed_agent_ids=allowed_agent_ids, + get_agent=_get_agent, + user_api_key_auth=user_api_key_dict, + send_message=asend_message, + ) + + litellm_params = { + **litellm_params, + "callable_agents": callable_agents, + "call_agent": _call_agent, + } + # URL is required unless using completion bridge with a provider that derives endpoint from model # (e.g., bedrock/agentcore derives endpoint from ARN in model string) if not agent_url and not custom_llm_provider: diff --git a/tests/test_litellm/a2a_protocol/test_agentic_loop.py b/tests/test_litellm/a2a_protocol/test_agentic_loop.py new file mode 100644 index 00000000000..405525c946a --- /dev/null +++ b/tests/test_litellm/a2a_protocol/test_agentic_loop.py @@ -0,0 +1,312 @@ +"""Tests for the completion-bridge agentic loop and its completion routing. + +Covers the helpers the loop is built on (turning MCP tool results into chat tool +messages, exposing other A2A agents as function tools), the loop driver itself +(MCP tool execution, agent delegation, and the non-empty-answer guarantee), and +the routing in ``_run_completion``. +""" + +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import pytest + +from litellm.a2a_protocol.litellm_completion_bridge import handler +from litellm.a2a_protocol.litellm_completion_bridge.agentic_loop import ( + AGENT_TOOL_PREFIX, + _build_agent_tools, + _mcp_results_to_tool_messages, + run_agentic_loop, +) + + +def _tool_call(call_id: str, name: str, arguments: str = "{}"): + return SimpleNamespace( + id=call_id, function=SimpleNamespace(name=name, arguments=arguments) + ) + + +def _response(content=None, tool_calls=None): + message = SimpleNamespace( + content=content, + tool_calls=tool_calls, + model_dump=lambda: {"role": "assistant", "content": content}, + ) + return SimpleNamespace(choices=[SimpleNamespace(message=message)]) + + +def test_mcp_results_become_tool_role_messages(): + """Regression: ``_execute_tool_calls`` returns ``{tool_call_id, result, + name}``, which is not a valid chat message. Feeding that back verbatim left + the model blind to the result, so it re-called the tool until the loop was + exhausted and produced an empty answer. The result must be reshaped into a + ``tool`` role message whose ``content`` holds the result text.""" + raw = [{"tool_call_id": "call_1", "result": "sunny, 21C", "name": "get_weather"}] + + messages = _mcp_results_to_tool_messages(raw) + + assert messages == ( + {"role": "tool", "tool_call_id": "call_1", "content": "sunny, 21C"}, + ) + # The original keys must not survive: leaking "result"/"name" or dropping + # "role" is exactly the bug. + assert "result" not in messages[0] + assert "name" not in messages[0] + + +def test_mcp_results_coerce_non_string_results_and_preserve_order(): + raw = [ + {"tool_call_id": "a", "result": {"city": "Paris", "temp_c": 21}}, + {"tool_call_id": "b", "result": None}, + ] + + messages = _mcp_results_to_tool_messages(raw) + + assert [m["tool_call_id"] for m in messages] == ["a", "b"] + assert all(m["role"] == "tool" for m in messages) + assert isinstance(messages[0]["content"], str) + assert "Paris" in messages[0]["content"] + + +def test_build_agent_tools_one_function_tool_per_agent(): + tools = _build_agent_tools( + [ + { + "agent_id": "translator", + "name": "Translator", + "description": "Translates text", + }, + {"agent_id": "summariser"}, + ] + ) + + assert len(tools) == 2 + first = tools[0]["function"] + assert first["name"] == f"{AGENT_TOOL_PREFIX}translator" + assert first["description"] == "Translates text" + # The model must supply a self-contained message for the sub-agent. + assert first["parameters"]["required"] == ["message"] + assert first["parameters"]["properties"]["message"]["type"] == "string" + # An agent without its own description still gets a usable default. + assert "summariser" in tools[1]["function"]["description"] + + +def test_build_agent_tools_skips_agents_without_id(): + tools = _build_agent_tools([{"name": "no id here"}, {"id": "kept"}]) + + names = [t["function"]["name"] for t in tools] + assert names == [f"{AGENT_TOOL_PREFIX}kept"] + + +@pytest.mark.asyncio +async def test_loop_delegates_to_agent_then_answers(): + """An agent tool call routes through call_agent, the reply is fed back as a + tool message, and the loop returns the model's final prose.""" + call_agent = AsyncMock(return_value="bonjour") + with patch("litellm.acompletion", new_callable=AsyncMock) as acompletion: + acompletion.side_effect = [ + _response( + tool_calls=[ + _tool_call( + "c1", + f"{AGENT_TOOL_PREFIX}translator", + json.dumps({"message": "say hi in french"}), + ) + ] + ), + _response(content="done"), + ] + result = await run_agentic_loop( + model="gemini/gemini-flash-latest", + messages=[{"role": "user", "content": "hi"}], + completion_kwargs={}, + callable_agents=[{"agent_id": "translator", "name": "Translator"}], + call_agent=call_agent, + ) + + assert result.choices[0].message.content == "done" + call_agent.assert_awaited_once_with("translator", "say hi in french") + fed_back = acompletion.await_args_list[1].kwargs["messages"] + assert {"role": "tool", "tool_call_id": "c1", "content": "bonjour"} in fed_back + + +@pytest.mark.asyncio +async def test_loop_executes_mcp_tool_and_feeds_result_back(): + """Regression at the loop level: an MCP tool result is reshaped to a tool + message so the model sees it and answers instead of re-calling the tool.""" + handler_path = ( + "litellm.responses.mcp.litellm_proxy_mcp_handler.LiteLLM_Proxy_MCP_Handler" + ) + with ( + patch("litellm.acompletion", new_callable=AsyncMock) as acompletion, + patch( + "litellm.a2a_protocol.litellm_completion_bridge.agentic_loop._load_mcp_tools", + new_callable=AsyncMock, + ) as load_mcp, + patch( + f"{handler_path}._execute_tool_calls", new_callable=AsyncMock + ) as execute_tool_calls, + ): + acompletion.side_effect = [ + _response(tool_calls=[_tool_call("t1", "get_weather")]), + _response(content="It is sunny"), + ] + load_mcp.return_value = ( + [{"type": "function", "function": {"name": "get_weather"}}], + {"get_weather": "server-1"}, + ) + execute_tool_calls.return_value = [ + {"tool_call_id": "t1", "result": "sunny", "name": "get_weather"} + ] + result = await run_agentic_loop( + model="gemini/gemini-flash-latest", + messages=[{"role": "user", "content": "weather?"}], + completion_kwargs={}, + callable_agents=[], + call_agent=AsyncMock(), + enable_mcp_tools=True, + ) + + assert result.choices[0].message.content == "It is sunny" + fed_back = acompletion.await_args_list[1].kwargs["messages"] + assert {"role": "tool", "tool_call_id": "t1", "content": "sunny"} in fed_back + + +@pytest.mark.asyncio +async def test_loop_makes_tool_free_final_call_when_last_turn_has_no_prose(): + """If the loop is exhausted on a tool-calling turn, one tool-free call is made + so the agent answers instead of returning an empty message.""" + with patch("litellm.acompletion", new_callable=AsyncMock) as acompletion: + acompletion.side_effect = [ + _response(tool_calls=[_tool_call("a1", f"{AGENT_TOOL_PREFIX}t")]), + _response(content="final summary"), + ] + result = await run_agentic_loop( + model="m", + messages=[{"role": "user", "content": "x"}], + completion_kwargs={}, + callable_agents=[{"agent_id": "t"}], + call_agent=AsyncMock(return_value="r"), + max_iterations=1, + ) + + assert result.choices[0].message.content == "final summary" + assert acompletion.await_count == 2 + # the final call is tool-free + assert "tools" not in acompletion.await_args_list[1].kwargs + + +@pytest.mark.asyncio +async def test_loop_surfaces_sub_agent_failure_without_crashing(): + """If a delegated agent call raises, the loop records the error as the tool + reply and still returns a final answer instead of propagating the exception.""" + call_agent = AsyncMock(side_effect=RuntimeError("boom")) + with patch("litellm.acompletion", new_callable=AsyncMock) as acompletion: + acompletion.side_effect = [ + _response(tool_calls=[_tool_call("c1", f"{AGENT_TOOL_PREFIX}flaky")]), + _response(content="handled"), + ] + result = await run_agentic_loop( + model="m", + messages=[{"role": "user", "content": "x"}], + completion_kwargs={}, + callable_agents=[{"agent_id": "flaky"}], + call_agent=call_agent, + ) + + assert result.choices[0].message.content == "handled" + fed_back = acompletion.await_args_list[1].kwargs["messages"] + assert any( + m.get("role") == "tool" + and "Error calling agent 'flaky'" in str(m.get("content")) + for m in fed_back + ) + + +@pytest.mark.asyncio +async def test_run_completion_routes_to_agentic_loop_when_agent_can_delegate(): + with patch( + "litellm.a2a_protocol.litellm_completion_bridge.agentic_loop.run_agentic_loop", + new_callable=AsyncMock, + ) as loop: + loop.return_value = "LOOP" + result = await handler._run_completion( + {"model": "m", "messages": [{"role": "user", "content": "x"}]}, + {"callable_agents": [{"agent_id": "a"}], "call_agent": AsyncMock()}, + ) + + assert result == "LOOP" + loop.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_run_completion_routes_to_mcp_when_only_tools_enabled(): + with patch( + "litellm.responses.mcp.chat_completions_handler.acompletion_with_mcp", + new_callable=AsyncMock, + ) as acompletion_with_mcp: + acompletion_with_mcp.return_value = "MCP" + result = await handler._run_completion( + {"model": "m", "messages": []}, {"enable_mcp_tools": True} + ) + + assert result == "MCP" + acompletion_with_mcp.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_run_completion_plain_when_no_tools_or_agents(): + with patch("litellm.acompletion", new_callable=AsyncMock) as acompletion: + acompletion.return_value = "PLAIN" + result = await handler._run_completion({"model": "m", "messages": []}, {}) + + assert result == "PLAIN" + acompletion.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_run_completion_strips_agent_tools_before_agentic_loop(): + """Regression: an agent carrying `tools`/`tool_choice` in its litellm_params + must not reach the loop's completion_kwargs, or they collide with the explicit + tools= the loop passes (TypeError: got multiple values for 'tools').""" + with patch( + "litellm.a2a_protocol.litellm_completion_bridge.agentic_loop.run_agentic_loop", + new_callable=AsyncMock, + ) as loop: + loop.return_value = "LOOP" + await handler._run_completion( + { + "model": "m", + "messages": [], + "tools": [{"type": "function", "function": {"name": "x"}}], + "tool_choice": "auto", + }, + {"callable_agents": [{"agent_id": "a"}], "call_agent": AsyncMock()}, + ) + + completion_kwargs = loop.await_args.kwargs["completion_kwargs"] + assert "tools" not in completion_kwargs + assert "tool_choice" not in completion_kwargs + + +@pytest.mark.asyncio +async def test_run_completion_strips_agent_tools_on_mcp_path(): + """Same collision guard on the MCP path: acompletion_with_mcp passes tools= + explicitly, so a tools key from litellm_params must be dropped (else the call + raises before the request is even made).""" + with patch( + "litellm.responses.mcp.chat_completions_handler.acompletion_with_mcp", + new_callable=AsyncMock, + ) as acompletion_with_mcp: + acompletion_with_mcp.return_value = "MCP" + result = await handler._run_completion( + { + "model": "m", + "messages": [], + "tools": [{"type": "function", "function": {"name": "x"}}], + }, + {"enable_mcp_tools": True}, + ) + + assert result == "MCP" diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py index 07e878401e0..f61064ac635 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py @@ -1560,3 +1560,187 @@ async def test_caller_identity_headers_cannot_be_spoofed_via_forwarded_headers() assert ( posted_headers.get("X-LiteLLM-Team-Id") == "real-team" ), "authenticated team id must not be overridden by forwarded client headers" + + +def test_resolve_model_and_provider_resolves_bare_alias(): + from litellm.proxy.agent_endpoints.a2a_endpoints import _resolve_model_and_provider + + model_list = [ + { + "model_name": "gemini-flash", + "litellm_params": {"model": "gemini/gemini-flash-latest"}, + } + ] + params, provider = _resolve_model_and_provider( + {"model": "gemini-flash"}, model_list + ) + assert params["model"] == "gemini/gemini-flash-latest" + assert provider == "gemini" + + +def test_resolve_model_and_provider_keeps_explicit_provider(): + from litellm.proxy.agent_endpoints.a2a_endpoints import _resolve_model_and_provider + + params, provider = _resolve_model_and_provider( + {"model": "anything", "custom_llm_provider": "openai"}, None + ) + assert provider == "openai" + assert params["model"] == "anything" + + +def test_resolve_model_and_provider_derives_provider_from_prefixed_model(): + from litellm.proxy.agent_endpoints.a2a_endpoints import _resolve_model_and_provider + + params, provider = _resolve_model_and_provider( + {"model": "anthropic/claude-x"}, None + ) + assert provider == "anthropic" + + +def test_resolve_model_and_provider_leaves_unknown_alias_unresolved(): + from litellm.proxy.agent_endpoints.a2a_endpoints import _resolve_model_and_provider + + params, provider = _resolve_model_and_provider({"model": "mystery"}, []) + assert provider is None + assert params["model"] == "mystery" + + +@pytest.mark.asyncio +async def test_delegate_to_agent_rejects_target_outside_whitelist(): + """Security regression: the target id is model-supplied (prompt-injectable) + and send_message bypasses the endpoint auth gate, so an agent id outside the + auth-filtered whitelist must be refused before any dispatch or registry + lookup.""" + from litellm.proxy.agent_endpoints.a2a_endpoints import _delegate_to_agent + + send_message = AsyncMock() + get_agent = MagicMock() + + reply = await _delegate_to_agent( + target_id="evil-agent", + message="hi", + allowed_agent_ids={"allowed-agent"}, + get_agent=get_agent, + user_api_key_auth=None, + send_message=send_message, + ) + + assert "not available" in reply + send_message.assert_not_awaited() + get_agent.assert_not_called() + + +@pytest.mark.asyncio +async def test_delegate_to_agent_does_not_dispatch_when_target_missing(): + from litellm.proxy.agent_endpoints.a2a_endpoints import _delegate_to_agent + + send_message = AsyncMock() + + reply = await _delegate_to_agent( + target_id="known", + message="hi", + allowed_agent_ids={"known"}, + get_agent=lambda _id: None, + user_api_key_auth=None, + send_message=send_message, + ) + + assert "not found" in reply + send_message.assert_not_awaited() + + +def test_scope_auth_to_agent_binds_agent_id_without_mutating_caller(): + """The agent's MCP object_permission is only enforced when agent_id is set on + the auth, so the bridge must receive a copy bound to the agent; the shared + caller auth must not be mutated.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.agent_endpoints.a2a_endpoints import _scope_auth_to_agent + + caller = UserAPIKeyAuth(api_key="sk-x", user_id="u") + scoped = _scope_auth_to_agent(caller, "agent-123") + + assert scoped.agent_id == "agent-123" + assert caller.agent_id != "agent-123" + + +@pytest.mark.asyncio +async def test_delegate_to_agent_refuses_target_with_guardrails(): + """A target with proxy-side guardrails must not be reachable via delegation: + this path bypasses the guardrail/post-call enforcement a direct call runs, so + it is refused rather than returning output a guardrail would have blocked.""" + from types import SimpleNamespace + + from litellm.proxy.agent_endpoints.a2a_endpoints import _delegate_to_agent + + send_message = AsyncMock() + guarded = SimpleNamespace( + agent_id="guarded", + agent_card_params={"url": "http://x"}, + litellm_params={"guardrails": ["pii-mask"]}, + ) + + reply = await _delegate_to_agent( + target_id="guarded", + message="hi", + allowed_agent_ids={"guarded"}, + get_agent=lambda _id: guarded, + user_api_key_auth=None, + send_message=send_message, + ) + + assert "cannot be called indirectly" in reply + send_message.assert_not_awaited() + + +_MCP_OBJ_PERM = ( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp." + "MCPRequestHandler._get_agent_object_permission" +) +_ACCESS_GROUPS = ( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler." + "AgentRequestHandler._get_agents_from_access_groups" +) + + +@pytest.mark.asyncio +async def test_source_agent_targets_unrestricted_without_object_permission(): + from litellm.proxy.agent_endpoints.a2a_endpoints import ( + _source_agent_allowed_targets, + ) + + with patch(_MCP_OBJ_PERM, new_callable=AsyncMock, return_value=None): + assert await _source_agent_allowed_targets(MagicMock()) is None + + +@pytest.mark.asyncio +async def test_source_agent_targets_intersect_agents_and_access_groups(): + """Regression: delegation must be limited to the source agent's own allowed + agents (object_permission.agents + access-group agents), not just the caller's.""" + from types import SimpleNamespace + + from litellm.proxy.agent_endpoints.a2a_endpoints import ( + _source_agent_allowed_targets, + ) + + obj_perm = SimpleNamespace(agents=["a", "b"], agent_access_groups=["grp"]) + with ( + patch(_MCP_OBJ_PERM, new_callable=AsyncMock, return_value=obj_perm), + patch(_ACCESS_GROUPS, new_callable=AsyncMock, return_value=["c"]), + ): + assert await _source_agent_allowed_targets(MagicMock()) == {"a", "b", "c"} + + +@pytest.mark.asyncio +async def test_source_agent_targets_empty_permission_is_unrestricted(): + from types import SimpleNamespace + + from litellm.proxy.agent_endpoints.a2a_endpoints import ( + _source_agent_allowed_targets, + ) + + obj_perm = SimpleNamespace(agents=[], agent_access_groups=[]) + with ( + patch(_MCP_OBJ_PERM, new_callable=AsyncMock, return_value=obj_perm), + patch(_ACCESS_GROUPS, new_callable=AsyncMock, return_value=[]), + ): + assert await _source_agent_allowed_targets(MagicMock()) is None diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/ModelConfigCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/ModelConfigCard.tsx index 66ddb90bea3..62009bef4a7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/ModelConfigCard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/ModelConfigCard.tsx @@ -28,7 +28,14 @@ const ModelConfigCard: React.FC = ({ return (
- +
) : agentType === "a2a" ? ( - + ) : selectedAgentTypeInfo?.use_a2a_form_fields ? ( <> - + {selectedAgentTypeInfo.credential_fields.length > 0 && (

diff --git a/ui/litellm-dashboard/src/components/agents/agent_config.ts b/ui/litellm-dashboard/src/components/agents/agent_config.ts index 14b6729bf93..97b33873b32 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_config.ts +++ b/ui/litellm-dashboard/src/components/agents/agent_config.ts @@ -6,7 +6,7 @@ export interface FieldConfig { name: string; label: string; - type: "text" | "textarea" | "url" | "switch" | "list"; + type: "text" | "textarea" | "url" | "switch" | "list" | "select"; required?: boolean; tooltip?: string; placeholder?: string; @@ -146,6 +146,25 @@ export const AGENT_FORM_CONFIG: { label: "Make Public", type: "switch", }, + { + name: "prompt_id", + label: "Prompt", + type: "select", + placeholder: "Select a stored prompt to apply (optional)", + tooltip: "The agent answers in this prompt's persona. Create prompts under Experimental > Prompts.", + }, + { + name: "enable_mcp_tools", + label: "Enable MCP Tools", + type: "switch", + tooltip: "Let the agent call the MCP tools it is entitled to, executing them in a loop.", + }, + { + name: "enable_agent_calls", + label: "Enable Agent-to-Agent Calls", + type: "switch", + tooltip: "Let the agent delegate to other agents it can access (agents-as-tools over A2A).", + }, ], }, cost: { @@ -275,6 +294,9 @@ export const buildAgentDataFromForm = (values: any, existingAgent?: any) => { if (values.model) params.model = values.model; if (values.make_public !== undefined) params.make_public = values.make_public; + if (values.prompt_id) params.prompt_id = values.prompt_id; + if (values.enable_mcp_tools !== undefined) params.enable_mcp_tools = values.enable_mcp_tools; + if (values.enable_agent_calls !== undefined) params.enable_agent_calls = values.enable_agent_calls; if (values.cost_per_query) params.cost_per_query = parseFloat(values.cost_per_query); if (values.input_cost_per_token) params.input_cost_per_token = parseFloat(values.input_cost_per_token); if (values.output_cost_per_token) params.output_cost_per_token = parseFloat(values.output_cost_per_token); @@ -333,6 +355,9 @@ export const parseAgentForForm = (agent: any) => { documentationUrl: agent.agent_card_params?.documentationUrl, supportsAuthenticatedExtendedCard: agent.agent_card_params?.supportsAuthenticatedExtendedCard, model: agent.litellm_params?.model, + prompt_id: agent.litellm_params?.prompt_id, + enable_mcp_tools: agent.litellm_params?.enable_mcp_tools, + enable_agent_calls: agent.litellm_params?.enable_agent_calls, make_public: agent.litellm_params?.make_public, cost_per_query: agent.litellm_params?.cost_per_query, input_cost_per_token: agent.litellm_params?.input_cost_per_token, diff --git a/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx b/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx index 7965af6322d..50ea12dc935 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx +++ b/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx @@ -11,13 +11,18 @@ const { Panel } = Collapse; interface AgentFormFieldsProps { showAgentName?: boolean; visiblePanels?: string[]; + promptOptions?: string[]; } /** * Reusable form fields component for agent forms * Uses shared configuration from agent_config.ts */ -const AgentFormFields: React.FC = ({ showAgentName = true, visiblePanels }) => { +const AgentFormFields: React.FC = ({ + showAgentName = true, + visiblePanels, + promptOptions = [], +}) => { const shouldShow = (key: string) => !visiblePanels || visiblePanels.includes(key); return ( <> @@ -179,8 +184,21 @@ const AgentFormFields: React.FC = ({ showAgentName = true, label={field.label} name={field.name} valuePropName={field.type === "switch" ? "checked" : undefined} + tooltip={field.tooltip} > - {field.type === "switch" ? : } + {field.type === "switch" ? ( + + ) : field.type === "select" ? ( + + )} ))} diff --git a/ui/litellm-dashboard/src/components/agents/agent_info.tsx b/ui/litellm-dashboard/src/components/agents/agent_info.tsx index d4420931a20..81ef98d09f7 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_info.tsx +++ b/ui/litellm-dashboard/src/components/agents/agent_info.tsx @@ -3,7 +3,7 @@ import { Card, Title, Text, Button as TremorButton, Tab, TabGroup, TabList, TabP import { Form, Input, InputNumber, Button as AntButton, Spin, Descriptions, Divider } from "antd"; import MessageManager from "@/components/molecules/message_manager"; import { ArrowLeftIcon } from "@heroicons/react/outline"; -import { getAgentInfo, patchAgentCall, getAgentCreateMetadata, AgentCreateInfo } from "../networking"; +import { getAgentInfo, patchAgentCall, getAgentCreateMetadata, AgentCreateInfo, getPromptsList } from "../networking"; import { Agent } from "./types"; import AgentFormFields from "./agent_form_fields"; import DynamicAgentFormFields, { buildDynamicAgentData } from "./dynamic_agent_form_fields"; @@ -28,6 +28,7 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT const [form] = Form.useForm(); const [agentTypeMetadata, setAgentTypeMetadata] = useState([]); const [detectedAgentType, setDetectedAgentType] = useState("a2a"); + const [promptOptions, setPromptOptions] = useState([]); const [appliedDiscoveredSelection, setAppliedDiscoveredSelection] = useState( null, ); @@ -44,6 +45,13 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT fetchMetadata(); }, []); + useEffect(() => { + if (!accessToken) return; + getPromptsList(accessToken) + .then((res) => setPromptOptions(res.prompts.map((p) => p.prompt_id))) + .catch((error) => console.error("Error fetching prompts:", error)); + }, [accessToken]); + useEffect(() => { fetchAgentInfo(); }, [agentId, accessToken]); @@ -350,11 +358,11 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT {detectedAgentType === "a2a" ? ( - + ) : selectedAgentTypeInfo ? ( ) : ( - + )} {discoveryRequest && ( diff --git a/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx b/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx index 060a8f0f111..6c9b2441a86 100644 --- a/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx +++ b/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx @@ -14,6 +14,7 @@ interface ModelSelectorProps { className?: string; showLabel?: boolean; labelText?: string; + allowClear?: boolean; } const ModelSelector: React.FC = ({ @@ -26,6 +27,7 @@ const ModelSelector: React.FC = ({ className, showLabel = true, labelText = "Select Model", + allowClear = false, }) => { const [selectedModel, setSelectedModel] = useState(value); const [showCustomModelInput, setShowCustomModelInput] = useState(false); @@ -55,16 +57,16 @@ const ModelSelector: React.FC = ({ loadModels(); }, [accessToken]); - const onModelChange = (value: string) => { + const onModelChange = (value?: string) => { if (value === "custom") { setShowCustomModelInput(true); setSelectedModel(undefined); - } else { - setShowCustomModelInput(false); - setSelectedModel(value); - if (onChange) { - onChange(value); - } + return; + } + setShowCustomModelInput(false); + setSelectedModel(value); + if (onChange) { + onChange(value ?? ""); } }; @@ -103,6 +105,7 @@ const ModelSelector: React.FC = ({ ]} style={{ width: "100%", ...style }} showSearch={true} + allowClear={allowClear} className={`rounded-md ${className || ""}`} disabled={disabled} /> From e391a28550578017b036141a3a51fbfa7382855a Mon Sep 17 00:00:00 2001 From: "A.Kebriaei" Date: Tue, 23 Jun 2026 09:41:09 +0100 Subject: [PATCH 2/3] docs(a2a): document the agentic loop and the empty-answer fix --- .../litellm_completion_bridge/README.md | 72 +++++++++++++++++++ 1 file changed, 72 insertions(+) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/README.md b/litellm/a2a_protocol/litellm_completion_bridge/README.md index a809e9bf55e..dec7aa335bd 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/README.md +++ b/litellm/a2a_protocol/litellm_completion_bridge/README.md @@ -72,3 +72,75 @@ When an A2A request hits `/a2a/{agent_id}/message/send`, the bridge: - `A2ACompletionBridgeTransformation` - Static methods for message format conversion - `A2ACompletionBridgeHandler` - Static methods for handling requests (streaming/non-streaming) +## Agentic loop: MCP tools and agent-to-agent delegation + +A plain bridge agent maps one A2A `message/send` to one `litellm.acompletion`, so +it can answer from the model alone but cannot use tools or other agents. +`agentic_loop.run_agentic_loop` turns that single call into a loop: the model is +offered the agent's tools, any tool or agent calls it makes are executed, their +results are fed back, and the loop repeats until the model stops asking for tools +and returns prose. + +Two kinds of tool are offered to the model: + +- MCP tools, when the agent has `enable_mcp_tools` set. The agent only ever sees + the servers its key is authorised for; entitlement is enforced downstream by + `user_api_key_auth`, not by the loop. +- One function tool per other agent the caller may reach, when the agent has + `enable_agent_calls` set. The tool name carries the target agent id, and a call + to it is dispatched as a real A2A `message/send` to that agent via + `asend_message`, so a delegated agent behaves exactly as if invoked directly. + Delegation is limited to one level so a sub-agent cannot recurse into further + hops. + +### The empty-answer bug this fixes + +`LiteLLM_Proxy_MCP_Handler._execute_tool_calls` returns each result as +`{tool_call_id, result, name}`, which is not a valid chat message. Feeding those +dicts straight back into the conversation left the model unable to see the +result, so it re-issued the same tool call on every iteration until the loop hit +its cap and returned an empty final answer. `_mcp_results_to_tool_messages` +reshapes each result into the `{role: "tool", tool_call_id, content}` form the +chat API expects, so the model sees the result, stops re-calling the tool, and +answers. The regression is covered in +`tests/test_litellm/a2a_protocol/test_agentic_loop.py`. + +### Configuring an agent + +Set these on the agent's `litellm_params` (via `POST /v1/agents`, the agent edit +API, or the dashboard agent create / Settings form): + +- `prompt_id`: a stored prompt whose persona the agent answers in. A prompt that + carries a model overrides the agent's model; a prompt with no model leaves the + agent's own model in place. Clear the model in the prompt editor (or omit the + `model:` line in the dotprompt) to have the agent use its own model. +- `enable_mcp_tools`: allow the agent to call its entitled MCP tools. +- `enable_agent_calls`: allow the agent to delegate to other agents it can reach. + +The agent's `model` may be provider-prefixed (`gemini/gemini-flash-latest`) or a +bare proxy model-list alias (`gemini-flash`); the endpoint resolves the alias +against the model list and derives the provider so the bridge can route without +an explicit URL. + +### Invoking an agent + +```bash +curl -sS -X POST "http://localhost:4000/a2a/$AGENT_ID" \ + -H "Authorization: Bearer $LITELLM_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "jsonrpc": "2.0", + "id": "1", + "method": "message/send", + "params": { + "message": { + "role": "user", + "parts": [{ "kind": "text", "text": "your question here" }], + "messageId": "m-1" + } + } + }' +``` + +The answer comes back in `result.parts[].text`. + From a45ae9a04d95f805960bf878853d6ae3350e2173 Mon Sep 17 00:00:00 2001 From: "A.Kebriaei" Date: Thu, 25 Jun 2026 14:08:52 +0100 Subject: [PATCH 3/3] fix(a2a): type auth params and bundle delegation context for the strict gates Replaces the Any annotations on the auth params with UserAPIKeyAuth (ANN401), makes _source_agent_allowed_targets accept an optional auth and bundles _delegate_to_agent's dispatch args into a frozen _DelegationContext (PLR0913), and casts the a2a message dict to Message (reportArgumentType). Behaviour is unchanged. --- .../litellm_completion_bridge/agentic_loop.py | 23 +++++-- .../proxy/agent_endpoints/a2a_endpoints.py | 62 ++++++++++++------- .../agent_endpoints/test_a2a_endpoints.py | 45 +++++++++----- 3 files changed, 89 insertions(+), 41 deletions(-) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/agentic_loop.py b/litellm/a2a_protocol/litellm_completion_bridge/agentic_loop.py index b64648ad71e..f712cbe8eec 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/agentic_loop.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/agentic_loop.py @@ -7,11 +7,23 @@ supplied by the proxy layer via ``call_agent``. """ import json -from typing import Any, Awaitable, Callable, Iterable, Mapping, Optional +from typing import ( + TYPE_CHECKING, + Any, + Awaitable, + Callable, + Iterable, + Mapping, + Optional, + cast, +) import litellm from litellm._logging import verbose_logger +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + AGENT_TOOL_PREFIX = "agent__" DEFAULT_MAX_ITERATIONS = 8 @@ -69,7 +81,7 @@ def _build_agent_tools(callable_agents: list[dict[str, Any]]) -> list[dict[str, async def _load_mcp_tools( - user_api_key_auth: Any, + user_api_key_auth: Optional["UserAPIKeyAuth"], litellm_trace_id: Optional[str], ): """Returns (openai_tools, tool_server_map) for the agent's allowed MCP servers.""" @@ -101,7 +113,7 @@ async def run_agentic_loop( completion_kwargs: dict[str, Any], callable_agents: list[dict[str, Any]], call_agent: Callable[[str, str], Awaitable[str]], - user_api_key_auth: Any = None, + user_api_key_auth: Optional["UserAPIKeyAuth"] = None, enable_mcp_tools: bool = False, max_iterations: int = DEFAULT_MAX_ITERATIONS, ): @@ -140,7 +152,10 @@ async def run_agentic_loop( break messages.append( - message.model_dump() if hasattr(message, "model_dump") else message + cast( + "dict[str, Any]", + message.model_dump() if hasattr(message, "model_dump") else message, + ) ) agent_calls = [ diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 41310947b57..a89bdcced9c 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -7,7 +7,8 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM import json from collections.abc import Awaitable, Callable -from typing import Any, AsyncGenerator, Dict, List, Optional +from dataclasses import dataclass +from typing import Any, AsyncGenerator, Dict, List, Optional, cast from urllib.parse import urlparse from fastapi import APIRouter, Depends, HTTPException, Request, Response @@ -525,7 +526,9 @@ def _resolve_model_and_provider( return litellm_params, custom_llm_provider -def _scope_auth_to_agent(user_api_key_auth: Any, agent_id: str) -> Any: +def _scope_auth_to_agent( + user_api_key_auth: Optional[UserAPIKeyAuth], agent_id: str +) -> Optional[UserAPIKeyAuth]: """Bind the caller's auth to ``agent_id`` so the MCP permission handler applies the agent's own object_permission (caller perms intersected with the agent's allowlist). Without ``agent_id`` set only the caller's key/team perms gate MCP, @@ -536,12 +539,17 @@ def _scope_auth_to_agent(user_api_key_auth: Any, agent_id: str) -> Any: return user_api_key_auth -async def _source_agent_allowed_targets(source_auth: Any) -> set | None: +async def _source_agent_allowed_targets( + source_auth: Optional[UserAPIKeyAuth], +) -> set | None: """Agent ids the source agent itself may call, from its own object_permission.agents plus agents resolved from its access groups. None means the source agent has no agent restriction. Delegation targets are intersected with this (not only the caller's perms), so a prompt-injected delegation cannot exceed what the source agent was granted.""" + if source_auth is None: + return None + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) @@ -564,14 +572,19 @@ async def _source_agent_allowed_targets(source_auth: Any) -> set | None: return allowed or None +@dataclass(frozen=True, slots=True) +class _DelegationContext: + allowed_agent_ids: set + get_agent: Callable[[str], Any] + user_api_key_auth: Optional[UserAPIKeyAuth] + send_message: Callable[..., Awaitable[Any]] + + async def _delegate_to_agent( *, target_id: str, message: str, - allowed_agent_ids: set, - get_agent: Callable[[str], Any], - user_api_key_auth: Any, - send_message: Callable[..., Awaitable[Any]], + context: _DelegationContext, ) -> str: """Dispatch a sub-agent call and return its text reply. @@ -579,10 +592,10 @@ async def _delegate_to_agent( via prompt injection, and ``send_message`` bypasses the endpoint's auth gate. The auth-filtered ``allowed_agent_ids`` whitelist is therefore re-enforced here so a caller can never reach an agent it was not granted.""" - if target_id not in allowed_agent_ids: + if target_id not in context.allowed_agent_ids: return f"Agent '{target_id}' is not available." - target = get_agent(target_id) + target = context.get_agent(target_id) if target is None: return f"Agent '{target_id}' not found." @@ -595,7 +608,7 @@ async def _delegate_to_agent( from uuid import uuid4 - from a2a.types import MessageSendParams, SendMessageRequest + from a2a.types import Message, MessageSendParams, SendMessageRequest target_card = target.agent_card_params or {} # Strip delegation flags so a sub-agent cannot recurse into further @@ -606,19 +619,22 @@ async def _delegate_to_agent( if k not in ("enable_agent_calls", "callable_agents", "call_agent") } target_params["user_api_key_auth"] = _scope_auth_to_agent( - user_api_key_auth, target.agent_id + context.user_api_key_auth, target.agent_id ) sub_request = SendMessageRequest( - id="", + id=uuid4().hex, params=MessageSendParams( - message={ - "role": "user", - "parts": [{"kind": "text", "text": message}], - "messageId": uuid4().hex, - } + message=cast( + Message, + { + "role": "user", + "parts": [{"kind": "text", "text": message}], + "messageId": uuid4().hex, + }, + ) ), ) - sub_response = await send_message( + sub_response = await context.send_message( request=sub_request, api_base=target_card.get("url"), litellm_params=target_params, @@ -807,10 +823,12 @@ async def invoke_agent_a2a( return await _delegate_to_agent( target_id=target_id, message=message, - allowed_agent_ids=allowed_agent_ids, - get_agent=_get_agent, - user_api_key_auth=user_api_key_dict, - send_message=asend_message, + context=_DelegationContext( + allowed_agent_ids=allowed_agent_ids, + get_agent=_get_agent, + user_api_key_auth=user_api_key_dict, + send_message=asend_message, + ), ) litellm_params = { diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py index f61064ac635..cc6f96ec1a8 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py @@ -1611,7 +1611,10 @@ async def test_delegate_to_agent_rejects_target_outside_whitelist(): and send_message bypasses the endpoint auth gate, so an agent id outside the auth-filtered whitelist must be refused before any dispatch or registry lookup.""" - from litellm.proxy.agent_endpoints.a2a_endpoints import _delegate_to_agent + from litellm.proxy.agent_endpoints.a2a_endpoints import ( + _DelegationContext, + _delegate_to_agent, + ) send_message = AsyncMock() get_agent = MagicMock() @@ -1619,10 +1622,12 @@ async def test_delegate_to_agent_rejects_target_outside_whitelist(): reply = await _delegate_to_agent( target_id="evil-agent", message="hi", - allowed_agent_ids={"allowed-agent"}, - get_agent=get_agent, - user_api_key_auth=None, - send_message=send_message, + context=_DelegationContext( + allowed_agent_ids={"allowed-agent"}, + get_agent=get_agent, + user_api_key_auth=None, + send_message=send_message, + ), ) assert "not available" in reply @@ -1632,17 +1637,22 @@ async def test_delegate_to_agent_rejects_target_outside_whitelist(): @pytest.mark.asyncio async def test_delegate_to_agent_does_not_dispatch_when_target_missing(): - from litellm.proxy.agent_endpoints.a2a_endpoints import _delegate_to_agent + from litellm.proxy.agent_endpoints.a2a_endpoints import ( + _DelegationContext, + _delegate_to_agent, + ) send_message = AsyncMock() reply = await _delegate_to_agent( target_id="known", message="hi", - allowed_agent_ids={"known"}, - get_agent=lambda _id: None, - user_api_key_auth=None, - send_message=send_message, + context=_DelegationContext( + allowed_agent_ids={"known"}, + get_agent=lambda _id: None, + user_api_key_auth=None, + send_message=send_message, + ), ) assert "not found" in reply @@ -1670,7 +1680,10 @@ async def test_delegate_to_agent_refuses_target_with_guardrails(): it is refused rather than returning output a guardrail would have blocked.""" from types import SimpleNamespace - from litellm.proxy.agent_endpoints.a2a_endpoints import _delegate_to_agent + from litellm.proxy.agent_endpoints.a2a_endpoints import ( + _DelegationContext, + _delegate_to_agent, + ) send_message = AsyncMock() guarded = SimpleNamespace( @@ -1682,10 +1695,12 @@ async def test_delegate_to_agent_refuses_target_with_guardrails(): reply = await _delegate_to_agent( target_id="guarded", message="hi", - allowed_agent_ids={"guarded"}, - get_agent=lambda _id: guarded, - user_api_key_auth=None, - send_message=send_message, + context=_DelegationContext( + allowed_agent_ids={"guarded"}, + get_agent=lambda _id: guarded, + user_api_key_auth=None, + send_message=send_message, + ), ) assert "cannot be called indirectly" in reply