From 03f5e25f96d2b924297a0d20107f1d95ce7fb386 Mon Sep 17 00:00:00 2001 From: dajiaohuang Date: Fri, 14 Aug 2026 02:30:59 +0800 Subject: [PATCH] feat(agent): budget model-visible MCP schemas --- docs/configuration.md | 22 ++- nanobot/agent/loop.py | 1 + nanobot/agent/runner.py | 12 +- nanobot/agent/subagent.py | 1 + nanobot/agent/tool_schema_selection.py | 180 ++++++++++++++++++++++ nanobot/config/schema.py | 5 + tests/agent/test_tool_schema_selection.py | 152 ++++++++++++++++++ 7 files changed, 371 insertions(+), 2 deletions(-) create mode 100644 nanobot/agent/tool_schema_selection.py create mode 100644 tests/agent/test_tool_schema_selection.py diff --git a/docs/configuration.md b/docs/configuration.md index 0a86c6e2959..66c05861ab8 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -2067,10 +2067,30 @@ Use `enabledTools` to register only a subset of tools from an MCP server: - Set `enabledTools` to `[]` to register no tools from that server. Resources and prompts are also skipped, since they have no per-name filter. - Set `enabledTools` to a non-empty list of names to register only those tools — resources and prompts are not registered. -MCP tools are automatically discovered and registered on startup. The LLM can use them alongside built-in tools — no extra configuration needed. +Large MCP setups can optionally limit the schemas sent to the model for each request: +```json +{ + "tools": { + "mcpSchemaBudgetBytes": 12000, + "mcpServers": { + "filesystem": { + "command": "npx", + "args": ["-y", "@modelcontextprotocol/server-filesystem", "/path/to/dir"] + } + } + } +} +``` +`mcpSchemaBudgetBytes` counts compact JSON schema bytes and is disabled by default (`0`). +When enabled, all built-in tools remain visible. If the registered MCP schemas exceed the +budget, nanobot chooses a deterministic subset from the latest user request and keeps that +view stable until another user message arrives. If relevance is unclear or the strongest +match cannot fit, it sends the full MCP schema set instead of hiding a potentially required +tool. This affects only model-visible definitions; all registered MCP tools remain executable. +MCP tools are automatically discovered and registered on startup. The LLM can use them alongside built-in tools — no extra configuration needed. ## Security diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 2532471687c..261d8c000e2 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -1088,6 +1088,7 @@ def _goal_continue() -> str | None: workspace=effective_scope.project_path, session_key=session.key if session else None, context_block_limit=self.context_block_limit, + mcp_schema_budget_bytes=self.tools_config.mcp_schema_budget_bytes, provider_retry_mode=self.provider_retry_mode, progress_callback=on_progress, stream_progress_deltas=on_stream is not None, diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index 7d03674b350..9455b461845 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -18,6 +18,7 @@ ContextGovernor, ) from nanobot.agent.hook import AgentHook, AgentHookContext, AgentRunHookContext +from nanobot.agent.tool_schema_selection import select_model_visible_tools from nanobot.agent.tools.registry import ToolRegistry, is_tool_error_result from nanobot.providers.base import ( LLMProvider, @@ -104,6 +105,7 @@ class AgentRunSpec: workspace: Path | None = None session_key: str | None = None context_block_limit: int | None = None + mcp_schema_budget_bytes: int = 0 provider_retry_mode: str = "standard" progress_callback: ProgressCallback | None = None stream_progress_deltas: bool = True @@ -445,6 +447,11 @@ async def _run_core( messages=messages, state=spec.provider_state, ) + model_visible_tools = select_model_visible_tools( + spec.tools.get_definitions(), + spec.initial_messages, + spec.mcp_schema_budget_bytes, + ) governance_config = ContextGovernanceConfig( provider=spec.runtime.provider, model=spec.runtime.model, @@ -485,6 +492,7 @@ async def _run_core( messages_for_model, hook, context, + model_visible_tools=model_visible_tools, conversation_state=conversation_state, provider_context=provider_context, ) @@ -899,6 +907,7 @@ async def _request_model( hook: AgentHook, context: AgentHookContext, *, + model_visible_tools: list[dict[str, Any]], malformed_retry: bool = False, conversation_state: ProviderConversationStateController, provider_context: ProviderCallContext | None = None, @@ -919,7 +928,7 @@ async def _request_model( kwargs = self._build_request_kwargs( spec, messages, - tools=spec.tools.get_definitions(), + tools=model_visible_tools, ) wants_streaming = hook.wants_streaming() progress_callback = spec.progress_callback @@ -1075,6 +1084,7 @@ async def _stream_progress(delta: str) -> None: ) return await self._request_model( spec, retry_messages, hook, context, + model_visible_tools=model_visible_tools, malformed_retry=True, conversation_state=conversation_state, provider_context=conversation_state.independent_request_context( diff --git a/nanobot/agent/subagent.py b/nanobot/agent/subagent.py index e46146ea1e6..074fa0f7057 100644 --- a/nanobot/agent/subagent.py +++ b/nanobot/agent/subagent.py @@ -416,6 +416,7 @@ async def _on_checkpoint(payload: dict[str, Any]) -> None: checkpoint_callback=_on_checkpoint, session_key=sess_key, workspace=root, + mcp_schema_budget_bytes=self.tools_config.mcp_schema_budget_bytes, llm_timeout_s=llm_timeout, )) finally: diff --git a/nanobot/agent/tool_schema_selection.py b/nanobot/agent/tool_schema_selection.py new file mode 100644 index 00000000000..b435976323c --- /dev/null +++ b/nanobot/agent/tool_schema_selection.py @@ -0,0 +1,180 @@ +"""Build a bounded model-visible view of registered MCP tool schemas.""" + +from __future__ import annotations + +import json +import re +from typing import Any, cast + +from nanobot.session.history_visibility import is_hidden_history_message + +_WORD_RE = re.compile(r"[a-zA-Z0-9]+") +_CAMEL_BOUNDARY_RE = re.compile(r"(?<=[a-z0-9])(?=[A-Z])") +_STOP_WORDS = frozenset({ + "a", + "an", + "and", + "are", + "as", + "at", + "be", + "by", + "can", + "do", + "for", + "from", + "i", + "in", + "is", + "it", + "me", + "mcp", + "of", + "on", + "or", + "please", + "the", + "this", + "to", + "tool", + "tools", + "use", + "with", + "you", +}) + + +def _tool_name(schema: dict[str, Any]) -> str: + function = schema.get("function") + if isinstance(function, dict): + name = cast(dict[str, Any], function).get("name") + if isinstance(name, str): + return name + name = schema.get("name") + return name if isinstance(name, str) else "" + + +def _terms(text: str) -> frozenset[str]: + expanded = _CAMEL_BOUNDARY_RE.sub(" ", text) + return frozenset( + token + for token in (match.group(0).lower() for match in _WORD_RE.finditer(expanded)) + if len(token) > 1 and token not in _STOP_WORDS + ) + + +def _text_content(content: Any) -> str: + if isinstance(content, str): + return content + if not isinstance(content, list): + return "" + parts: list[str] = [] + for item in cast(list[Any], content): + if not isinstance(item, dict): + continue + block = cast(dict[str, Any], item) + if block.get("type") == "text" and isinstance(block.get("text"), str): + parts.append(cast(str, block["text"])) + return "\n".join(parts) + + +def _latest_user_terms(messages: list[dict[str, Any]]) -> frozenset[str]: + for message in reversed(messages): + if message.get("role") != "user" or is_hidden_history_message(message): + continue + return _terms(_text_content(message.get("content"))) + return frozenset() + + +def schema_size_bytes(schema: dict[str, Any]) -> int: + """Return the deterministic compact-JSON UTF-8 size of one tool schema.""" + payload = json.dumps( + schema, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + ) + return len(payload.encode("utf-8")) + + +def schema_list_size_bytes(schemas: list[dict[str, Any]]) -> int: + """Return the compact-JSON UTF-8 size of a list of tool schemas.""" + if not schemas: + return 2 # [] + return 2 + sum(schema_size_bytes(schema) for schema in schemas) + len(schemas) - 1 + + +def _relevance(schema: dict[str, Any], query_terms: frozenset[str]) -> tuple[int, bool]: + name = _tool_name(schema) + name_terms = _terms(name) + name_overlap = query_terms & name_terms + + function = schema.get("function") + definition = cast(dict[str, Any], function) if isinstance(function, dict) else schema + description = definition.get("description") + description_terms: frozenset[str] = ( + _terms(description) if isinstance(description, str) else frozenset() + ) + description_overlap = query_terms & description_terms + + # Require two independent intent terms. A server name or generic action on + # its own is not enough evidence to silently hide other MCP capabilities. + clear = len(name_overlap | description_overlap) >= 2 + score = 100 * len(name_overlap) + 10 * len(description_overlap) + return score, clear + + +def select_model_visible_tools( + definitions: list[dict[str, Any]], + messages: list[dict[str, Any]], + mcp_schema_budget_bytes: int, +) -> list[dict[str, Any]]: + """Return built-ins plus a deterministic, budgeted MCP subset. + + A non-positive budget disables selection. Selection also fails open when + user intent has no clear lexical match or when the best match cannot fit. + The registry and executable tool set remain unchanged. + """ + if mcp_schema_budget_bytes <= 0: + return definitions + + builtins: list[dict[str, Any]] = [] + mcp_tools: list[dict[str, Any]] = [] + for schema in definitions: + target = mcp_tools if _tool_name(schema).startswith("mcp_") else builtins + target.append(schema) + + try: + sizes = {id(schema): schema_size_bytes(schema) for schema in mcp_tools} + except (TypeError, ValueError): + return definitions + if 2 + sum(sizes.values()) + max(0, len(mcp_tools) - 1) <= mcp_schema_budget_bytes: + return definitions + + query_terms = _latest_user_terms(messages) + if not query_terms: + return definitions + + ranked: list[tuple[int, str, dict[str, Any]]] = [] + for schema in mcp_tools: + score, clear = _relevance(schema, query_terms) + if clear: + ranked.append((score, _tool_name(schema), schema)) + if not ranked: + return definitions + + ranked.sort(key=lambda item: (-item[0], item[1])) + if sizes[id(ranked[0][2])] + 2 > mcp_schema_budget_bytes: + return definitions + + selected: list[dict[str, Any]] = [] + used = 2 # JSON list brackets + for _score, _name, schema in ranked: + size = sizes[id(schema)] + separator_size = 1 if selected else 0 + if used + separator_size + size <= mcp_schema_budget_bytes: + selected.append(schema) + used += separator_size + size + + selected.sort(key=_tool_name) + return builtins + selected diff --git a/nanobot/config/schema.py b/nanobot/config/schema.py index 92c97783ced..806ff8639d2 100644 --- a/nanobot/config/schema.py +++ b/nanobot/config/schema.py @@ -407,6 +407,11 @@ class ToolsConfig(Base): image_generation: ImageGenerationToolConfig = Field( default_factory=lambda: _lazy_default("nanobot.agent.tools.image_generation", "ImageGenerationToolConfig"), ) + mcp_schema_budget_bytes: int = Field( + default=0, + ge=0, + validation_alias=AliasChoices("mcpSchemaBudgetBytes", "mcp_schema_budget_bytes"), + ) # Opt-in byte budget for model-visible MCP schemas; 0 keeps all schemas visible restrict_to_workspace: bool = False # policy intent: keep tool access inside workspace when possible webui_allow_local_service_access: bool = Field( default=True, diff --git a/tests/agent/test_tool_schema_selection.py b/tests/agent/test_tool_schema_selection.py new file mode 100644 index 00000000000..7a455c49980 --- /dev/null +++ b/tests/agent/test_tool_schema_selection.py @@ -0,0 +1,152 @@ +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from agent.runner_helpers import make_run_spec +from nanobot.agent.runner import AgentRunner +from nanobot.agent.tool_schema_selection import ( + schema_list_size_bytes, + select_model_visible_tools, +) +from nanobot.config.schema import AgentDefaults, Config, ToolsConfig +from nanobot.providers.base import LLMProvider, LLMResponse + + +def _schema(name: str, description: str, detail: str = "") -> dict: + return { + "type": "function", + "function": { + "name": name, + "description": description, + "parameters": { + "type": "object", + "properties": { + "detail": {"type": "string", "description": detail}, + }, + }, + }, + } + + +def _names(definitions: list[dict]) -> list[str]: + return [schema["function"]["name"] for schema in definitions] + + +def test_mcp_schema_budget_is_opt_in_and_accepts_camel_case() -> None: + assert ToolsConfig().mcp_schema_budget_bytes == 0 + + config = Config.model_validate({"tools": {"mcpSchemaBudgetBytes": 4096}}) + + assert config.tools.mcp_schema_budget_bytes == 4096 + assert config.model_dump(by_alias=True)["tools"]["mcpSchemaBudgetBytes"] == 4096 + + +def test_selection_reduces_only_mcp_schemas_and_recalls_required_tool() -> None: + builtin = _schema("read_file", "Read a local file") + forecast = _schema( + "mcp_weather_forecast", + "Return the weather forecast for a location", + "Include hourly weather details", + ) + history = _schema( + "mcp_weather_history", + "Return historical weather observations", + "Include archived observations", + ) + calendar = _schema("mcp_calendar_events", "List calendar events", "Date range") + definitions = [builtin, forecast, history, calendar] + budget = schema_list_size_bytes([forecast]) + + selected = select_model_visible_tools( + definitions, + [{"role": "user", "content": "Show the weather forecast for Singapore"}], + budget, + ) + + assert _names(selected) == ["read_file", "mcp_weather_forecast"] + assert schema_list_size_bytes(selected[1:]) <= budget + assert schema_list_size_bytes(selected[1:]) < schema_list_size_bytes(definitions[1:]) + + +def test_selection_is_deterministic_for_latest_user_message() -> None: + calendar = _schema("mcp_calendar_events", "List calendar events") + weather = _schema("mcp_weather_forecast", "Return a weather forecast") + definitions = [calendar, weather] + messages = [ + {"role": "user", "content": "Show calendar events"}, + {"role": "assistant", "content": "What next?"}, + {"role": "user", "content": "Now show the weather forecast"}, + ] + budget = schema_list_size_bytes([weather]) + + first = select_model_visible_tools(definitions, messages, budget) + second = select_model_visible_tools(definitions, messages, budget) + + assert _names(first) == ["mcp_weather_forecast"] + assert first == second + + +@pytest.mark.parametrize( + "message,budget", + [ + ("Help me with this task", 100), + ("Show the weather forecast", 1), + ], +) +def test_selection_fails_open_when_relevance_or_budget_is_unsafe( + message: str, + budget: int, +) -> None: + definitions = [ + _schema("mcp_weather_forecast", "Return a weather forecast"), + _schema("mcp_calendar_events", "List calendar events"), + ] + + selected = select_model_visible_tools( + definitions, + [{"role": "user", "content": message}], + budget, + ) + + assert selected is definitions + + +def test_selection_fails_open_for_server_name_without_an_operation() -> None: + definitions = [ + _schema("mcp_github_create_issue", "Create an issue in a repository"), + _schema("mcp_github_merge_pull_request", "Merge a pull request"), + ] + + selected = select_model_visible_tools( + definitions, + [{"role": "user", "content": "Use GitHub for this"}], + schema_list_size_bytes([definitions[0]]), + ) + + assert selected is definitions + + +@pytest.mark.asyncio +async def test_runner_sends_budgeted_view_without_changing_registry() -> None: + provider = MagicMock(spec=LLMProvider) + provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="done")) + builtin = _schema("read_file", "Read a local file") + weather = _schema("mcp_weather_forecast", "Return a weather forecast") + calendar = _schema("mcp_calendar_events", "List calendar events") + definitions = [builtin, weather, calendar] + tools = MagicMock() + tools.get_definitions.return_value = definitions + + await AgentRunner().run(make_run_spec( + provider, + initial_messages=[{"role": "user", "content": "Show the weather forecast"}], + tools=tools, + model="test-model", + max_iterations=1, + max_tool_result_chars=AgentDefaults().max_tool_result_chars, + mcp_schema_budget_bytes=schema_list_size_bytes([weather]), + )) + + sent = provider.chat_with_retry.await_args.kwargs["tools"] + assert _names(sent) == ["read_file", "mcp_weather_forecast"] + assert tools.get_definitions.return_value is definitions