diff --git a/agent/protocol/agent_stream.py b/agent/protocol/agent_stream.py index 1eb96e50..920112b2 100644 --- a/agent/protocol/agent_stream.py +++ b/agent/protocol/agent_stream.py @@ -379,6 +379,12 @@ class AgentStreamExecutor: self._emit_event("agent_start") + # Reset the run-scoped MCP tool-retrieval accumulator. On-demand tool + # retrieval only grows this set within a run, so a tool that already + # produced a tool_use never disappears from the schema mid-run (which + # would make Claude/MiniMax raise a message-format error). + self._retrieved_mcp_names = set() + final_response = "" turn = 0 @@ -702,6 +708,70 @@ class AgentStreamExecutor: return final_response + def _select_tools_for_injection(self) -> list: + """Decide which tools to inject into the current LLM turn. + + Built-in tools are ALWAYS injected in full (skills and core flows hard + depend on them). MCP tools are also injected in full UNLESS on-demand + retrieval is enabled AND the MCP tool count exceeds the configured + threshold — then only the most relevant MCP tools are injected, unioned + with those already selected earlier in this run (only-grows, so a tool + that already produced a tool_use never vanishes from the schema). + + Degrades safely: disabled feature, no embedding provider, embedding + failure, count below threshold, or any error → inject all tools. Tools + are never silently dropped. + """ + all_tools = list(self.tools.values()) + try: + from config import conf + if not conf().get("mcp_tool_retrieval_enabled", False): + return all_tools + + from agent.tools.mcp.mcp_tool import McpTool + mcp_tools = [t for t in all_tools if isinstance(t, McpTool)] + builtin_tools = [t for t in all_tools if not isinstance(t, McpTool)] + + threshold = int(conf().get("mcp_tool_retrieval_threshold", 20) or 20) + if len(mcp_tools) <= threshold: + return all_tools + + top_k = int(conf().get("mcp_tool_retrieval_top_k", 10) or 10) + + from agent.tools import ToolManager + from agent.tools.mcp.tool_retrieval import ( + build_retrieval_query, + select_mcp_tools, + ) + + tm = ToolManager() + tool_vectors = tm.get_mcp_tool_vectors() + query = build_retrieval_query(self.messages) + query_vector = tm.embed_query(query) + + selected = select_mcp_tools( + query_vector, + tool_vectors, + top_k, + getattr(self, "_retrieved_mcp_names", set()), + ) + if selected is None: + # No provider / empty index / error → full injection. + return all_tools + + # Persist the accumulated selection for subsequent turns. + self._retrieved_mcp_names = selected + + selected_mcp = [t for t in mcp_tools if t.name in selected] + logger.info( + f"[ToolRetrieval] Injecting {len(builtin_tools)} built-in + " + f"{len(selected_mcp)}/{len(mcp_tools)} MCP tool(s) (top_k={top_k})" + ) + return builtin_tools + selected_mcp + except Exception as e: + logger.debug(f"[ToolRetrieval] full injection (retrieval skipped): {e}") + return all_tools + def _call_llm_stream(self, retry_on_empty=True, retry_count=0, max_retries=3, _overflow_retry: bool = False) -> Tuple[str, List[Dict]]: """ @@ -742,7 +812,7 @@ class AgentStreamExecutor: tools_schema = None if self.tools: tools_schema = [] - for tool in self.tools.values(): + for tool in self._select_tools_for_injection(): input_schema = tool.params try: dynamic = (tool.get_json_schema() or {}).get("parameters") or {} diff --git a/config-template.json b/config-template.json index 1fa50b24..ad3583bf 100644 --- a/config-template.json +++ b/config-template.json @@ -41,5 +41,8 @@ "enable_thinking": false, "reasoning_effort": "high", "knowledge": true, - "self_evolution_enabled": true + "self_evolution_enabled": true, + "mcp_tool_retrieval_enabled": false, + "mcp_tool_retrieval_threshold": 20, + "mcp_tool_retrieval_top_k": 10 } diff --git a/config.py b/config.py index 921feffc..6fe2db03 100644 --- a/config.py +++ b/config.py @@ -269,6 +269,13 @@ available_setting = { "deep_dream_enabled": True, # scheduled deep dream switch; manual /memory dream is unaffected "skill": {}, # Per-skill runtime config; nested keys flatten to SKILL__ env vars at startup "mcp_servers": [], # MCP server list; each entry supports type "stdio" (local process) or "sse" (remote URL) + # On-demand MCP tool retrieval: when many MCP tools are connected, inject + # only the most query-relevant ones instead of all of them. Built-in tools + # are always injected in full; degrades to full injection when disabled, + # below threshold, or when no embedding provider is available. + "mcp_tool_retrieval_enabled": False, # switch for on-demand MCP tool retrieval + "mcp_tool_retrieval_threshold": 20, # only retrieve when MCP tool count exceeds this + "mcp_tool_retrieval_top_k": 10, # max relevant MCP tools injected per turn }