feat: load deferred tools on demand via tool.describe

src/bub/builtin/agent.py

@@ -7,7 +7,7 @@
 import re
 import shlex
 import time
-from collections.abc import AsyncGenerator, AsyncIterator, Collection, Iterable, Mapping
+from collections.abc import AsyncGenerator, AsyncIterator, Collection, Iterable
 from contextlib import AsyncExitStack, aclosing
 from dataclasses import dataclass, replace
 from datetime import UTC, datetime
@@ -29,15 +29,13 @@
 from bub.store import AsyncTapeStore, AsyncTapeStoreAdapter, InMemoryTapeStore, TapeStore, is_async_tape_store
 from bub.streaming import AsyncStreamEvents, StreamEvent, StreamState
 from bub.tape import Tape
-from bub.tools import REGISTRY, Tool, ToolContext, ToolProvider, model_tools
+from bub.tools import REGISTRY, Tool, ToolContext, model_tools
 from bub.tracing import Span, current_span
 from bub.turn import TurnState
 from bub.utils import workspace_from_state
 
 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:
@@ -75,22 +73,10 @@ 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.
@@ -260,14 +246,13 @@ async def _run_command(self, tape: Tape, *, line: str) -> str:
         output = ""
         status = "ok"
         try:
-            known_tools = self.known_tools
-            if name not in known_tools:
-                if "bash" not in known_tools:
+            if name not in self.tools:
+                if "bash" not in self.tools:
                     raise ValueError("bash tool is not available")  # noqa: TRY301
-                bash_tool = known_tools["bash"]
+                bash_tool = self.tools["bash"]
                 output = bash_tool.render(await bash_tool.run(context=context, command=line))
             else:
-                command_tool = known_tools[name]
+                command_tool = self.tools[name]
                 args = _parse_args(arg_tokens)
                 if command_tool.context:
                     args.kwargs["context"] = context
@@ -459,27 +444,50 @@ 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=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 known_tools.values() if tool.name in allowed_tools]
-        else:
-            tools = list(known_tools.values())
         return await self._run_once_stream(
             tape=tape,
             prompt=prompt,
             prompt_text=prompt_text,
             model=model,
             allowed_skills=allowed_skills,
-            tools=tools,
+            tools=self._allowed_tools(tape, allowed_tools),
+        )
+
+    def _allowed_tools(self, tape: Tape, allowed_tools: Collection[str] | None) -> list[Tool]:
+        """Filter direct and deferred tools alike by ``allowed_tools`` and record the scope in state."""
+        from bub.builtin.tools import ALLOWED_TOOLS_STATE_KEY, resolve_tool_names
+
+        state = tape.context.state
+        if allowed_tools is None:
+            state.pop(ALLOWED_TOOLS_STATE_KEY, None)
+            return list(self.tools.values())
+        allowed = resolve_tool_names(allowed_tools, all_names=self.tools)
+        state[ALLOWED_TOOLS_STATE_KEY] = sorted(allowed)
+        return [tool for name, tool in self.tools.items() if name in allowed]
+
+    async def _prepare_deferred_tools(self, tools: list[Tool], tape: Tape) -> tuple[list[Tool], str]:
+        """Return direct tools followed by deferred tools loaded on the tape, plus the deferred tools prompt.
+
+        ``tool.describe`` is exposed exactly when a deferred tool remains, since it is the only way to load one.
+        """
+        from bub.builtin.tools import (
+            TOOL_DESCRIBE_TOOL_NAME,
+            loaded_tool_names,
+            render_deferred_tools_prompt,
+            tool_describe,
         )
 
+        direct = [tool for tool in tools if not tool.deferred and tool.name != TOOL_DESCRIBE_TOOL_NAME]
+        deferred = {tool.name: tool for tool in tools if tool.deferred and tool.agent_use}
+        if not deferred:
+            return direct, ""
+        direct.append(self.tools.get(TOOL_DESCRIBE_TOOL_NAME, tool_describe))
+        loaded = [deferred[name] for name in await loaded_tool_names(tape) if name in deferred]
+        return [*direct, *loaded], render_deferred_tools_prompt(deferred.values())
+
     async def _run_once_stream(
         self,
         *,
@@ -490,20 +498,14 @@ async def _run_once_stream(
         allowed_skills: set[str] | None,
         tools: list[Tool],
     ) -> AsyncStreamEvents:
-        tools_prompts: list[str] = []
-        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)
+        # Code mode sees every allowed tool, deferred ones included; deferral applies to model-facing tools.
+        tools, code_mode_prompt = await self._prepare_code_mode(tools, tape)
+        tools, deferred_prompt = await self._prepare_deferred_tools(tools, tape)
         system_prompt = self._system_prompt(
             prompt_text,
             state=tape.context.state,
             allowed_skills=allowed_skills,
-            tools_prompt="\n\n".join(tools_prompts),
+            tools_prompt="\n\n".join(block for block in (code_mode_prompt, deferred_prompt) if block),
         )
         resolved_model = model or self.settings.model
 

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 REQUEST_TOOLS_STATE_KEY, Agent
+from bub.builtin.agent import 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
@@ -356,16 +356,29 @@ async def before_tool_call(
         replace it with a guidance ``tool_result`` so the model can re-issue a
         valid call on the next step.
         """
+        from bub.builtin.codemode import CODE_TOOLS_STATE_KEY
+        from bub.builtin.tools import TOOL_DESCRIBE_TOOL_NAME
         from bub.tools import model_tools
 
-        if (request_tools := state.get(REQUEST_TOOLS_STATE_KEY)) is not None:
-            available_tools = tuple(request_tools)
+        agent = self._get_agent(state)
+
+        agent_tools = model_tools(agent.tools.values())
+        if "_runtime_tool_names" in state:
+            code_tools = state.get(CODE_TOOLS_STATE_KEY) or ()
+            available_tools = (*state["_runtime_tool_names"], *(tool_item.name for tool_item in code_tools))
         else:
-            agent = self._get_agent(state)
-            available_tools = tuple(tool_item.name for tool_item in model_tools(agent.known_tools.values()))
+            available_tools = tuple(tool_item.name for tool_item in agent_tools)
         if call.tool in available_tools:
             return None
 
+        describe_name = TOOL_DESCRIBE_TOOL_NAME.replace(".", "_")
+        if describe_name in available_tools and any(
+            tool_item.deferred and tool_item.name == call.tool for tool_item in agent_tools
+        ):
+            return ToolCallDecision.replace(
+                f"Tool `{call.tool}` is not loaded. Call `{describe_name}` with its name first."
+            )
+
         matches = get_close_matches(call.tool, available_tools, n=3, cutoff=0.6)
         if matches:
             suggestions = "\n".join(f"- {name}" for name in matches)

src/bub/builtin/tools.py

@@ -4,7 +4,7 @@
 import contextlib
 import json
 import uuid
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping
 from contextlib import aclosing
 from dataclasses import asdict
 from typing import TYPE_CHECKING, Any, TypedDict, cast, final
@@ -19,10 +19,14 @@
 
 if TYPE_CHECKING:
     from bub.builtin.agent import Agent
+    from bub.tape import Tape
 
 DEFAULT_COMMAND_TIMEOUT_SECONDS = 30
 DEFAULT_HEADERS = {"accept": "text/markdown"}
 DEFAULT_REQUEST_TIMEOUT_SECONDS = 10
+TOOL_DESCRIBE_TOOL_NAME = "tool.describe"
+TOOL_LOADED_EVENT = "tool.loaded"
+ALLOWED_TOOLS_STATE_KEY = "_runtime_allowed_tools"
 
 
 def _to_model_name(name: str) -> str:
@@ -104,6 +108,31 @@ def render_tools_prompt(tools: Iterable[Tool]) -> str:
     return f"<available_tools>\n{'\n'.join(lines)}\n</available_tools>"
 
 
+def render_deferred_tools_prompt(tools: Iterable[Tool]) -> str:
+    """List deferred tools by model-facing name so the model knows what it can load."""
+    names = "\n".join(f"- {_to_model_name(tool_item.name)}" for tool_item in tools)
+    if not names:
+        return ""
+    return (
+        "<deferred_tools>\n"
+        f"These tools are not loaded yet. Call {_to_model_name(TOOL_DESCRIBE_TOOL_NAME)} with their names to load "
+        "their definitions; loaded tools become callable from the next step.\n"
+        f"{names}\n</deferred_tools>"
+    )
+
+
+async def loaded_tool_names(tape: Tape) -> list[str]:
+    """Return runtime names of tools loaded on this tape, in load order."""
+    names: dict[str, None] = {}
+    for entry in await tape.search(tape.query().kinds("event")):
+        if entry.payload.get("name") != TOOL_LOADED_EVENT:
+            continue
+        data = entry.payload.get("data")
+        if isinstance(data, Mapping):
+            names.update(dict.fromkeys(str(name) for name in data.get("names") or ()))
+    return list(names)
+
+
 def _raise_for_failed_shell(returncode: int | None, output: str) -> None:
     if returncode in (None, 0):
         return
@@ -167,6 +196,17 @@ class ToolFailure(TypedDict):
     error: str
 
 
+class ToolDefinition(TypedDict):
+    name: str
+    description: str
+    parameters: dict[str, Any]
+
+
+class ToolDescriptions(TypedDict):
+    tools: list[ToolDefinition]
+    unknown: list[str]
+
+
 class TapeInfoResult(TypedDict):
     name: str
     entries: int
@@ -358,6 +398,40 @@ def skill_describe(name: str | None = None, *, context: ToolContext) -> SkillLis
     return {"name": skill.name, "location": str(skill.location), "content": skill.body() or ""}
 
 
+@tool(context=True, name=TOOL_DESCRIBE_TOOL_NAME, preserve=True)
+async def tool_describe(names: list[str], *, context: ToolContext) -> ToolDescriptions:
+    """Load tools by name and return their definitions. Deferred tools become callable from the next step."""
+    agent = _get_agent(context)
+    allowed_tools = context.state.get(ALLOWED_TOOLS_STATE_KEY)
+    available = {
+        name: tool_item
+        for name, tool_item in agent.tools.items()
+        if tool_item.agent_use and (allowed_tools is None or name in allowed_tools)
+    }
+    index = _tool_name_index(available)
+    loaded = set(await loaded_tool_names(context.tape))
+    definitions: list[ToolDefinition] = []
+    unknown: list[str] = []
+    newly_loaded: list[str] = []
+    for name in names:
+        resolved = index.get(name.strip().casefold())
+        if resolved is None:
+            unknown.append(name)
+            continue
+        tool_item = available[resolved]
+        definitions.append({
+            "name": _to_model_name(resolved),
+            "description": tool_item.description,
+            "parameters": tool_item.parameters,
+        })
+        if tool_item.deferred and resolved not in loaded:
+            loaded.add(resolved)
+            newly_loaded.append(resolved)
+    if newly_loaded:
+        await context.tape.append_event(TOOL_LOADED_EVENT, {"names": newly_loaded})
+    return {"tools": definitions, "unknown": unknown}
+
+
 @tool(context=True, name="tape.info", renderer=_render_tape_info)
 async def tape_info(context: ToolContext) -> TapeInfoResult:
     """Get information about the current tape, such as number of entries and anchors."""
@@ -434,7 +508,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.known_tools)
+    allowed_tools = resolve_tool_names(param.allowed_tools or None, exclude={"subagent"}, all_names=agent.tools)
     output = ""
     errors: list[str] = []
     stream = await agent.run_stream(

src/bub/tools.py

@@ -149,6 +149,8 @@ class Tool:
     """Keep the tool directly callable by the model in code mode; others are reachable only from code."""
     output_schema: dict[str, Any] | None = None
     """JSON schema of the structured result, used to describe the tool to model-written code."""
+    deferred: bool = False
+    """Load the tool on demand: only its name is listed until ``tool.describe`` loads its definition."""
 
     @property
     def code_use(self) -> bool:
@@ -189,6 +191,7 @@ def from_callable(
         agent_use: bool = True,
         renderer: Callable[[Any], str] | None = None,
         preserve: bool = False,
+        deferred: bool = False,
     ) -> Tool:
         signature = inspect.signature(func)
         if context and "context" not in signature.parameters:
@@ -216,13 +219,10 @@ def validated(*args: Any, **kwargs: Any) -> Any:
             renderer=renderer,
             preserve=preserve,
             output_schema=_output_schema(func),
+            deferred=deferred,
         )
 
 
-type ToolProvider = Callable[[list[Tool], Tape], Awaitable[tuple[list[Tool], str]]]
-"""Prepare registered tools and a prompt fragment for one model request."""
-
-
 def model_tools(tools: Iterable[Tool]) -> list[Tool]:
     """Convert agent-enabled runtime tools into model-safe aliases."""
     return [replace(tool_item, name=tool_item.name.replace(".", "_")) for tool_item in tools if tool_item.agent_use]
@@ -562,6 +562,7 @@ def tool(
     agent_use: bool = ...,
     renderer: Callable[[Any], str] | None = ...,
     preserve: bool = ...,
+    deferred: bool = ...,
 ) -> Tool: ...
 
 
@@ -576,6 +577,7 @@ def tool(
     agent_use: bool = ...,
     renderer: Callable[[Any], str] | None = ...,
     preserve: bool = ...,
+    deferred: bool = ...,
 ) -> Callable[[Callable], Tool]: ...
 
 
@@ -589,11 +591,13 @@ def tool(
     agent_use: bool = True,
     renderer: Callable[[Any], str] | None = None,
     preserve: bool = False,
+    deferred: bool = False,
 ) -> Tool | Callable[[Callable], Tool]:
     """Decorator to convert a function into a Tool instance.
 
     Tools should return structured results; ``renderer`` turns such a result into the
     plain text shown to the model outside code mode (defaults to JSON for non-strings).
+    ``deferred`` tools are loaded on demand through ``tool.describe``.
     """
 
     def decorator(func: Callable) -> Tool:
@@ -618,6 +622,7 @@ def handler(*args: Any, **kwargs: Any) -> Any:
                 renderer=renderer,
                 preserve=preserve,
                 output_schema=_output_schema(func),
+                deferred=deferred,
             )
         else:
             result = Tool.from_callable(
@@ -628,6 +633,7 @@ def handler(*args: Any, **kwargs: Any) -> Any:
                 agent_use=agent_use,
                 renderer=renderer,
                 preserve=preserve,
+                deferred=deferred,
             )
         tool_instance = _add_logging(result)
         REGISTRY[tool_instance.name] = tool_instance