diff --git a/docs/my-website/docs/mcp_chat_completions_orchestration.md b/docs/my-website/docs/mcp_chat_completions_orchestration.md new file mode 100644 index 00000000000..0c9b50be542 --- /dev/null +++ b/docs/my-website/docs/mcp_chat_completions_orchestration.md @@ -0,0 +1,307 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# MCP + Agent Orchestration in /chat/completions + +Use your registered MCP servers and A2A agents directly from `/v1/chat/completions` — no hardcoded URLs required. + +## The problem it solves + +Without this, every client has to know which MCP server to call: + +```json +// ❌ Every request hardcodes a specific server URL +{ + "tools": [{"type": "mcp", "server_url": "http://my-zapier-server/mcp", ...}] +} +``` + +That means updating every client when servers change, no central access control, and no way to let the LLM pick across multiple servers. + +## How it works + +Register your servers once via `POST /v1/mcp/server`. Then point any chat request at the proxy's registry using `"server_url": "litellm_proxy/mcp"` — the proxy fetches available tools, injects them into the LLM context, and executes tool calls on the model's behalf. + +``` +POST /v1/chat/completions + │ + ├── type:"mcp", server_url:"litellm_proxy/mcp" + │ └── expand → all servers registered via POST /v1/mcp/server + │ └── fetch tool schemas from each server + │ └── inject into LLM context + │ + └── type:"a2a_agent", server_url:"litellm_proxy/agents" + └── expand → all agents registered via POST /v1/agents + └── wrap each agent as a callable function tool + │ + ▼ + LLM decides which tools to call + │ + ▼ + Proxy executes tool calls, returns results + │ + ▼ + Follow-up LLM call → final answer +``` + +## Quickstart + +### 1. Register an MCP server + +```bash +curl -X POST http://localhost:4000/v1/mcp/server \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "server_name": "math", + "url": "http://localhost:8001/mcp", + "transport": "http" + }' +``` + +### 2. Call `/v1/chat/completions` with `litellm_proxy/mcp` + + + + +```bash +curl http://localhost:4000/v1/chat/completions \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "What is 1250 + 18?"}], + "tools": [ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp", + "require_approval": "never" + } + ] + }' +``` + + + + +```python +import openai + +client = openai.OpenAI( + api_key="sk-1234", + base_url="http://localhost:4000", +) + +response = client.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "What is 1250 + 18?"}], + tools=[ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp", + "require_approval": "never", + } + ], +) +print(response.choices[0].message.content) +``` + + + + +```python +import litellm + +response = await litellm.acompletion( + model="gpt-4o-mini", + api_base="http://localhost:4000", + api_key="sk-1234", + messages=[{"role": "user", "content": "What is 1250 + 18?"}], + tools=[ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp", + "require_approval": "never", + } + ], +) +print(response.choices[0].message.content) +``` + + + + +## Target a specific server + +Append the server name to `litellm_proxy/mcp/` to restrict tool injection to one server: + +```json +{ + "type": "mcp", + "server_url": "litellm_proxy/mcp/math", + "require_approval": "never" +} +``` + +## A2A Agent orchestration + +Agents registered via `POST /v1/agents` are exposed as callable function tools using the same pattern. + +### 1. Register an agent + +```bash +curl -X POST http://localhost:4000/v1/agents \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "agent_name": "FX_Converter", + "agent_url": "http://my-fx-agent/a2a", + "description": "Converts currency amounts using live exchange rates." + }' +``` + +### 2. Use MCP tools and agents together + +```json +{ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "Add 500 + 250, then convert the result to EUR."}], + "tools": [ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp", + "require_approval": "never" + }, + { + "type": "a2a_agent", + "server_url": "litellm_proxy/agents", + "require_approval": "never" + } + ] +} +``` + +The proxy wraps each registered agent as a function tool. When the LLM calls it, the proxy sends a JSON-RPC `message/send` to the agent and returns the result as a tool message. + +## Real-world example — Finance MCP + Compliance Agent + +Register a finance calculation MCP server and a compliance analyst A2A agent once. Every request can then use both without knowing any server URLs. + +```python +import openai + +client = openai.OpenAI( + api_key="sk-1234", + base_url="http://localhost:4000", +) + +# MCP only — financial calculation +response = client.chat.completions.create( + model="gpt-4o-mini", + messages=[{ + "role": "user", + "content": "What is the monthly repayment on a £250,000 mortgage at 4.5% APR over 25 years?" + }], + tools=[{ + "type": "mcp", + "server_url": "litellm_proxy/mcp", + "require_approval": "never", + }], +) +# → calls calculate_loan_payment tool → £1,389.58/mo + +# Both MCP + Agent in a single call +response = client.chat.completions.create( + model="gpt-4o-mini", + messages=[{ + "role": "user", + "content": ( + "Calculate compound interest on £100,000 at 2.8% over 3 years, " + "then draft a compliance note summarising the outcome for the audit file." + ) + }], + tools=[ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp", + "require_approval": "never", + }, + { + "type": "a2a_agent", + "server_url": "litellm_proxy/agents", + "require_approval": "never", + }, + ], +) +# → calls calculate_compound_interest (£8,637.40 interest) AND compliance_analyst agent +# → final answer includes both the numbers and the audit-ready compliance note +``` + +### Demo results (10 scenarios, local proxy, gpt-4o-mini) + +MCP server registered: `finance` — `calculate_compound_interest`, `convert_currency`, `calculate_loan_payment`, `calculate_var` +Agent registered: `compliance_analyst` — Basel III, KYC, VaR, earnings, trade summaries + +| # | Scenario | MCP | Agent | Tool Called | Result | +|---|----------|:---:|:---:|-------------|--------| +| 1 | Mortgage repayment | ✓ | — | `calculate_loan_payment(£250k, 4.5%, 25yr)` | **£1,389.58/mo** | +| 2 | FX conversion GBP→USD | ✓ | — | `convert_currency(£1.25M, 1.2738)` | £1,592,250 USD | +| 3 | Compound interest | ✓ | — | `calculate_compound_interest(£50k, 3.5%, 5yr)` | **£9,384 interest** | +| 4 | Basel III notice | — | ✓ | `compliance_analyst` | CET1 ≥4.5%, Tier1 ≥6% — review capital position | +| 5 | KYC note | — | ✓ | `compliance_analyst` | Entity verified, no sanctions, onboarding approved | +| 6 | VaR calculation | ✓ | ✓ | `calculate_var(£5M, 0.8% vol, 99%)` | 1-day VaR **£93,040**, 10-day **£294,218** | +| 7 | Interest calc + audit note | ✓ | ✓ | `calculate_compound_interest` + `compliance_analyst` | **£8,637 interest** + audit-ready compliance note | +| 8 | Mortgage refinance | ✓ | ✓ | `calculate_loan_payment(£180k, 3.9%, 20yr)` | **£1,081.30/mo** | +| 9 | Large FX GBP→JPY | ✓ | — | `convert_currency(£2.5M, 191.45)` | **¥478,625,000** | +| 10 | Earnings summary | — | ✓ | `compliance_analyst` | NII +8% YoY, CET1=13.8%, guidance reaffirmed | + +Row 7 demonstrates the orchestrator routing a single request to **both** the MCP finance server and the compliance analyst agent — the LLM received the calculation result from MCP and the formatted audit note from the agent in one turn, with no URL configuration in the client. + +## Semantic filter + +Add `"semantic_filter": true` to only inject tools relevant to the user's query. Useful when you have many registered servers and want to keep the LLM context lean. + +```json +{ + "type": "mcp", + "server_url": "litellm_proxy/mcp", + "require_approval": "never", + "semantic_filter": true +} +``` + +Configure top-k and similarity threshold in your proxy config: + +```yaml +mcp_semantic_tool_filter: + top_k: 10 + similarity_threshold: 0.3 + embedding_model: "text-embedding-3-small" +``` + +See [MCP Semantic Filter](./mcp_semantic_filter) for setup details. + +## Streaming + +Works with `"stream": true` — tokens arrive as they're generated, tool execution happens between LLM turns. + +```python +stream = await litellm.acompletion( + model="gpt-4o-mini", + api_base="http://localhost:4000", + api_key="sk-1234", + messages=[{"role": "user", "content": "What is 42 × 13?"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp", "require_approval": "never"}], + stream=True, +) +async for chunk in stream: + delta = chunk.choices[0].delta + if delta.content: + print(delta.content, end="", flush=True) +``` + +## Access control + +Tool visibility follows your existing key and team permissions. A virtual key scoped to specific MCP servers will only see those servers when it calls `litellm_proxy/mcp` — no extra config needed. + +See [MCP Zero Trust](./mcp_zero_trust) for per-key and per-team tool restrictions. diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 6446e227d99..15acebf8ffb 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -308,6 +308,7 @@ const sidebars = { "a2a_cost_tracking", "a2a_agent_permissions", "a2a_iteration_budgets", + "mcp_chat_completions_orchestration", ], }, { diff --git a/litellm/proxy/agent_endpoints/registry_orchestrator.py b/litellm/proxy/agent_endpoints/registry_orchestrator.py new file mode 100644 index 00000000000..6d265cf307f --- /dev/null +++ b/litellm/proxy/agent_endpoints/registry_orchestrator.py @@ -0,0 +1,307 @@ +""" +RegistryOrchestrator — centralises all registry-backed orchestration logic. + +Responsibilities +---------------- +- Parsing ``a2a_agent`` tool configs out of a request's ``tools`` list. +- Resolving registered A2A agents from the global agent registry and wrapping each + one as an OpenAI function tool that the LLM can call. +- Executing a single A2A tool call via JSON-RPC 2.0 ``message/send``. +- Applying the semantic MCP tool filter when the caller opts in via + ``"semantic_filter": true`` on an MCP tool config. +""" + +import re +from typing import Any, Dict, Iterable, List, Optional, Tuple + +from litellm._logging import verbose_logger + +# NOTE: Kept broad to avoid coupling to optional OpenAI SDK typing symbols. +ToolParam = Any + +LITELLM_PROXY_AGENTS_URL = "litellm_proxy/agents" + +# Import hoisted out of the callback loop to avoid re-evaluating on every iteration. +try: + from litellm.proxy.hooks.mcp_semantic_filter.hook import ( # noqa: E501 + SemanticToolFilterHook as _SemanticToolFilterHook, + ) +except ImportError: + _SemanticToolFilterHook = None # type: ignore + + +# --------------------------------------------------------------------------- +# Module-level helper +# --------------------------------------------------------------------------- + + +def _parse_a2a_response(data: Dict[str, Any]) -> str: + """Extract text content from an A2A JSON-RPC message/send response. + + Handles both ``"type": "text"`` (older A2A SDK) and ``"kind": "text"`` + (A2A SDK >= 0.3) part schemas. + """ + if "error" in data: + err = data["error"] + return f"Agent error: {err.get('message', str(err))}" + + result = data.get("result", {}) + + def _is_text_part(p: Dict[str, Any]) -> bool: + return (p.get("kind") == "text" or p.get("type") == "text") and bool( + p.get("text") + ) + + # A2A spec: result.artifacts[].parts[].text + for artifact in result.get("artifacts", []): + texts = [p["text"] for p in artifact.get("parts", []) if _is_text_part(p)] + if texts: + return "\n".join(texts) + + # Fallback: status.message.parts[].text + status = result.get("status", {}) + if isinstance(status, dict): + msg = status.get("message") or {} + for p in msg.get("parts", []): + if _is_text_part(p): + return p["text"] + + return str(result) if result else "Agent executed successfully" + + +# --------------------------------------------------------------------------- +# RegistryOrchestrator +# --------------------------------------------------------------------------- + + +class RegistryOrchestrator: + """ + Static-method class that owns all registry-backed orchestration concerns: + + * Parsing A2A agent tool configs from a request. + * Resolving registered agents and wrapping them as function tools. + * Executing A2A tool calls via JSON-RPC. + * Applying the per-request semantic MCP tool filter. + """ + + @staticmethod + def parse_agent_tool_configs( + tools: Optional[Iterable[ToolParam]], + ) -> Tuple[List[ToolParam], List[Any]]: + """ + Separate ``a2a_agent`` registry tool configs from all other tools. + + Returns: + (agent_tool_configs, other_tools) + """ + agent_tool_configs: List[ToolParam] = [] + other_tools: List[Any] = [] + + if tools: + for tool in tools: + if isinstance(tool, dict) and tool.get("type") == "a2a_agent": + server_url = tool.get("server_url", "") + if ( + isinstance(server_url, str) + and server_url == LITELLM_PROXY_AGENTS_URL + ): + agent_tool_configs.append(tool) + else: + other_tools.append(tool) + else: + other_tools.append(tool) + + return agent_tool_configs, other_tools + + @staticmethod + async def resolve_agent_tools( + user_api_key_auth: Any, + ) -> Tuple[List[Dict[str, Any]], Dict[str, Dict[str, str]]]: + """ + Read all registered A2A agents and expose each as an OpenAI function tool. + + Returns: + (function_tools, agent_tool_map) + + * ``function_tools``: list of ``{"type": "function", "function": {...}}`` dicts + * ``agent_tool_map``: mapping of sanitized function name → ``{"url": str, "agent_name": str}`` + """ + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + # NOTE: user_api_key_auth is accepted for future per-key agent filtering + # (mirroring get_allowed_mcp_servers). Agent-level access control is not yet + # implemented in AgentRegistry; all registered agents are returned for now. + _ = user_api_key_auth + + agents = global_agent_registry.get_agent_list() + function_tools: List[Dict[str, Any]] = [] + agent_tool_map: Dict[str, Dict[str, str]] = {} + + for agent in agents: + card = agent.agent_card_params or {} + agent_url = card.get("url", "") + agent_name = card.get("name") or agent.agent_name + + if not agent_url: + verbose_logger.warning( + "Agent '%s' has no URL configured, skipping", agent_name + ) + continue + + description = card.get("description") or f"A2A agent: {agent_name}" + + # Enrich description with up to 3 skill descriptions + skills = card.get("skills") or [] + skill_descs = [ + s.get("description", "") + for s in skills[:3] + if isinstance(s, dict) and s.get("description") + ] + if skill_descs: + description += " Skills: " + "; ".join(skill_descs) + + # Sanitize to a valid OpenAI function name (^[a-zA-Z0-9_-]{1,64}$) + func_name = ( + re.sub(r"[^a-zA-Z0-9_-]", "_", agent_name)[:64] + or f"agent_{agent.agent_id[:8]}" + ) + + # Deduplicate: if two agents produce the same sanitized name, append the + # agent_id suffix so neither is silently dropped. + if func_name in agent_tool_map: + func_name = f"{func_name}_{agent.agent_id[:8]}"[:64] + verbose_logger.warning( + "Agent name collision: renamed to '%s' to avoid overwrite", + func_name, + ) + + function_tools.append( + { + "type": "function", + "function": { + "name": func_name, + "description": description, + "parameters": { + "type": "object", + "properties": { + "message": { + "type": "string", + "description": "The message or task to send to this agent", + } + }, + "required": ["message"], + "additionalProperties": False, + }, + }, + } + ) + agent_tool_map[func_name] = {"url": agent_url, "agent_name": agent_name} + + verbose_logger.debug( + "Wrapped %d registered agents as function tools: %s", + len(function_tools), + list(agent_tool_map.keys()), + ) + return function_tools, agent_tool_map + + @staticmethod + async def execute_a2a_tool_call( + agent_url: str, + agent_name: str, + message: str, + tool_call_id: str, + tool_name: str, + litellm_trace_id: Optional[str] = None, + ) -> Dict[str, Any]: + """Send a message to an A2A agent via LiteLLM's asend_message and return the result.""" + import uuid + + try: + from a2a.types import ( + Message, + MessageSendParams, + Part, + Role, + SendMessageRequest, + TextPart, + ) + except ImportError as exc: + raise ImportError( + "The 'a2a' package is required for A2A agent calls. " + "Install it with: pip install a2a-sdk" + ) from exc + + from litellm.a2a_protocol.main import asend_message + + a2a_message = Message( + role=Role.user, + parts=[Part(root=TextPart(text=message))], + message_id=uuid.uuid4().hex, + context_id=litellm_trace_id, + ) + request = SendMessageRequest( + id=str(uuid.uuid4()), + params=MessageSendParams(message=a2a_message), + ) + + try: + response = await asend_message( + api_base=agent_url, + request=request, + agent_id=agent_name, + ) + result_text = _parse_a2a_response( + response.model_dump(mode="json", exclude_none=True) + ) + verbose_logger.debug( + "A2A agent '%s' returned: %s", agent_name, result_text[:200] + ) + return { + "tool_call_id": tool_call_id, + "result": result_text, + "name": tool_name, + } + except Exception as e: + verbose_logger.exception("Error calling A2A agent '%s': %s", agent_name, e) + return { + "tool_call_id": tool_call_id, + "result": f"Error calling agent {agent_name}: {str(e)}", + "name": tool_name, + } + + @staticmethod + async def apply_semantic_filter( + tools: List[Any], + messages: List[Any], + ) -> List[Any]: + """ + Filter MCP tools semantically based on the user query. + + Uses the global ``SemanticToolFilterHook`` if configured; otherwise returns + all tools unchanged. + """ + try: + import litellm + + for callback in litellm.callbacks or []: + if _SemanticToolFilterHook is None: + break + if isinstance(callback, _SemanticToolFilterHook): + query = callback.filter.extract_user_query(messages) + if query: + filtered = await callback.filter.filter_tools( + query=query, + available_tools=tools, + ) + verbose_logger.debug( + "Semantic filter (per-tool flag): %d → %d tools for query '%s...'", + len(tools), + len(filtered), + query[:60], + ) + return filtered + except Exception as e: + verbose_logger.warning( + "semantic_filter flag: filter failed (%s), using all tools", e + ) + return tools diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 24b5db28571..f8e1139d075 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -2,17 +2,19 @@ from typing import ( Any, + Dict, List, Optional, Union, cast, ) +from litellm._logging import verbose_logger from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.types.utils import ModelResponse +from litellm.types.utils import ModelResponse, ModelResponseStream, StreamingChoices from litellm.utils import CustomStreamWrapper @@ -77,6 +79,356 @@ def _add_mcp_metadata_to_response( setattr(message, "provider_specific_fields", provider_fields) +class _SyncIteratorWrapper: + """Wraps an async iterator for synchronous iteration.""" + + def __init__(self, async_iterator: Any, loop: Any) -> None: + self._async_iterator = async_iterator + self._loop = loop + self._iterator: Any = None + + def __iter__(self) -> "_SyncIteratorWrapper": + return self + + def __next__(self) -> Any: + import asyncio + + if self._iterator is None: + aiter_result = self._async_iterator.__aiter__() + if hasattr(aiter_result, "__await__"): + self._iterator = self._loop.run_until_complete(aiter_result) + else: + self._iterator = aiter_result + try: + return self._loop.run_until_complete(self._iterator.__anext__()) + except StopAsyncIteration: + raise StopIteration + + +class MCPStreamingIterator: + """ + Async iterator that drives the MCP tool-execution loop for streaming responses. + + Phases: + 1. Yield chunks from the initial LLM stream. + 2. When the stream ends, execute any tool calls (MCP or A2A). + 3. Yield chunks from the follow-up LLM stream. + """ + + def __init__( + self, + stream_wrapper: Any, + messages: List, + tool_server_map: Any, + user_api_key_auth: Any, + mcp_auth_header: Optional[str], + mcp_server_auth_headers: Any, + oauth2_headers: Any, + raw_headers: Any, + litellm_call_id: Optional[str], + litellm_trace_id: Optional[str], + openai_tools: List, + base_call_args: Dict[str, Any], + agent_tool_map: Optional[Dict[str, Any]] = None, + ) -> None: + self.stream_wrapper = stream_wrapper + self.messages = messages + self.tool_server_map = tool_server_map + self.user_api_key_auth = user_api_key_auth + self.mcp_auth_header = mcp_auth_header + self.mcp_server_auth_headers = mcp_server_auth_headers + self.oauth2_headers = oauth2_headers + self.raw_headers = raw_headers + self.litellm_call_id = litellm_call_id + self.litellm_trace_id = litellm_trace_id + self.openai_tools = openai_tools + self.base_call_args = base_call_args + self.agent_tool_map = agent_tool_map or {} + self.collected_chunks: List[ModelResponseStream] = [] + self.tool_calls: Optional[List] = None + self.tool_results: Optional[List] = None + self.complete_response: Optional[ModelResponse] = None + self.stream_exhausted = False + self.tool_execution_done = False + self.follow_up_stream: Optional[CustomStreamWrapper] = None + self.follow_up_non_stream: Optional[ModelResponse] = None + self.follow_up_iterator: Any = None + self.follow_up_exhausted = False + + def __aiter__(self) -> "MCPStreamingIterator": + return self + + def _add_mcp_list_tools_to_chunk( + self, chunk: ModelResponseStream + ) -> ModelResponseStream: + """Add mcp_list_tools to the first chunk.""" + from litellm.types.utils import add_provider_specific_fields + + if not self.openai_tools: + return chunk + + if hasattr(chunk, "choices") and chunk.choices: + for choice in chunk.choices: + if ( + isinstance(choice, StreamingChoices) + and hasattr(choice, "delta") + and choice.delta + ): + provider_fields = dict( + getattr(choice.delta, "provider_specific_fields", None) or {} + ) + provider_fields["mcp_list_tools"] = self.openai_tools + add_provider_specific_fields(choice.delta, provider_fields) + + return chunk + + def _add_mcp_tool_metadata_to_final_chunk( + self, chunk: ModelResponseStream + ) -> ModelResponseStream: + """Add mcp_tool_calls and mcp_call_results to the final chunk.""" + from litellm.types.utils import add_provider_specific_fields + + if hasattr(chunk, "choices") and chunk.choices: + for choice in chunk.choices: + if ( + isinstance(choice, StreamingChoices) + and hasattr(choice, "delta") + and choice.delta + ): + attr_value = getattr(choice.delta, "provider_specific_fields", None) + provider_fields = ( + dict(attr_value) if isinstance(attr_value, dict) else {} + ) + + if self.tool_calls: + provider_fields["mcp_tool_calls"] = self.tool_calls + if self.tool_results: + provider_fields["mcp_call_results"] = self.tool_results + + add_provider_specific_fields(choice.delta, provider_fields) + + return chunk + + async def __anext__(self) -> Any: + # Phase 1: Collect and yield initial stream chunks + if not self.stream_exhausted: + if not hasattr(self, "_stream_iterator"): + self._stream_iterator = self.stream_wrapper.__aiter__() + _add_mcp_metadata_to_response( + response=self.stream_wrapper, + openai_tools=self.openai_tools, + ) + + try: + chunk = await self._stream_iterator.__anext__() + self.collected_chunks.append(chunk) + + if len(self.collected_chunks) == 1: + chunk = self._add_mcp_list_tools_to_chunk(chunk) + + is_final = ( + hasattr(chunk, "choices") + and chunk.choices + and hasattr(chunk.choices[0], "finish_reason") + and chunk.choices[0].finish_reason is not None + ) + + if is_final: + self.stream_exhausted = True + await self._process_tool_calls() + chunk = self._add_mcp_tool_metadata_to_final_chunk(chunk) + if self.tool_results and self.complete_response: + await self._prepare_follow_up_call() + + return chunk + except StopAsyncIteration: + self.stream_exhausted = True + await self._process_tool_calls() + if self.collected_chunks: + final_chunk = self.collected_chunks[-1] + final_chunk = self._add_mcp_tool_metadata_to_final_chunk( + final_chunk + ) + if self.tool_results and self.complete_response: + await self._prepare_follow_up_call() + return final_chunk + + # Phase 2: Yield follow-up stream chunks if available + if self.follow_up_stream and not self.follow_up_exhausted: + if not self.follow_up_iterator: + self.follow_up_iterator = self.follow_up_stream.__aiter__() + verbose_logger.debug("Follow-up stream iterator created") + + try: + chunk = await self.follow_up_iterator.__anext__() + verbose_logger.debug("Follow-up chunk yielded: %s", chunk) + return chunk + except StopAsyncIteration: + self.follow_up_exhausted = True + verbose_logger.debug("Follow-up stream exhausted") + raise StopAsyncIteration + + # Phase 3: emit non-streaming follow-up answer as a synthetic final chunk + if self.follow_up_non_stream is not None: + from litellm.types.utils import ModelResponseStream, StreamingChoices + + non_stream = self.follow_up_non_stream + self.follow_up_non_stream = None + # Build a minimal streaming chunk from the ModelResponse + content = "" + if non_stream.choices: + content = getattr(non_stream.choices[0].message, "content", "") or "" + synthetic = ModelResponseStream( + id=non_stream.id, + model=non_stream.model or "", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta={"role": "assistant", "content": content}, # type: ignore[arg-type] + ) + ], + ) + return synthetic + + raise StopAsyncIteration + + async def _process_tool_calls(self) -> None: + """Build complete response from collected chunks and execute any tool calls.""" + from litellm.main import stream_chunk_builder + + if self.tool_execution_done: + return + + self.tool_execution_done = True + + if not self.collected_chunks: + return + + complete_response = stream_chunk_builder( + chunks=self.collected_chunks, + messages=self.messages, + ) + + if isinstance(complete_response, ModelResponse): + self.complete_response = complete_response + self.tool_calls = ( + LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response( + response=complete_response + ) + ) + + if self.tool_calls: + self.tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map=self.tool_server_map, + tool_calls=self.tool_calls, + user_api_key_auth=self.user_api_key_auth, + mcp_auth_header=self.mcp_auth_header, + mcp_server_auth_headers=self.mcp_server_auth_headers, + oauth2_headers=self.oauth2_headers, + raw_headers=self.raw_headers, + litellm_call_id=self.litellm_call_id, + litellm_trace_id=self.litellm_trace_id, + agent_tool_map=self.agent_tool_map, + ) + + async def _prepare_follow_up_call(self) -> None: + """Initiate the follow-up streaming call with tool results.""" + if self.follow_up_stream is not None: + return + + if not self.tool_results or not self.complete_response: + return + + follow_up_messages = ( + LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( + original_messages=self.messages, + response=self.complete_response, + tool_results=self.tool_results, + ) + ) + + follow_up_call_args = { + **self.base_call_args, + "messages": follow_up_messages, + "stream": True, + "_skip_mcp_handler": True, + } + + import litellm + + follow_up_response = await litellm.acompletion(**follow_up_call_args) + + if isinstance(follow_up_response, CustomStreamWrapper): + self.follow_up_stream = follow_up_response + verbose_logger.debug("Follow-up stream created successfully") + elif isinstance(follow_up_response, ModelResponse): + # Provider returned a non-streaming response despite stream=True. + # Store it so __anext__ can yield it as a synthetic final chunk rather + # than silently dropping the follow-up answer. + self.follow_up_non_stream = follow_up_response + verbose_logger.debug( + "Follow-up response is non-streaming ModelResponse; will emit as final chunk" + ) + else: + verbose_logger.warning( + "Follow-up response is unexpected type %s, answer may be dropped", + type(follow_up_response), + ) + + +class MCPStreamWrapper(CustomStreamWrapper): + """ + Thin ``CustomStreamWrapper`` subclass that delegates async iteration to + an ``MCPStreamingIterator`` so that the MCP tool-execution loop is + transparent to callers that consume the stream normally. + """ + + def __init__( + self, + original_wrapper: CustomStreamWrapper, + custom_iterator: MCPStreamingIterator, + ) -> None: + super().__init__( + completion_stream=None, + model=getattr(original_wrapper, "model", "unknown"), + logging_obj=getattr(original_wrapper, "logging_obj", None), + custom_llm_provider=getattr(original_wrapper, "custom_llm_provider", None), + stream_options=getattr(original_wrapper, "stream_options", None), + make_call=getattr(original_wrapper, "make_call", None), + _response_headers=getattr(original_wrapper, "_response_headers", None), + ) + self._original_wrapper = original_wrapper + self._custom_iterator = custom_iterator + if hasattr(original_wrapper, "_hidden_params"): + self._hidden_params = original_wrapper._hidden_params + self._sync_iterator: Optional[_SyncIteratorWrapper] = None + self._sync_loop: Any = None + + def __aiter__(self) -> MCPStreamingIterator: + return self._custom_iterator + + def __iter__(self) -> _SyncIteratorWrapper: + import asyncio + + if self._sync_iterator is None: + self._sync_loop = asyncio.new_event_loop() + asyncio.set_event_loop(self._sync_loop) + self._sync_iterator = _SyncIteratorWrapper( + self._custom_iterator, self._sync_loop + ) + return self._sync_iterator + + def __next__(self) -> Any: + if self._sync_iterator is None: + self.__iter__() + assert self._sync_iterator is not None + return next(self._sync_iterator) + + def __getattr__(self, name: str) -> Any: + return getattr(self._original_wrapper, name) + + async def acompletion_with_mcp( # noqa: PLR0915 model: str, messages: List, @@ -97,6 +449,7 @@ async def acompletion_with_mcp( # noqa: PLR0915 5. Make a follow-up call with the tool results """ from litellm import acompletion as litellm_acompletion + from litellm.proxy.agent_endpoints.registry_orchestrator import RegistryOrchestrator # Parse MCP tools and separate from other tools ( @@ -104,8 +457,14 @@ async def acompletion_with_mcp( # noqa: PLR0915 other_tools, ) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) - if not mcp_tools_with_litellm_proxy: - # No MCP tools, proceed with regular completion + # Parse A2A agent tools from what remains + ( + agent_tool_configs, + other_tools, + ) = RegistryOrchestrator.parse_agent_tool_configs(other_tools) + + if not mcp_tools_with_litellm_proxy and not agent_tool_configs: + # No MCP or agent tools, proceed with regular completion return await litellm_acompletion( model=model, messages=messages, @@ -141,17 +500,41 @@ async def acompletion_with_mcp( # noqa: PLR0915 mcp_server_auth_headers=mcp_server_auth_headers, ) + # Apply per-tool semantic filter if any MCP or agent tool config has semantic_filter=true + if any( + isinstance(t, dict) and t.get("semantic_filter") + for t in list(mcp_tools_with_litellm_proxy) + list(agent_tool_configs) + ): + deduplicated_mcp_tools = await RegistryOrchestrator.apply_semantic_filter( + tools=deduplicated_mcp_tools, + messages=messages, + ) + openai_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai( deduplicated_mcp_tools, target_format="chat", ) - # Combine with other tools - all_tools = openai_tools + other_tools if (openai_tools or other_tools) else None + # Wrap registered A2A agents as function tools + agent_function_tools: List = [] + agent_tool_map: dict = {} + if agent_tool_configs: + agent_function_tools, agent_tool_map = ( + await RegistryOrchestrator.resolve_agent_tools( + user_api_key_auth=user_api_key_auth, + ) + ) - # Determine if we should auto-execute tools + # Combine all tool types + combined = openai_tools + agent_function_tools + other_tools + all_tools: Optional[List] = combined if combined else None + + # Determine if we should auto-execute tools (MCP or agent tools with require_approval="never") should_auto_execute = LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools( mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy + ) or any( + isinstance(t, dict) and t.get("require_approval") == "never" + for t in agent_tool_configs ) # Prepare call parameters @@ -186,6 +569,7 @@ async def acompletion_with_mcp( # noqa: PLR0915 initial_call_args["stream"] = True if mock_tool_calls is not None: initial_call_args["mock_tool_calls"] = mock_tool_calls + _agent_tool_map = agent_tool_map # capture for closure # Make initial streaming call initial_stream = await litellm_acompletion(**initial_call_args) @@ -199,312 +583,7 @@ async def acompletion_with_mcp( # noqa: PLR0915 ) return initial_stream - # Create a custom async generator that collects chunks and handles tool execution - from litellm.main import stream_chunk_builder - from litellm.types.utils import ModelResponseStream - - class MCPStreamingIterator: - """Custom iterator that collects chunks, detects tool calls, and adds MCP metadata to final chunk.""" - - def __init__( - self, - stream_wrapper, - messages, - tool_server_map, - user_api_key_auth, - mcp_auth_header, - mcp_server_auth_headers, - oauth2_headers, - raw_headers, - litellm_call_id, - litellm_trace_id, - openai_tools, - base_call_args, - ): - self.stream_wrapper = stream_wrapper - self.messages = messages - self.tool_server_map = tool_server_map - self.user_api_key_auth = user_api_key_auth - self.mcp_auth_header = mcp_auth_header - self.mcp_server_auth_headers = mcp_server_auth_headers - self.oauth2_headers = oauth2_headers - self.raw_headers = raw_headers - self.litellm_call_id = litellm_call_id - self.litellm_trace_id = litellm_trace_id - self.openai_tools = openai_tools - self.base_call_args = base_call_args - self.collected_chunks: List[ModelResponseStream] = [] - self.tool_calls: Optional[List] = None - self.tool_results: Optional[List] = None - self.complete_response: Optional[ModelResponse] = None - self.stream_exhausted = False - self.tool_execution_done = False - self.follow_up_stream = None - self.follow_up_iterator = None - self.follow_up_exhausted = False - - async def __aiter__(self): - return self - - def _add_mcp_list_tools_to_chunk( - self, chunk: ModelResponseStream - ) -> ModelResponseStream: - """Add mcp_list_tools to the first chunk.""" - from litellm.types.utils import ( - StreamingChoices, - add_provider_specific_fields, - ) - - if not self.openai_tools: - return chunk - - if hasattr(chunk, "choices") and chunk.choices: - for choice in chunk.choices: - if ( - isinstance(choice, StreamingChoices) - and hasattr(choice, "delta") - and choice.delta - ): - # Get existing provider_specific_fields or create new dict - existing_fields = ( - getattr(choice.delta, "provider_specific_fields", None) - or {} - ) - provider_fields = dict( - existing_fields - ) # Create a copy to avoid mutating the original - - # Add only mcp_list_tools to first chunk - provider_fields["mcp_list_tools"] = self.openai_tools - - # Use add_provider_specific_fields to ensure proper setting - # This function handles Pydantic model attribute setting correctly - add_provider_specific_fields(choice.delta, provider_fields) - - return chunk - - def _add_mcp_tool_metadata_to_final_chunk( - self, chunk: ModelResponseStream - ) -> ModelResponseStream: - """Add mcp_tool_calls and mcp_call_results to the final chunk.""" - from litellm.types.utils import ( - StreamingChoices, - add_provider_specific_fields, - ) - - if hasattr(chunk, "choices") and chunk.choices: - for choice in chunk.choices: - if ( - isinstance(choice, StreamingChoices) - and hasattr(choice, "delta") - and choice.delta - ): - # Get existing provider_specific_fields or create new dict - # Access the attribute directly to handle Pydantic model attributes correctly - existing_fields = {} - if hasattr(choice.delta, "provider_specific_fields"): - attr_value = getattr( - choice.delta, "provider_specific_fields", None - ) - if attr_value is not None: - # Create a copy to avoid mutating the original - existing_fields = ( - dict(attr_value) - if isinstance(attr_value, dict) - else {} - ) - - provider_fields = existing_fields - - # Add tool_calls and tool_results if available - if self.tool_calls: - provider_fields["mcp_tool_calls"] = self.tool_calls - if self.tool_results: - provider_fields["mcp_call_results"] = self.tool_results - - # Use add_provider_specific_fields to ensure proper setting - # This function handles Pydantic model attribute setting correctly - add_provider_specific_fields(choice.delta, provider_fields) - - return chunk - - async def __anext__(self): - # Phase 1: Collect and yield initial stream chunks - if not self.stream_exhausted: - # Get the iterator from the stream wrapper - if not hasattr(self, "_stream_iterator"): - self._stream_iterator = self.stream_wrapper.__aiter__() - # Add mcp_list_tools to the first chunk (available from the start) - _add_mcp_metadata_to_response( - response=self.stream_wrapper, - openai_tools=self.openai_tools, - ) - - try: - chunk = await self._stream_iterator.__anext__() - self.collected_chunks.append(chunk) - - # Add mcp_list_tools to the first chunk - if len(self.collected_chunks) == 1: - chunk = self._add_mcp_list_tools_to_chunk(chunk) - - # Check if this is the final chunk (has finish_reason) - is_final = ( - hasattr(chunk, "choices") - and chunk.choices - and hasattr(chunk.choices[0], "finish_reason") - and chunk.choices[0].finish_reason is not None - ) - - if is_final: - # This is the final chunk, mark stream as exhausted - self.stream_exhausted = True - # Process tool calls after we've collected all chunks - await self._process_tool_calls() - # Apply MCP metadata (tool_calls and tool_results) to final chunk - chunk = self._add_mcp_tool_metadata_to_final_chunk(chunk) - # If we have tool results, prepare follow-up call immediately - if self.tool_results and self.complete_response: - await self._prepare_follow_up_call() - - return chunk - except StopAsyncIteration: - self.stream_exhausted = True - # Process tool calls after stream is exhausted - await self._process_tool_calls() - # If we have chunks, yield the final one with metadata - if self.collected_chunks: - final_chunk = self.collected_chunks[-1] - final_chunk = self._add_mcp_tool_metadata_to_final_chunk( - final_chunk - ) - # If we have tool results, prepare follow-up call - if self.tool_results and self.complete_response: - await self._prepare_follow_up_call() - return final_chunk - - # Phase 2: Yield follow-up stream chunks if available - if self.follow_up_stream and not self.follow_up_exhausted: - if not self.follow_up_iterator: - self.follow_up_iterator = self.follow_up_stream.__aiter__() - from litellm._logging import verbose_logger - - verbose_logger.debug("Follow-up stream iterator created") - - try: - chunk = await self.follow_up_iterator.__anext__() - from litellm._logging import verbose_logger - - verbose_logger.debug(f"Follow-up chunk yielded: {chunk}") - return chunk - except StopAsyncIteration: - self.follow_up_exhausted = True - from litellm._logging import verbose_logger - - verbose_logger.debug("Follow-up stream exhausted") - # After follow-up stream is exhausted, check if we need to raise StopAsyncIteration - raise StopAsyncIteration - - # If we're here and follow_up_stream is None but we expected it, log a warning - if ( - self.stream_exhausted - and self.tool_results - and self.complete_response - and self.follow_up_stream is None - ): - from litellm._logging import verbose_logger - - verbose_logger.warning( - "Follow-up stream was not created despite having tool results" - ) - - raise StopAsyncIteration - - async def _process_tool_calls(self): - """Process tool calls after streaming completes.""" - if self.tool_execution_done: - return - - self.tool_execution_done = True - - if not self.collected_chunks: - return - - # Build complete response from chunks - complete_response = stream_chunk_builder( - chunks=self.collected_chunks, - messages=self.messages, - ) - - if isinstance(complete_response, ModelResponse): - self.complete_response = complete_response - # Extract tool calls from complete response - self.tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response( - response=complete_response - ) - - if self.tool_calls: - # Execute tool calls - self.tool_results = ( - await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( - tool_server_map=self.tool_server_map, - tool_calls=self.tool_calls, - user_api_key_auth=self.user_api_key_auth, - mcp_auth_header=self.mcp_auth_header, - mcp_server_auth_headers=self.mcp_server_auth_headers, - oauth2_headers=self.oauth2_headers, - raw_headers=self.raw_headers, - litellm_call_id=self.litellm_call_id, - litellm_trace_id=self.litellm_trace_id, - ) - ) - - async def _prepare_follow_up_call(self): - """Prepare and initiate follow-up call with tool results.""" - if self.follow_up_stream is not None: - return # Already prepared - - if not self.tool_results or not self.complete_response: - return - - # Create follow-up messages with tool results - follow_up_messages = ( - LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( - original_messages=self.messages, - response=self.complete_response, - tool_results=self.tool_results, - ) - ) - - # Make follow-up call with streaming - follow_up_call_args = dict(self.base_call_args) - follow_up_call_args["messages"] = follow_up_messages - follow_up_call_args["stream"] = True - # Ensure follow-up call doesn't trigger MCP handler again - follow_up_call_args["_skip_mcp_handler"] = True - - # Import litellm here to ensure we get the patched version - # This ensures the patch works correctly in tests - import litellm - - follow_up_response = await litellm.acompletion(**follow_up_call_args) - - # Ensure follow-up response is a CustomStreamWrapper - if isinstance(follow_up_response, CustomStreamWrapper): - self.follow_up_stream = follow_up_response - from litellm._logging import verbose_logger - - verbose_logger.debug("Follow-up stream created successfully") - else: - # Unexpected response type - log and set to None - from litellm._logging import verbose_logger - - verbose_logger.warning( - f"Follow-up response is not a CustomStreamWrapper: {type(follow_up_response)}" - ) - self.follow_up_stream = None - - # Create the custom iterator + # Create the MCP streaming iterator (module-level class) iterator = MCPStreamingIterator( stream_wrapper=initial_stream, messages=messages, @@ -518,88 +597,9 @@ async def acompletion_with_mcp( # noqa: PLR0915 litellm_trace_id=kwargs.get("litellm_trace_id"), openai_tools=openai_tools, base_call_args=base_call_args, + agent_tool_map=_agent_tool_map, ) - # Create a wrapper class that delegates to our custom iterator - # We'll use a simple approach: just replace the __aiter__ method - class MCPStreamWrapper(CustomStreamWrapper): - def __init__(self, original_wrapper, custom_iterator): - # Initialize with the same parameters as original wrapper - super().__init__( - completion_stream=None, - model=getattr(original_wrapper, "model", "unknown"), - logging_obj=getattr(original_wrapper, "logging_obj", None), - custom_llm_provider=getattr( - original_wrapper, "custom_llm_provider", None - ), - stream_options=getattr(original_wrapper, "stream_options", None), - make_call=getattr(original_wrapper, "make_call", None), - _response_headers=getattr( - original_wrapper, "_response_headers", None - ), - ) - self._original_wrapper = original_wrapper - self._custom_iterator = custom_iterator - # Copy important attributes from original wrapper - if hasattr(original_wrapper, "_hidden_params"): - self._hidden_params = original_wrapper._hidden_params - # For synchronous iteration, we need to run the async iterator - self._sync_iterator = None - self._sync_loop = None - - def __aiter__(self): - return self._custom_iterator - - def __iter__(self): - # For synchronous iteration, create a sync wrapper - if self._sync_iterator is None: - import asyncio - - try: - self._sync_loop = asyncio.get_event_loop() - except RuntimeError: - self._sync_loop = asyncio.new_event_loop() - asyncio.set_event_loop(self._sync_loop) - self._sync_iterator = _SyncIteratorWrapper( - self._custom_iterator, self._sync_loop - ) - return self._sync_iterator - - def __next__(self): - # Delegate to sync iterator - if self._sync_iterator is None: - self.__iter__() - return next(self._sync_iterator) - - def __getattr__(self, name): - # Delegate all other attributes to original wrapper - return getattr(self._original_wrapper, name) - - # Helper class to wrap async iterator for sync iteration - class _SyncIteratorWrapper: - def __init__(self, async_iterator, loop): - self._async_iterator = async_iterator - self._loop = loop - self._iterator = None - - def __iter__(self): - return self - - def __next__(self): - if self._iterator is None: - # __aiter__ might be async, so we need to await it - aiter_result = self._async_iterator.__aiter__() - if hasattr(aiter_result, "__await__"): - # It's a coroutine, await it - self._iterator = self._loop.run_until_complete(aiter_result) - else: - # It's already an iterator - self._iterator = aiter_result - try: - return self._loop.run_until_complete(self._iterator.__anext__()) - except StopAsyncIteration: - raise StopIteration - return cast(CustomStreamWrapper, MCPStreamWrapper(initial_stream, iterator)) # Non-streaming mode: use existing logic @@ -626,7 +626,7 @@ async def acompletion_with_mcp( # noqa: PLR0915 ) return initial_response - # Execute tool calls + # Execute tool calls (MCP + A2A agents) tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( tool_server_map=tool_server_map, tool_calls=tool_calls, @@ -637,6 +637,7 @@ async def acompletion_with_mcp( # noqa: PLR0915 raw_headers=raw_headers, litellm_call_id=kwargs.get("litellm_call_id"), litellm_trace_id=kwargs.get("litellm_trace_id"), + agent_tool_map=agent_tool_map, ) if not tool_results: diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 7a3934ffdaa..ed1759d4526 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -148,7 +148,8 @@ class LiteLLM_Proxy_MCP_Handler: _get_tools_from_mcp_servers, ) - mcp_servers: List[str] = [] + # None means "fetch from all allowed servers"; a non-empty list means specific servers only. + mcp_servers: Optional[List[str]] = None if mcp_tools_with_litellm_proxy: for _tool in mcp_tools_with_litellm_proxy: # if user specifies servers as server_url: litellm_proxy/mcp/zapier,github then return zapier,github @@ -158,7 +159,14 @@ class LiteLLM_Proxy_MCP_Handler: if isinstance(server_url, str) and server_url.startswith( LITELLM_PROXY_MCP_SERVER_URL_PREFIX ): - mcp_servers.append(server_url.split("/")[-1]) + # "litellm_proxy/mcp/github" → specific server "github" + # "litellm_proxy/mcp" → no server name suffix → fetch all (leave mcp_servers=None) + server_name = server_url[len(LITELLM_PROXY_MCP_SERVER_URL_PREFIX) :] + if server_name: + if mcp_servers is None: + mcp_servers = [] + mcp_servers.append(server_name) + # else: bare "litellm_proxy/mcp" means all servers → keep None tools = await _get_tools_from_mcp_servers( user_api_key_auth=user_api_key_auth, @@ -537,20 +545,30 @@ class LiteLLM_Proxy_MCP_Handler: raw_headers: Optional[Dict[str, str]] = None, litellm_call_id: Optional[str] = None, litellm_trace_id: Optional[str] = None, + agent_tool_map: Optional[Dict[str, Dict[str, str]]] = None, ) -> List[Dict[str, Any]]: """Execute tool calls and return results.""" - from fastapi import HTTPException + try: + from fastapi import HTTPException + except ImportError: + # FastAPI is a proxy-only dependency; fall back to a plain exception so + # SDK users (without fastapi installed) can still call MCP tools. + # The .detail access below is already guarded with hasattr(). + HTTPException = Exception # type: ignore[assignment,misc] from litellm._uuid import uuid from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) + from litellm.proxy.agent_endpoints.registry_orchestrator import ( + RegistryOrchestrator, + ) from litellm.proxy.proxy_server import proxy_logging_obj + rules_obj = Rules() tool_results = [] tool_call_id: Optional[str] = None - rules_obj = Rules() for tool_call in tool_calls: logging_request_data: Dict[str, Any] = {} tool_name: Optional[str] = None @@ -569,10 +587,111 @@ class LiteLLM_Proxy_MCP_Handler: tool_arguments ) - # Import here to avoid circular import - from litellm.proxy.proxy_server import proxy_logging_obj + # Route A2A agent tool calls through the same logging path as MCP tools + if agent_tool_map and tool_name in agent_tool_map: + agent_info = agent_tool_map[tool_name] + message = parsed_arguments.get("message") or str(parsed_arguments) + start_time = datetime.now() + logging_request_data = { + "model": f"A2A: {tool_name}", + "metadata": { + "tool_call_id": tool_call_id, + "tool_name": tool_name, + "agent_name": agent_info["agent_name"], + }, + "input": [{"role": "tool", "content": message}], + "call_type": CallTypes.call_mcp_tool.value, + "litellm_call_id": litellm_call_id or str(uuid.uuid4()), + "proxy_server_request": { + "url": agent_info["url"], + "method": "POST", + "headers": {}, + "body": {"message": message}, + }, + } + if litellm_trace_id: + logging_request_data["litellm_trace_id"] = litellm_trace_id + if user_api_key_auth is not None: + user_api_key = getattr(user_api_key_auth, "api_key", None) + if user_api_key: + logging_request_data["metadata"][ + "user_api_key" + ] = user_api_key - server_name = tool_server_map[tool_name] + litellm_logging_obj = None + try: + litellm_logging_obj, _ = function_setup( + original_function="call_mcp_tool", + rules_obj=rules_obj, + start_time=start_time, + **logging_request_data, + ) + except Exception as _log_err: + verbose_logger.debug( + "Failed to init logging for A2A tool call %s: %s", + tool_name, + _log_err, + ) + + standard_logging_a2a_tool_call: StandardLoggingMCPToolCall = { + "name": tool_name, + "arguments": parsed_arguments, + "namespaced_tool_name": tool_name, + "mcp_server_name": agent_info["agent_name"], + } + if litellm_logging_obj: + litellm_logging_obj.model_call_details[ + "mcp_tool_call_metadata" + ] = standard_logging_a2a_tool_call + litellm_logging_obj.model = f"A2A: {tool_name}" + litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value + try: + litellm_logging_obj.pre_call(input=[message], api_key="") + except Exception: + pass + + result = await RegistryOrchestrator.execute_a2a_tool_call( + agent_url=agent_info["url"], + agent_name=agent_info["agent_name"], + message=message, + tool_call_id=tool_call_id or "", + tool_name=tool_name, + litellm_trace_id=litellm_trace_id, + ) + + if litellm_logging_obj: + try: + litellm_logging_obj.post_call( + original_response=result.get("result", "") + ) + end_time = datetime.now() + await litellm_logging_obj.async_post_mcp_tool_call_hook( + kwargs=litellm_logging_obj.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) + await litellm_logging_obj.async_success_handler( + result=result, + start_time=start_time, + end_time=end_time, + ) + except Exception: + verbose_logger.exception( + "Failed to log A2A tool call success for %s", tool_name + ) + + tool_results.append(result) + continue + + server_name = tool_server_map.get(tool_name) + if server_name is None: + verbose_logger.warning( + "Tool '%s' not found in tool_server_map — skipping (possible " + "hallucinated tool name)", + tool_name, + ) + continue # Remove the server name prefix if the tool name includes it. sanitized_tool_name = tool_name @@ -682,14 +801,14 @@ class LiteLLM_Proxy_MCP_Handler: standard_logging_mcp_tool_call["mcp_server_logo_url"] = logo_url cost_info = mcp_info.get("mcp_server_cost_info") if cost_info: - standard_logging_mcp_tool_call[ - "mcp_server_cost_info" - ] = cost_info + standard_logging_mcp_tool_call["mcp_server_cost_info"] = ( + cost_info + ) if litellm_logging_obj: - litellm_logging_obj.model_call_details[ - "mcp_tool_call_metadata" - ] = standard_logging_mcp_tool_call + litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = ( + standard_logging_mcp_tool_call + ) litellm_logging_obj.model = f"MCP: {tool_name}" litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value @@ -779,7 +898,7 @@ class LiteLLM_Proxy_MCP_Handler: error=e, ) verbose_logger.error(f"HTTPException in MCP tool call: {str(e)}") - error_message = f"Tool call failed: {str(e.detail) if hasattr(e, 'detail') else str(e)}" + error_message = f"Tool call failed: {str(e.detail) if hasattr(e, 'detail') else str(e)}" # type: ignore[union-attr] tool_results.append( { "tool_call_id": tool_call_id, diff --git a/tests/test_litellm/responses/mcp/test_registry_orchestration.py b/tests/test_litellm/responses/mcp/test_registry_orchestration.py new file mode 100644 index 00000000000..54a504c43b7 --- /dev/null +++ b/tests/test_litellm/responses/mcp/test_registry_orchestration.py @@ -0,0 +1,675 @@ +""" +Registry-backed orchestration tests for /v1/chat/completions. + +Validates the feature where: + - server_url: "litellm_proxy/mcp" → expands to ALL registered MCP servers + - server_url: "litellm_proxy/agents" → expands to ALL registered A2A agents + - Both MCP tool calls and A2A agent calls share the same trace + - semantic_filter: true pre-filters MCP tools by query relevance + - Streaming (stream=True) works identically to non-streaming +""" + +from types import SimpleNamespace +from typing import Any, Dict, List, Optional +from unittest.mock import AsyncMock, patch + +import pytest + +import litellm +from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.agents import AgentResponse +from litellm.types.utils import ModelResponse + + +def _mcp_tool_to_openai(tool): + """Convert a SimpleNamespace MCP tool to OpenAI function tool format without importing mcp.""" + return { + "type": "function", + "function": { + "name": tool.name, + "description": tool.description, + "parameters": tool.inputSchema, + }, + } + + +# --------------------------------------------------------------------------- +# Shared fixtures +# --------------------------------------------------------------------------- + +MATH_MCP_TOOL = SimpleNamespace( + name="add", + description="Add two numbers", + inputSchema={ + "type": "object", + "properties": { + "a": {"type": "integer", "description": "First operand"}, + "b": {"type": "integer", "description": "Second operand"}, + }, + "required": ["a", "b"], + }, +) + +CURRENCY_AGENT = AgentResponse( + agent_id="agent-fx-001", + agent_name="FX_Converter", + agent_card_params={ + "url": "http://mock-agent.internal/a2a", + "name": "FX_Converter", + "description": "Converts amounts between currencies using live rates.", + "skills": [ + { + "id": "fx-convert", + "name": "Currency Conversion", + "description": "Convert a numeric amount from one currency to another", + "tags": ["finance", "fx"], + } + ], + }, +) + + +def _make_fake_process(mcp_tools=None, tool_server_map=None): + """Return a fake _process_mcp_tools_without_openai_transform.""" + _tools = mcp_tools or [MATH_MCP_TOOL] + _map = tool_server_map or {MATH_MCP_TOOL.name: "math_server"} + + async def fake_process(user_api_key_auth, mcp_tools_with_litellm_proxy, **kwargs): + return _tools, _map + + return fake_process + + +def _no_mcp_headers(secret_fields, tools): + return (None, None, None, None) + + +# --------------------------------------------------------------------------- +# Test 1 – Non-streaming: MCP + A2A tool calls in the same trace +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_registry_orchestration_nonstreaming(monkeypatch): + """ + One LLM turn triggers both an MCP tool call (add) and an A2A agent call + (FX_Converter). Both are executed in a single trace and the final answer + is returned as a ModelResponse. + """ + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + # ── Setup registries ────────────────────────────────────────────────── + original_agents = list(global_agent_registry.agent_list) + global_agent_registry.agent_list = [CURRENCY_AGENT] + + executed: List[Dict[str, Any]] = [] + + async def fake_execute(**kwargs): + tool_calls: List[Any] = kwargs.get("tool_calls") or [] + agent_tool_map: Dict[str, Any] = kwargs.get("agent_tool_map") or {} + + results = [] + for tc in tool_calls: + fn = tc.get("function") or {} + name = fn.get("name") or tc.get("name") or "" + call_id = tc.get("id") or "tc-unknown" + + if name == "add": + executed.append({"type": "mcp", "tool": "add"}) + results.append({"tool_call_id": call_id, "result": "12", "name": "add"}) + elif name in agent_tool_map or name == "FX_Converter": + executed.append({"type": "a2a", "tool": name}) + results.append( + { + "tool_call_id": call_id, + "result": "12 USD = 9.48 GBP", + "name": name, + } + ) + return results + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_process_mcp_tools_without_openai_transform", + _make_fake_process(), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_transform_mcp_tools_to_openai", + staticmethod(lambda tools, **kw: [_mcp_tool_to_openai(t) for t in tools]), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_execute_tool_calls", + fake_execute, + ) + monkeypatch.setattr( + ResponsesAPIRequestUtils, + "extract_mcp_headers_from_request", + staticmethod(_no_mcp_headers), + ) + + try: + response = await litellm.acompletion( + model="gpt-4o-mini", + messages=[ + { + "role": "user", + "content": "Add 5 and 7, then convert the result to GBP.", + } + ], + tools=[ + # All registered MCP servers + { + "type": "mcp", + "server_url": "litellm_proxy/mcp", + "require_approval": "never", + }, + # All registered A2A agents + { + "type": "a2a_agent", + "server_url": "litellm_proxy/agents", + "require_approval": "never", + }, + ], + # First LLM response: call both tools + mock_tool_calls=[ + { + "id": "tc-mcp-1", + "type": "function", + "function": { + "name": "add", + "arguments": '{"a": 5, "b": 7}', + }, + }, + { + "id": "tc-a2a-1", + "type": "function", + "function": { + "name": "FX_Converter", + "arguments": '{"message": "Convert 12 USD to GBP"}', + }, + }, + ], + # Second LLM response after tool results are fed back + mock_response="5 + 7 = 12. The FX Converter confirms: 12 USD = 9.48 GBP.", + ) + finally: + global_agent_registry.agent_list = original_agents + + # ── Assertions ──────────────────────────────────────────────────────── + assert isinstance(response, ModelResponse), "Expected a ModelResponse" + assert "12 USD = 9.48 GBP" in response.choices[0].message.content + + mcp_calls = [e for e in executed if e["type"] == "mcp"] + a2a_calls = [e for e in executed if e["type"] == "a2a"] + assert mcp_calls, "MCP tool 'add' was never executed" + assert a2a_calls, "A2A agent 'FX_Converter' was never executed" + + mcp_metadata = ( + response.choices[0].message.provider_specific_fields or {} + if hasattr(response.choices[0].message, "provider_specific_fields") + else {} + ) + # Both MCP list and agent tools should appear in provider metadata + assert ( + "mcp_list_tools" in mcp_metadata + ), f"Expected mcp_list_tools in provider_specific_fields, got: {list(mcp_metadata.keys())}" + + +# --------------------------------------------------------------------------- +# Test 2 – litellm_proxy/mcp bare URL expands to ALL registered servers +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_bare_mcp_url_expands_to_all_servers(monkeypatch): + """ + server_url: 'litellm_proxy/mcp' (no /server_name suffix) must call + _process_mcp_tools_without_openai_transform with mcp_servers=None so that + ALL registered MCP servers are queried, not just one. + """ + captured: Dict[str, Any] = {} + + async def spy_process(user_api_key_auth, mcp_tools_with_litellm_proxy, **kwargs): + # Record the tool config to check server_url below + captured["tools"] = mcp_tools_with_litellm_proxy + return [MATH_MCP_TOOL], {MATH_MCP_TOOL.name: "math_server"} + + async def fake_execute(**kwargs): + return [{"tool_call_id": "tc-1", "result": "8", "name": "add"}] + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_process_mcp_tools_without_openai_transform", + spy_process, + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_transform_mcp_tools_to_openai", + staticmethod(lambda tools, **kw: [_mcp_tool_to_openai(t) for t in tools]), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_execute_tool_calls", + fake_execute, + ) + monkeypatch.setattr( + ResponsesAPIRequestUtils, + "extract_mcp_headers_from_request", + staticmethod(_no_mcp_headers), + ) + + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "add 3 and 5"}], + tools=[ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp", # bare – no server suffix + "require_approval": "never", + } + ], + mock_tool_calls=[ + { + "id": "tc-1", + "type": "function", + "function": {"name": "add", "arguments": '{"a": 3, "b": 5}'}, + } + ], + mock_response="3 + 5 = 8", + ) + + # The tool config passed into _process_mcp_tools must include the bare URL + assert captured.get("tools"), "spy_process was never called" + bare_url_tools = [ + t + for t in captured["tools"] + if isinstance(t, dict) and t.get("server_url") == "litellm_proxy/mcp" + ] + assert bare_url_tools, ( + "Expected a tool entry with server_url='litellm_proxy/mcp' " + f"(all-servers sentinel). Got: {captured['tools']}" + ) + + +# --------------------------------------------------------------------------- +# Test 3 – Agent wrapping: agents exposed as OpenAI function tools +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_agents_wrapped_as_function_tools(monkeypatch): + """ + When agent_tool_configs are present, _wrap_agents_as_function_tools reads + global_agent_registry and converts each agent to an OpenAI function tool + with a sanitized name and enriched description. + """ + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + original_agents = list(global_agent_registry.agent_list) + global_agent_registry.agent_list = [CURRENCY_AGENT] + + try: + from litellm.proxy.agent_endpoints.registry_orchestrator import ( + RegistryOrchestrator, + ) + + function_tools, agent_tool_map = await RegistryOrchestrator.resolve_agent_tools( + user_api_key_auth=None + ) + finally: + global_agent_registry.agent_list = original_agents + + assert len(function_tools) == 1 + ft = function_tools[0] + assert ft["type"] == "function" + fn = ft["function"] + + # Name must be a valid OpenAI function name (alphanumeric + _ + -) + import re + + assert re.match( + r"^[a-zA-Z0-9_-]{1,64}$", fn["name"] + ), f"Function name '{fn['name']}' is not a valid OpenAI function name" + + # Description should be enriched with skill description + assert ( + "Convert a numeric amount" in fn["description"] + ), f"Skill description missing from function description: {fn['description']}" + + # Parameters schema must include 'message' field + params = fn["parameters"] + assert params["type"] == "object" + assert "message" in params["properties"] + assert params["required"] == ["message"] + + # agent_tool_map maps the sanitized name to the agent URL + assert fn["name"] in agent_tool_map + assert agent_tool_map[fn["name"]]["url"] == "http://mock-agent.internal/a2a" + + +# --------------------------------------------------------------------------- +# Test 4 – A2A response parsing +# --------------------------------------------------------------------------- + + +def test_parse_a2a_response_artifacts(): + """Extracts text from A2A result.artifacts[].parts[].""" + from litellm.proxy.agent_endpoints.registry_orchestrator import _parse_a2a_response + + data = { + "jsonrpc": "2.0", + "id": "req-1", + "result": { + "artifacts": [ + { + "parts": [ + {"type": "text", "text": "12 USD = 9.48 GBP"}, + {"type": "text", "text": "Rate: 0.79"}, + ] + } + ] + }, + } + result = _parse_a2a_response(data) + assert "12 USD = 9.48 GBP" in result + assert "Rate: 0.79" in result + + +def test_parse_a2a_response_status_message(): + """Falls back to result.status.message.parts[] when no artifacts.""" + from litellm.proxy.agent_endpoints.registry_orchestrator import _parse_a2a_response + + data = { + "jsonrpc": "2.0", + "id": "req-2", + "result": { + "status": { + "state": "completed", + "message": { + "role": "agent", + "parts": [{"type": "text", "text": "Done: 9.48 GBP"}], + }, + } + }, + } + assert _parse_a2a_response(data) == "Done: 9.48 GBP" + + +def test_parse_a2a_response_error(): + """Error responses surface the error message.""" + from litellm.proxy.agent_endpoints.registry_orchestrator import _parse_a2a_response + + data = { + "jsonrpc": "2.0", + "id": "req-3", + "error": {"code": -32600, "message": "Invalid Request"}, + } + result = _parse_a2a_response(data) + assert "Invalid Request" in result + + +# --------------------------------------------------------------------------- +# Test 5 – Streaming mode: MCP + A2A in same trace, stream=True +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_registry_orchestration_streaming(monkeypatch): + """ + With stream=True the handler wraps streaming in MCPStreamingIterator. + Collecting all chunks must yield a final text response that includes + both the MCP tool result and the A2A agent result. + """ + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.utils import CustomStreamWrapper + + original_agents = list(global_agent_registry.agent_list) + global_agent_registry.agent_list = [CURRENCY_AGENT] + + executed: List[Dict[str, Any]] = [] + + async def fake_execute(**kwargs): + tool_calls: List[Any] = kwargs.get("tool_calls") or [] + agent_tool_map: Dict[str, Any] = kwargs.get("agent_tool_map") or {} + results = [] + for tc in tool_calls: + fn = tc.get("function") or {} + name = fn.get("name") or tc.get("name") or "" + call_id = tc.get("id") or "tc-s" + if name == "add": + executed.append({"type": "mcp", "tool": "add"}) + results.append({"tool_call_id": call_id, "result": "12", "name": "add"}) + elif name in agent_tool_map or name == "FX_Converter": + executed.append({"type": "a2a", "tool": name}) + results.append( + { + "tool_call_id": call_id, + "result": "12 USD = 9.48 GBP", + "name": name, + } + ) + return results + + # Mock tool calls to be "found" after stream collection. + # mock_tool_calls in streaming mode are not reliably embedded in chunk deltas, + # so we inject them directly via _extract_tool_calls_from_chat_response. + _stream_tool_calls = [ + { + "id": "tc-s-mcp", + "type": "function", + "function": {"name": "add", "arguments": '{"a": 5, "b": 7}'}, + }, + { + "id": "tc-s-a2a", + "type": "function", + "function": { + "name": "FX_Converter", + "arguments": '{"message": "Convert 12 USD to GBP"}', + }, + }, + ] + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_process_mcp_tools_without_openai_transform", + _make_fake_process(), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_transform_mcp_tools_to_openai", + staticmethod(lambda tools, **kw: [_mcp_tool_to_openai(t) for t in tools]), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_extract_tool_calls_from_chat_response", + staticmethod(lambda response: _stream_tool_calls), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_execute_tool_calls", + fake_execute, + ) + monkeypatch.setattr( + ResponsesAPIRequestUtils, + "extract_mcp_headers_from_request", + staticmethod(_no_mcp_headers), + ) + + try: + response = await litellm.acompletion( + model="gpt-4o-mini", + messages=[ + { + "role": "user", + "content": "Add 5 and 7, then convert the result to GBP.", + } + ], + tools=[ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp", + "require_approval": "never", + }, + { + "type": "a2a_agent", + "server_url": "litellm_proxy/agents", + "require_approval": "never", + }, + ], + stream=True, + mock_tool_calls=[ + { + "id": "tc-s-mcp", + "type": "function", + "function": { + "name": "add", + "arguments": '{"a": 5, "b": 7}', + }, + }, + { + "id": "tc-s-a2a", + "type": "function", + "function": { + "name": "FX_Converter", + "arguments": '{"message": "Convert 12 USD to GBP"}', + }, + }, + ], + mock_response="5 + 7 = 12. The FX Converter confirms: 12 USD = 9.48 GBP.", + ) + finally: + global_agent_registry.agent_list = original_agents + + # Collect all chunks from the stream + chunks = [] + final_text = "" + if isinstance(response, CustomStreamWrapper): + async for chunk in response: + chunks.append(chunk) + delta = chunk.choices[0].delta if chunk.choices else None + if delta and getattr(delta, "content", None): + final_text += delta.content + elif isinstance(response, ModelResponse): + # Non-streaming fallback (shouldn't happen but handle gracefully) + final_text = response.choices[0].message.content or "" + + assert chunks, "No streaming chunks received" + assert ( + "12 USD = 9.48 GBP" in final_text + ), f"Expected final text to contain FX result. Got: {final_text!r}" + + # Both MCP and A2A calls must have fired during the stream loop + assert any(e["type"] == "mcp" for e in executed), "MCP tool not executed in stream" + assert any(e["type"] == "a2a" for e in executed), "A2A agent not executed in stream" + + +# --------------------------------------------------------------------------- +# Test 6 – semantic_filter flag: filter hook reduces tool count +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_semantic_filter_reduces_tools(monkeypatch): + """ + When semantic_filter: true is set on the MCP tool config, _apply_semantic_filter + is invoked. This test verifies the hook integration: if the filter is applied, + the tool list passed downstream is reduced. + """ + from litellm.responses.mcp import chat_completions_handler + + # Two MCP tools available + add_tool = MATH_MCP_TOOL + multiply_tool = SimpleNamespace( + name="multiply", + description="Multiply two numbers", + inputSchema={ + "type": "object", + "properties": { + "a": {"type": "integer"}, + "b": {"type": "integer"}, + }, + "required": ["a", "b"], + }, + ) + + async def fake_process(user_api_key_auth, mcp_tools_with_litellm_proxy, **kwargs): + return [add_tool, multiply_tool], { + "add": "math_server", + "multiply": "math_server", + } + + async def fake_execute(**kwargs): + return [{"tool_call_id": "tc-1", "result": "8", "name": "add"}] + + # Semantic filter: keep only the first tool (simulates "add" being most relevant) + async def fake_semantic_filter(tools, messages): + return tools[:1] # keep only 'add', drop 'multiply' + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_process_mcp_tools_without_openai_transform", + fake_process, + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_transform_mcp_tools_to_openai", + staticmethod(lambda tools, **kw: [_mcp_tool_to_openai(t) for t in tools]), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_execute_tool_calls", + fake_execute, + ) + monkeypatch.setattr( + ResponsesAPIRequestUtils, + "extract_mcp_headers_from_request", + staticmethod(_no_mcp_headers), + ) + # Patch RegistryOrchestrator.apply_semantic_filter (moved from module-level) + from litellm.proxy.agent_endpoints.registry_orchestrator import RegistryOrchestrator + + monkeypatch.setattr( + RegistryOrchestrator, + "apply_semantic_filter", + staticmethod(fake_semantic_filter), + ) + + response = await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "add 3 and 5"}], + tools=[ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp", + "require_approval": "never", + "semantic_filter": True, # ← trigger filter + } + ], + mock_tool_calls=[ + { + "id": "tc-1", + "type": "function", + "function": {"name": "add", "arguments": '{"a": 3, "b": 5}'}, + } + ], + mock_response="3 + 5 = 8", + ) + + assert isinstance(response, ModelResponse) + assert "8" in response.choices[0].message.content + + # Verify semantic filter was applied: only 'add' tool should appear in metadata + mcp_metadata = ( + response.choices[0].message.provider_specific_fields or {} + if hasattr(response.choices[0].message, "provider_specific_fields") + else {} + ) + listed = mcp_metadata.get("mcp_list_tools", []) + tool_names = [t.get("function", {}).get("name") for t in listed] + assert "add" in tool_names, f"Expected 'add' in mcp_list_tools, got: {tool_names}" + assert ( + "multiply" not in tool_names + ), f"Expected 'multiply' to be filtered out, got: {tool_names}"