diff --git a/packages/bub-acp-server/src/bub_acp_server/agent.py b/packages/bub-acp-server/src/bub_acp_server/agent.py index d4c1003..85bfe54 100644 --- a/packages/bub-acp-server/src/bub_acp_server/agent.py +++ b/packages/bub-acp-server/src/bub_acp_server/agent.py @@ -962,7 +962,7 @@ async def _run_streaming_subagent( ) child_state = {**context.state, "session_id": child_session} allowed_tools = resolve_tool_names( - param.allowed_tools or None, exclude={"subagent"}, all_names=agent.tools + param.allowed_tools or None, exclude={"subagent"}, all_names=agent.known_tools ) await emitter.started( description=_subagent_description(param.prompt), model_id=param.model diff --git a/packages/bub-acp-server/tests/test_mcp.py b/packages/bub-acp-server/tests/test_mcp.py index 8b8390f..d8bd3d0 100644 --- a/packages/bub-acp-server/tests/test_mcp.py +++ b/packages/bub-acp-server/tests/test_mcp.py @@ -112,7 +112,7 @@ async def call_echo(agent, session_id): async def runtime_tools(agent, session_id): inbound = agent._build_inbound([], agent._sessions[session_id]) state = await agent.framework.build_state(inbound, inbound.session_id) - return state["_runtime_agent"].tools + return state["_runtime_agent"].known_tools def test_transport_configs_preserve_arguments_environment_headers_and_cwd(tmp_path): @@ -202,8 +202,8 @@ async def test_model_receives_only_current_session_mcp_tools( seen = [] - def run(self, *, tools, **kwargs): - seen.append({tool.name for tool in tools}) + def run(self, *, tools, system_prompt, **kwargs): + seen.append(({tool.name for tool in tools}, system_prompt)) async def events(): yield StreamEvent("text", {"delta": "ok"}) @@ -228,8 +228,14 @@ async def events(): session_id=session.session_id, prompt=[TextContentBlock(text="hello")] ) # Bub aliases dots for the model, but execution still uses each session's tool object. - assert "mcp_first_echo" in seen[0] and "mcp_second_echo" not in seen[0] - assert "mcp_second_echo" in seen[1] and "mcp_first_echo" not in seen[1] + for (names, prompt), own, other in zip( + seen, + ["mcp_first_echo", "mcp_second_echo"], + ["mcp_second_echo", "mcp_first_echo"], + strict=True, + ): + assert "mcp_describe" in names and not {own, other} & names + assert own in prompt and other not in prompt finally: await agent.shutdown() @@ -474,7 +480,7 @@ async def test_session_tools_override_configured_mcp_and_restore_on_close( assert len(clients[1].calls) == 2 runtime_agent = agent._runtime_agents[session.session_id] await agent.close_session(session.session_id) - assert await runtime_agent.tools["mcp.shared_echo"].run() == "configured" + assert await runtime_agent.known_tools["mcp.shared_echo"].run() == "configured" await configured._manager.stop() assert "mcp.shared_echo" not in runtime_agent.tools finally: diff --git a/packages/bub-acp-server/tests/test_subagents.py b/packages/bub-acp-server/tests/test_subagents.py index 1e1f264..bb5b80b 100644 --- a/packages/bub-acp-server/tests/test_subagents.py +++ b/packages/bub-acp-server/tests/test_subagents.py @@ -73,6 +73,7 @@ def __init__(self, events: list[StreamEvent], *, error: Exception | None = None) self.error = error self.calls: list[dict[str, Any]] = [] self.tools = {"subagent": REGISTRY["subagent"], "bash": REGISTRY["bash"]} + self.known_tools = self.tools async def run_stream(self, **kwargs: Any) -> AsyncStreamEvents: self.calls.append(kwargs) diff --git a/packages/bub-mcp/README.md b/packages/bub-mcp/README.md index e1f8cab..8a50437 100644 --- a/packages/bub-mcp/README.md +++ b/packages/bub-mcp/README.md @@ -32,6 +32,22 @@ The file must contain a top-level `mcpServers` mapping: } ``` +## Tool definitions and selection + +In ordinary mode, MCP tools start as names and short summaries. `mcp_describe` +exposes selected native definitions on the next model call. Load missing tools +from any source together in one call; their original handlers remain unchanged. +Definitions follow first-discovery order in each tape context and survive restart. +New sessions, reset and handoff start fresh. Summaries stay stable as definitions +load. Builtins remain directly available; code mode uses its complete allowed stub. + +Configure `mcp.allowed_tools` and `mcp.excluded_tools` in Bub's `config.yml`, or set +`BUB_MCP_ALLOWED_TOOLS` and `BUB_MCP_EXCLUDED_TOOLS` to JSON arrays. Patterns accept +runtime names (`mcp.notes_lookup`), model aliases (`mcp_notes_lookup`) and shell +wildcards. Exclusions win; an omitted allowlist allows all, `[]` allows none. +Restart Bub after changing configuration. Filtering limits discovery, native calls +and code mode; server discovery and listing remain unchanged. + ## CLI Usage Use the CLI to inspect and manage `mcp.json`: @@ -97,12 +113,14 @@ finally: Discovered tools belong to the channel and are available through `channel.tools`; discovery does not modify `bub.tools.REGISTRY`. Bub snapshots that registry when an `Agent` is created, so changing it afterward would not update existing agents. The plugin's `load_state` hook waits for startup -discovery and binds tools to the turn's Agent (including an explicit `_runtime_agent`). Stopping -the channel removes its bindings and restores any tools it replaced. +discovery and binds the inventory and provider to the turn's Agent (including an explicit `_runtime_agent`). Remote +tools are selected per request without changing `Agent.tools`; comma commands resolve them +through the discovery inventory. Code mode loads the allowed catalog into its complete stub. +Stopping the channel removes its inventory, revealing any tools shadowed by that source. An embedding application can call `channel.bind_agent(agent)` after startup discovery completes, or `await channel.bind_runtime_tools(framework, message)` from its own `load_state` hook. Both paths -use instance tools. This requires Bub's instance-tool API introduced in upstream PR #311. +use instance tools. This requires Bub's tool-discovery API introduced in upstream PR #339. The default `MCPChannel()` behavior remains backed by Bub's `mcp.json`. Read-only channels reject `add()` and `remove()` so an embedding plugin remains the owner of its source configuration. The diff --git a/packages/bub-mcp/src/bub_mcp/config.py b/packages/bub-mcp/src/bub_mcp/config.py index 9288211..b4e2d86 100644 --- a/packages/bub-mcp/src/bub_mcp/config.py +++ b/packages/bub-mcp/src/bub_mcp/config.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +from fnmatch import fnmatchcase from pathlib import Path from typing import Any @@ -28,6 +29,28 @@ class MCPSettings(bub.Settings): config_path: Path = Field(default_factory=default_config_path) init_timeout_seconds: float | None = 20.0 + allowed_tools: list[str] | None = Field( + None, + description="Tool name patterns to expose; None allows all, an empty list allows none.", + ) + excluded_tools: list[str] = Field( + default_factory=list, + description="Tool name patterns to exclude, even when allowed.", + ) + + def allows_tool(self, name: str) -> bool: + """Match runtime names or model aliases without changing remote schemas.""" + alias = name.replace(".", "_") + + def matches(patterns: list[str]) -> bool: + return any( + fnmatchcase(name, pattern) or fnmatchcase(alias, pattern) + for pattern in patterns + ) + + return ( + self.allowed_tools is None or matches(self.allowed_tools) + ) and not matches(self.excluded_tools) def read_mcp_servers(self) -> dict[str, Any]: return read_config(self.config_path) diff --git a/packages/bub-mcp/src/bub_mcp/plugin.py b/packages/bub-mcp/src/bub_mcp/plugin.py index 1a15fd9..f482449 100644 --- a/packages/bub-mcp/src/bub_mcp/plugin.py +++ b/packages/bub-mcp/src/bub_mcp/plugin.py @@ -7,7 +7,7 @@ from copy import deepcopy from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal, Self, TypedDict, final -from weakref import WeakKeyDictionary +from weakref import WeakSet import fastmcp import mcp.types @@ -18,6 +18,7 @@ from bub.tools import Tool, ToolContext from bub.channels.contracts import MessageHandler from bub.envelope import Envelope, field_of +from bub.tape import Tape from bub.turn import TurnState from loguru import logger @@ -28,6 +29,8 @@ from bub.framework import BubFramework TOOL_PREFIX = "mcp." +MCP_TOOLS_STATE_KEY = "_mcp_tools" +DEFINITIONS_LOADED_EVENT = "mcp.definitions.loaded" TEXT_OUTPUT_SCHEMA: dict[str, Any] = {"type": "string"} type ResultMode = Literal["text", "structured", "wrapped"] @@ -159,9 +162,7 @@ def __init__(self) -> None: self._lock = asyncio.Lock() self._bootstrap_task: asyncio.Task[None] | None = None self._servers: dict[str, MCPServerState] = {} - self._bindings: WeakKeyDictionary[ - Agent, dict[str, tuple[Tool, Tool | None]] - ] = WeakKeyDictionary() + self._agents: WeakSet[Agent] = WeakSet() self._stop_event: asyncio.Event | None = None @classmethod @@ -210,9 +211,13 @@ async def stop(self) -> None: server.client = None server.connected = False - for agent, bindings in list(self._bindings.items()): - self._restore_tools(agent, bindings) - self._bindings.clear() + for agent in self._agents: + agent.tool_sources.pop(self, None) + if not any(isinstance(source, MCPChannel) for source in agent.tool_sources): + agent.tool_providers.remove(prepare_mcp_tools) + if agent.tools.get(mcp_describe.name) is mcp_describe: + agent.tools.pop(mcp_describe.name) + self._agents.clear() for client in clients: await self._close_client(client) @@ -225,35 +230,17 @@ def tools(self) -> dict[str, Tool]: for server in self._servers.values() if server.connected for tool in server.tools + if self.settings.allows_tool(tool.name) } def bind_agent(self, agent: Agent) -> None: - """Refresh this channel's tools on an Agent, preserving name collisions.""" - previous = self._bindings.pop(agent, {}) - tools = self.tools - self._restore_tools( - agent, - {name: binding for name, binding in previous.items() if name not in tools}, - ) - bindings = {} - for name, remote_tool in tools.items(): - original = previous[name][1] if name in previous else agent.tools.get(name) - bindings[name] = (remote_tool, original) - agent.tools[name] = remote_tool - if bindings: - self._bindings[agent] = bindings - - @staticmethod - def _restore_tools( - agent: Agent, bindings: dict[str, tuple[Tool, Tool | None]] - ) -> None: - for name, (installed, original) in bindings.items(): - if agent.tools.get(name) is not installed: - continue - if original is None: - agent.tools.pop(name, None) - else: - agent.tools[name] = original + """Register discovery definitions separately from request preparation.""" + self._agents.add(agent) + # Explicit tool sets may omit the globally registered discovery helper. + agent.tools.setdefault(mcp_describe.name, mcp_describe) + agent.tool_sources[self] = self.tools + if prepare_mcp_tools not in agent.tool_providers: + agent.tool_providers.append(prepare_mcp_tools) async def bind_runtime_tools( self, framework: BubFramework, message: Envelope @@ -542,3 +529,65 @@ def mcp_list(*, context: ToolContext) -> MCPServerList: for name, server in manager.list().items() } } + + +@tool(name="mcp.describe", context=True, preserve=True) +async def mcp_describe(names: list[str], *, context: ToolContext) -> str: + """Expose complete native definitions for selected MCP tools on the next model call. + + Use exact names from the mcp_tools catalog. Call already available tools directly. + Load the missing definitions needed for the current task together in one call. + """ + from bub.builtin.tools import resolve_tool_names + + available: dict[str, Tool] = context.state.get(MCP_TOOLS_STATE_KEY, {}) + resolved = resolve_tool_names(names, all_names=available) + if not resolved: + raise ValueError("provide at least one available MCP tool name") + await context.tape.append_event( + DEFINITIONS_LOADED_EVENT, {"names": sorted(resolved)}, context=False + ) + aliases = ", ".join(name.replace(".", "_") for name in sorted(resolved)) + return f"Complete native definitions are now available for: {aliases}. Call these tools directly." + + +async def prepare_mcp_tools(tools: list[Tool], tape: Tape) -> tuple[list[Tool], str]: + """Prepare scoped MCP tools across sources in the tape's discovery order.""" + agent = tape.context.state["_runtime_agent"] + known = agent.known_tools + owned = { + name: item + for source, definitions in agent.tool_sources.items() + if isinstance(source, MCPChannel) + for name, item in definitions.items() + if known.get(name) is item + } + available = { + item.name: item for item in tools if item.name in owned and item.agent_use + } + native = [ + item + for item in tools + if item.name not in owned and item.name != mcp_describe.name + ] + tape.context.state[MCP_TOOLS_STATE_KEY] = available + code_mode = tape.context.state.get("code_mode") and any( + item.name == "run_code" for item in native + ) + if code_mode: + return [*native, *available.values()], "" + if available: + native.append(mcp_describe) + loaded: dict[str, None] = {} + query = tape.context.build_query(tape.query()).kinds("event") + for entry in await tape.store.fetch_all(query): + if entry.payload.get("name") == DEFINITIONS_LOADED_EVENT: + loaded.update(dict.fromkeys(entry.payload["data"]["names"])) + native.extend(available[name] for name in loaded if name in available) + lines = [] + for name, item in available.items(): + summary = item.description.strip().split("\n", 1)[0].split(". ", 1)[0][:180] + lines.append( + f"- {name.replace('.', '_')}" + (f": {summary}" if summary else "") + ) + return native, f"\n{'\n'.join(lines)}\n" if lines else "" diff --git a/packages/bub-mcp/tests/test_plugin.py b/packages/bub-mcp/tests/test_plugin.py index fbdd158..1adcdf9 100644 --- a/packages/bub-mcp/tests/test_plugin.py +++ b/packages/bub-mcp/tests/test_plugin.py @@ -1,11 +1,17 @@ from __future__ import annotations import asyncio +import json +from dataclasses import replace from pathlib import Path +from typing import Any import pytest +from any_llm.types.completion import ChatCompletion from bub.builtin.agent import Agent +from bub.builtin.codemode import run_code from bub.framework import BubFramework +from bub.tape import Tape from bub.tools import REGISTRY from bub_mcp import plugin @@ -502,7 +508,7 @@ async def test_late_discovery_reaches_existing_agent_and_stop_restores_tools( assert output == ["forecast for Paris"] assert REGISTRY == original_registry assert original.name not in unrelated.tools - # Repeated binding must not save our own tool as the collision fallback. + # Repeated discovery keeps the same source precedence. await framework.build_state({}, "test:weather") finally: await channel.stop() @@ -523,7 +529,7 @@ async def test_mcp_binds_explicit_runtime_agent_without_changing_default( channel._servers["embedded"] = plugin.MCPServerState(tools=[remote], connected=True) await channel.bind_runtime_tools(framework, {"_runtime_agent": embedded}) - assert embedded.tools[remote.name] is remote + assert embedded.known_tools[remote.name].run() == "ok" assert remote.name not in default.tools assert remote.name not in REGISTRY # Shutdown must not remove a replacement installed by someone else. @@ -531,3 +537,172 @@ async def test_mcp_binds_explicit_runtime_agent_without_changing_default( embedded.tools[remote.name] = replacement await channel.stop() assert embedded.tools[remote.name] is replacement + + +@pytest.mark.parametrize("code_mode", [False, True]) +@pytest.mark.parametrize( + ("allowed", "excluded", "visible"), + [ + (None, [], ["alpha", "beta"]), + ([], [], []), + (["mcp.*"], ["mcp_beta_weather_get_forecast"], ["alpha"]), + (["mcp_alpha_weather_get_forecast"], ["mcp.alpha_weather_get_forecast"], []), + ], +) +async def test_mcp_discovery_scope_and_cleanup_with_independent_sources( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + code_mode: bool, + allowed: list[str] | None, + excluded: list[str], + visible: list[str], +) -> None: + monkeypatch.setenv("BUB_HOME", str(tmp_path)) + framework = BubFramework(config_file=tmp_path / "config.yml") + framework.workspace = tmp_path + framework.load_builtin_hooks() + monkeypatch.setattr(plugin, "_create_fastmcp_client", FakeClient) + replies: list[str | tuple[str, dict[str, Any]]] = [] + requests: list[dict[str, Any]] = [] + + class Provider: + SUPPORTS_COMPLETION_STREAMING = False + + async def acompletion(self, **kwargs: Any) -> ChatCompletion: + requests.append(kwargs) + reply = replies.pop(0) + message: dict[str, Any] = {"role": "assistant", "content": reply} + if isinstance(reply, tuple): + name, arguments = reply + function = {"name": name, "arguments": json.dumps(arguments)} + call = { + "id": str(len(requests)), + "type": "function", + "function": function, + } + message = {"role": "assistant", "tool_calls": [call]} + choice = { + "index": 0, + "message": message, + "finish_reason": "stop" if isinstance(reply, str) else "tool_calls", + } + return ChatCompletion( + id="reply", + model="test", + created=0, + object="chat.completion", + choices=[choice], + ) + + monkeypatch.setattr( + "bub.builtin.model_runner.AnyLLM.create", lambda *a, **k: Provider() + ) + channels = [ + plugin.MCPChannel.from_server_configs({name: {"command": name}}) + for name in ("alpha", "beta") + ] + for channel in channels: + channel.settings = plugin.MCPSettings( + allowed_tools=allowed, excluded_tools=excluded + ) + + async def uppercase( + tools: list[plugin.Tool], tape: Tape + ) -> tuple[list[plugin.Tool], str]: + return [ + replace(item, renderer=str.upper) if item.name.startswith("mcp.") else item + for item in tools + ], "" + + async def run(agent: Agent, session: str = "task") -> list[dict[str, Any]]: + requests.clear() + replies.append("done") + stream = await agent.run_stream( + session_id=session, + prompt="Read forecasts.", + model="openrouter:test", + state={"code_mode": code_mode}, + ) + async for _ in stream: + pass + return list(requests) + + def names(request: dict[str, Any]) -> list[str]: + return [item["function"]["name"] for item in request.get("tools") or []] + + def call(name: str) -> tuple[str, dict[str, Any]]: + alias = f"mcp_{name}_weather_get_forecast" + return ( + ("run_code", {"code": f"print(await tools.{alias}(city='{name}'))"}) + if code_mode + else (alias, {"city": name}) + ) + + try: + for channel in channels: + await channel.connect() + agent = Agent(framework, tools=[REGISTRY["tape.info"], run_code], skill_dirs=[]) + agent.tool_providers.append(uppercase) + for channel in channels: + channel.bind_agent(agent) + replies.append( + ("run_code", {"code": "print(await tools.tape_info())"}) + if code_mode + else ("tape_info", {}) + ) + for name in reversed(visible): + if not code_mode: + replies.append( + ("mcp_describe", {"names": [f"mcp_{name}_weather_get_forecast"]}) + ) + replies.append(call(name)) + first = await run(agent) + assert not any( + name.startswith("mcp_") and name != "mcp_describe" + for name in names(first[0]) + ) + if not code_mode: + assert names(first[-1]) == [ + *names(first[0]), + *(f"mcp_{name}_weather_get_forecast" for name in reversed(visible)), + ] + prompts = [ + next(m["content"] for m in r["messages"] if m["role"] == "system") + for r in first + ] + for name in ("alpha", "beta"): + assert (f"mcp_{name}_weather_get_forecast" in prompts[0]) == ( + name in visible + ) + assert all(prompt == prompts[0] for prompt in prompts) + results = "\n".join( + m.get("content", "") for m in first[-1]["messages"] if m["role"] == "tool" + ) + assert "entries" in results + for name in visible: + expected = f"forecast for {name}" + assert (expected if code_mode else expected.upper()) in results + if visible and not code_mode: + replies.append(call(visible[0])) + fresh = await run(agent, "fresh") + if visible and not code_mode: + assert any( + "does not exist" in m.get("content", "") + for m in fresh[-1]["messages"] + if m["role"] == "tool" + ) + for name, channel in zip(("alpha", "beta"), channels, strict=True): + assert channel.list()[name].client.tool_calls == ( + [("weather_get_forecast", {"city": name})] if name in visible else [] + ) + assert not any( + n.startswith("mcp_") and n != "mcp_describe" for n in names(fresh[0]) + ) + for channel, name in zip(channels, ("alpha", "beta"), strict=True): + await channel.stop() + remaining = await run(agent) + assert f"mcp_{name}_weather_get_forecast" not in names(remaining[0]) + assert "mcp_describe" not in names(remaining[0]) + finally: + for channel in channels: + await channel.stop() diff --git a/uv.lock b/uv.lock index 2cf6eed..f1be618 100644 --- a/uv.lock +++ b/uv.lock @@ -371,8 +371,8 @@ wheels = [ [[package]] name = "bub" -version = "0.5.0" -source = { git = "https://github.com/bubbuild/bub.git#8669780e7c6754f5d445b2f0c655ea7c01bf00c6" } +version = "0.5.1.dev10+g2b23812fa" +source = { git = "https://github.com/bubbuild/bub.git#2b23812faf7055669b52d880e74a37cf5a8761a2" } dependencies = [ { name = "aiohttp" }, { name = "any-llm-sdk" },