refactor: separate tool discovery from request preparation

src/bub/builtin/agent.py

@@ -7,7 +7,7 @@
 import re
 import shlex
 import time
-from collections.abc import AsyncGenerator, AsyncIterator, Collection, Iterable
+from collections.abc import AsyncGenerator, AsyncIterator, Collection, Iterable, Mapping
 from contextlib import AsyncExitStack, aclosing
 from dataclasses import dataclass, replace
 from datetime import UTC, datetime
@@ -36,6 +36,8 @@
 
 HINT_RE = re.compile(r"\$([A-Za-z0-9_.-]+)")
 MAX_AUTO_HANDOFF_RETRIES = 1
+# Model aliases of the tools prepared for the current request, used to recover unknown tool calls.
+REQUEST_TOOLS_STATE_KEY = "_runtime_request_tools"
 
 
 class Agent:
@@ -73,11 +75,22 @@ def __init__(
         )
         self.framework = framework
         self.tools = {tool.name: tool for tool in tools} if tools is not None else REGISTRY.copy()
+        self.tool_sources: dict[object, Mapping[str, Tool]] = {}
         self.tool_providers: list[ToolProvider] = []
         self.tape_store = tape_store
         self.skill_dirs = skill_dirs
         self.model_runner = ModelRunner(self.settings, hooks=framework.get_agent_hooks())
 
+    @property
+    def known_tools(self) -> dict[str, Tool]:
+        """Known runtime tools; later registered sources take precedence."""
+        tools = self.tools.copy()
+        for source in self.tool_sources.values():
+            for name, item in source.items():
+                tools.pop(name, None)
+                tools[name] = item
+        return tools
+
     @cached_property
     def tape(self) -> Tape:
         """Return the lazily constructed, cached tape factory for this agent.
@@ -247,13 +260,14 @@ async def _run_command(self, tape: Tape, *, line: str) -> str:
         output = ""
         status = "ok"
         try:
-            if name not in self.tools:
-                if "bash" not in self.tools:
+            known_tools = self.known_tools
+            if name not in known_tools:
+                if "bash" not in known_tools:
                     raise ValueError("bash tool is not available")  # noqa: TRY301
-                bash_tool = self.tools["bash"]
+                bash_tool = known_tools["bash"]
                 output = bash_tool.render(await bash_tool.run(context=context, command=line))
             else:
-                command_tool = self.tools[name]
+                command_tool = known_tools[name]
                 args = _parse_args(arg_tokens)
                 if command_tool.context:
                     args.kwargs["context"] = context
@@ -445,17 +459,18 @@ async def _run_once(
             prompt_text = ""
         else:
             prompt_text = _extract_text_from_parts(prompt)
+        known_tools = self.known_tools
         if allowed_tools is not None:
             from bub.builtin.tools import resolve_tool_names
 
-            allowed_tools = resolve_tool_names(allowed_tools, all_names=self.tools)
+            allowed_tools = resolve_tool_names(allowed_tools, all_names=known_tools)
         if allowed_skills is not None:
             allowed_skills = {name.casefold() for name in allowed_skills}
             tape.context.state["allowed_skills"] = list(allowed_skills)
         if allowed_tools is not None:
-            tools = [tool for tool in self.tools.values() if tool.name in allowed_tools]
+            tools = [tool for tool in known_tools.values() if tool.name in allowed_tools]
         else:
-            tools = list(self.tools.values())
+            tools = list(known_tools.values())
         return await self._run_once_stream(
             tape=tape,
             prompt=prompt,
@@ -476,10 +491,14 @@ async def _run_once_stream(
         tools: list[Tool],
     ) -> AsyncStreamEvents:
         tools_prompts: list[str] = []
-        for provider in (self._prepare_code_mode, *self.tool_providers):
+        for provider in self.tool_providers:
             tools, tools_prompt = await provider(tools, tape)
             if tools_prompt:
                 tools_prompts.append(tools_prompt)
+        tape.context.state[REQUEST_TOOLS_STATE_KEY] = [item.name for item in model_tools(tools)]
+        tools, tools_prompt = await self._prepare_code_mode(tools, tape)
+        if tools_prompt:
+            tools_prompts.append(tools_prompt)
         system_prompt = self._system_prompt(
             prompt_text,
             state=tape.context.state,

src/bub/builtin/hook_impl.py

@@ -9,7 +9,7 @@
 from loguru import logger
 
 from bub import inquirer as bub_inquirer
-from bub.builtin.agent import Agent
+from bub.builtin.agent import REQUEST_TOOLS_STATE_KEY, Agent
 from bub.builtin.commands import strip_command_prefix
 from bub.builtin.context import default_tape_context
 from bub.builtin.onboarding import collect_model_config
@@ -358,9 +358,11 @@ async def before_tool_call(
         """
         from bub.tools import model_tools
 
-        agent = self._get_agent(state)
-
-        available_tools = tuple(tool_item.name for tool_item in model_tools(agent.tools.values()))
+        if (request_tools := state.get(REQUEST_TOOLS_STATE_KEY)) is not None:
+            available_tools = tuple(request_tools)
+        else:
+            agent = self._get_agent(state)
+            available_tools = tuple(tool_item.name for tool_item in model_tools(agent.known_tools.values()))
         if call.tool in available_tools:
             return None
 

src/bub/builtin/tools.py

@@ -434,7 +434,7 @@ async def run_subagent(param: SubAgentInput, *, context: ToolContext) -> SubAgen
     else:
         subagent_session = param.session
     state = {**context.state, "session_id": subagent_session}
-    allowed_tools = resolve_tool_names(param.allowed_tools or None, exclude={"subagent"}, all_names=agent.tools)
+    allowed_tools = resolve_tool_names(param.allowed_tools or None, exclude={"subagent"}, all_names=agent.known_tools)
     output = ""
     errors: list[str] = []
     stream = await agent.run_stream(

tests/test_builtin_agent.py

@@ -55,6 +55,7 @@ async def build_prompt(message: dict[str, Any], session_id: str, state: dict[str
     agent.command_prefix = agent.settings.command_prefix
     agent.framework = framework
     agent.tools = REGISTRY.copy()
+    agent.tool_sources = {}
     agent.tool_providers = []
     agent.tape_store = None
     agent.skill_dirs = None

tests/test_builtin_hook_impl.py

@@ -41,6 +41,7 @@ def __init__(self, home: Path, *, tape: Tape | None = None) -> None:
         self.command_prefix = ","
         self.settings = SimpleNamespace(home=home)
         self.tools = REGISTRY.copy()
+        self.known_tools = self.tools
         # A real in-memory async tape so load_state's recovery path runs against
         # the same store the tests write `model_switch` events to.
         self.tape = tape if tape is not None else _fake_tape(home)
@@ -493,3 +494,21 @@ async def _do():
     assert decision is not None
     assert "fs_reed" in decision.result
     assert "fs_read" in decision.result
+
+
+def test_before_tool_call_recovers_tool_outside_current_request(tmp_path: Path) -> None:
+    _, impl, _ = _build_impl(tmp_path)
+    import asyncio
+
+    from bub.builtin.agent import REQUEST_TOOLS_STATE_KEY
+    from bub.hooks.interception import ToolCall
+
+    state = {REQUEST_TOOLS_STATE_KEY: ["bash"]}
+
+    async def _do(name: str):
+        return await impl.before_tool_call(ToolCall(run_id="r", tool=name, arguments={}), state=state)
+
+    assert asyncio.run(_do("bash")) is None
+    decision = asyncio.run(_do("bash_output"))
+    assert decision is not None and decision.action == "replace"
+    assert "bash_output" in decision.result

tests/test_subagent_tool.py

@@ -21,6 +21,7 @@ def __init__(self, state: dict[str, Any]) -> None:
 class FakeAgent:
     def __init__(self) -> None:
         self.tools = REGISTRY.copy()
+        self.known_tools = self.tools
         self.run_stream = AsyncMock(side_effect=self._run_stream)
 
     async def _run_stream(self, **kwargs: Any) -> AsyncStreamEvents:

tests/test_tool_providers.py

@@ -1,28 +1,33 @@
 from __future__ import annotations
 
 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 Tool
 
 
 @pytest.mark.asyncio
-async def test_provider_prompt_reaches_the_model_and_registered_tools_remain_callable(
-    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+@pytest.mark.parametrize("code_mode", [False, True])
+@pytest.mark.parametrize("reverse_providers", [False, True])
+async def test_discovery_precedence_is_independent_of_provider_order_for_native_and_code_calls(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, code_mode: bool, reverse_providers: bool
 ) -> None:
     monkeypatch.setenv("BUB_HOME", str(tmp_path))
     framework = BubFramework(config_file=tmp_path / "config.yml")
     framework.workspace = tmp_path
     framework.load_builtin_hooks()
     direct = Tool.from_callable(lambda: "direct", name="direct")
     denied = Tool.from_callable(lambda: "denied", name="denied")
+    pending = Tool.from_callable(lambda: "pending", name="provider.pending")
     calls: list[str] = []
 
     def lookup(name: str) -> str:
@@ -31,8 +36,13 @@ def lookup(name: str) -> str:
 
     supplied = Tool.from_callable(lookup, name="provider.lookup")
 
-    async def provide(tools: list[Tool], tape: Tape) -> tuple[list[Tool], str]:
-        return tools, "Use provider_lookup to greet the requested person."
+    async def select(tools: list[Tool], tape: Tape) -> tuple[list[Tool], str]:
+        return [item for item in tools if item is not pending], "Use provider_lookup to greet the requested person."
+
+    async def guide(tools: list[Tool], tape: Tape) -> tuple[list[Tool], str]:
+        return [
+            replace(item, renderer=str.upper) if item.name == supplied.name else item for item in tools
+        ], "Keep the greeting brief."
 
     requests: list[dict[str, Any]] = []
 
@@ -43,13 +53,15 @@ async def acompletion(self, **kwargs: Any) -> ChatCompletion:
             requests.append(kwargs)
             message: dict[str, Any] = {"role": "assistant", "content": "Hello Ada"}
             if len(requests) == 1:
+                name = "run_code" if code_mode else "provider_lookup"
+                arguments = {"code": "print(await tools.provider_lookup(name='Ada'))"} if code_mode else {"name": "Ada"}
                 message = {
                     "role": "assistant",
                     "tool_calls": [
                         {
                             "id": "lookup",
                             "type": "function",
-                            "function": {"name": "provider_lookup", "arguments": json.dumps({"name": "Ada"})},
+                            "function": {"name": name, "arguments": json.dumps(arguments)},
                         }
                     ],
                 }
@@ -68,24 +80,35 @@ async def acompletion(self, **kwargs: Any) -> ChatCompletion:
             })
 
     monkeypatch.setattr("bub.builtin.model_runner.AnyLLM.create", lambda *args, **kwargs: Provider())
-    agent = Agent(framework, tools=[direct, denied, supplied], skill_dirs=[])
-    agent.tool_providers.append(provide)
+    agent = Agent(framework, tools=[direct, run_code], skill_dirs=[])
+    agent.tool_sources["earlier"] = {supplied.name: Tool.from_callable(lambda: "wrong", name=supplied.name)}
+    agent.tool_sources["later"] = {supplied.name: supplied, denied.name: denied, pending.name: pending}
+    agent.tool_providers = [guide, select] if reverse_providers else [select, guide]
     stream = await agent.run_stream(
         session_id="provider",
         prompt="Greet Ada.",
         model="openrouter:test-model",
-        allowed_tools=["direct", "provider_lookup"],
+        allowed_tools=["direct", "provider_lookup", "provider_pending", "run_code"],
+        state={"code_mode": code_mode},
     )
     events = [event async for event in stream]
     assert any(event.data.get("text") == "Hello Ada" for event in events if event.kind == "final")
     assert calls == ["Ada"]
+    # Request preparation does not register discovered tools on the agent.
+    assert agent.tools.keys() == {"direct", "run_code"}
     definitions = {item["function"]["name"]: item["function"] for item in requests[0]["tools"]}
-    assert definitions.keys() == {"direct", "provider_lookup"}
-    assert any(
-        "Use provider_lookup" in message["content"]
-        for message in requests[0]["messages"]
-        if message["role"] == "system"
+    assert definitions.keys() == ({"run_code"} if code_mode else {"direct", "provider_lookup"})
+    if code_mode:
+        stub = next((tmp_path / "codemode").rglob("*.pyi")).read_text()
+        assert "provider_pending" not in stub
+    system = "\n".join(message["content"] for message in requests[0]["messages"] if message["role"] == "system")
+    guidance = (
+        ["Keep the greeting brief.", "Use provider_lookup"]
+        if reverse_providers
+        else ["Use provider_lookup", "Keep the greeting brief."]
     )
+    assert system.index(guidance[0]) < system.index(guidance[1])
+    expected = "Hello Ada" if code_mode else "HELLO ADA"
     assert any(
-        message.get("content") == "Hello Ada" for message in requests[1]["messages"] if message["role"] == "tool"
+        expected in message.get("content", "") for message in requests[1]["messages"] if message["role"] == "tool"
     )

website/src/content/docs/docs/build/tools.mdx

@@ -64,11 +64,13 @@ If `run_code` is not allowed (for example, a subagent restricted with `allowed_t
 
 `run_code` hands the code to the session [environment](#run-tools-in-an-environment)'s `run_code`, together with a `call_tool` callback. Tool calls always run on the host through that callback, so hooks still see every call. `timeout_seconds` cancels the environment call. The builtin `LocalEnvironment` starts a fresh Python process for every call (`sys.executable`, working directory is the workspace) and forwards tool calls as JSON lines over its stdin and stdout, so arguments and results must be JSON-serializable. When the code finishes, fails or times out, it kills the process and everything it started. That process runs on the host with the same permissions as `bash`, so enable code mode only where `bash` would be acceptable. The return annotation of a tool function (or `Tool.output_schema`, for tools built by hand) determines the result type shown in the stub.
 
-### Per-request tool presentation
+### Tool discovery and request preparation
 
-`Agent.tool_providers` accepts async callbacks `(tools, tape) -> (tools, tool_prompt)`. Bub applies the current tool scope and prepares code mode first, then runs these callbacks in registration order. They receive runtime names; model aliases are applied afterward. The returned tools and prompt fragment are used for that request.
+Register directly available tools in `Agent.tools`. For a separate discovery inventory, set `agent.tool_sources[source]` to a mapping of runtime names to `Tool` instances. `source` is a plugin-owned key used to update or remove that inventory. `Agent.known_tools` combines direct tools and source inventories for commands and `allowed_tools`; later registered sources win duplicate names, independently of provider order. Updating an existing source preserves its priority.
 
-Providers prepare already registered tools; they do not register tools or change connection lifecycles. A plugin binds its discovery helper alongside its tools and keeps discovery within the supplied scope. Full native definitions are sent through the tool schema instead of repeated in the system prompt.
+Append request-time preparation callables to `agent.tool_providers`. The `ToolProvider` contract in `bub.tools` is `async (tools, tape) -> (tools, prompt_fragment)`. Providers receive the filtered toolset, run in their supplied order, and pass their result to the next provider. They may select definitions and contribute discovery guidance, but must not reintroduce tools excluded by the scope. Builtin preparation runs last, so code mode uses the final selected toolset.
+
+Request preparation does not change `Agent.tools`. Known tools that pass the scope reach the providers, so a provider that defers discovered tools must remove them from the request itself. The tools left after the providers form the request toolset: only those calls run, and calls to other names receive a guidance result. Plugins own discovery state, source removal and connection cleanup. Full native definitions are sent through tool schemas instead of repeated in system. Keep discovery summaries stable and append native definitions in the conversation's first-discovery order; request serialization determines how much of the prefix remains cacheable.
 
 ## Run tools in an environment
 
@@ -91,7 +93,7 @@ The `REGISTRY` lives in [`bub.tools`](https://github.com/bubbuild/bub/blob/main/
 REGISTRY: dict[str, Tool] = {}
 ```
 
-Every `@tool` call mutates this dict at **import time**. Bub's builtin agent reads from `REGISTRY` when assembling the tool list for the model. There is no separate registration step.
+Every `@tool` call mutates this dict at **import time**. `Agent(tools=None)` snapshots it when created; later imports do not update existing agents. Add tools to an instance's `Agent.tools`, or register a discovery inventory in `agent.tool_sources` and a preparation callable in `agent.tool_providers`.
 
 ## Import the tools module from your plugin