diff --git a/src/bub/builtin/agent.py b/src/bub/builtin/agent.py index 1142ff05..2249f93a 100644 --- a/src/bub/builtin/agent.py +++ b/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, diff --git a/src/bub/builtin/hook_impl.py b/src/bub/builtin/hook_impl.py index 70f05ab3..2ec79708 100644 --- a/src/bub/builtin/hook_impl.py +++ b/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 diff --git a/src/bub/builtin/tools.py b/src/bub/builtin/tools.py index 12bfd7f3..ff6740c2 100644 --- a/src/bub/builtin/tools.py +++ b/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( diff --git a/tests/test_builtin_agent.py b/tests/test_builtin_agent.py index bedcec6a..836aa86c 100644 --- a/tests/test_builtin_agent.py +++ b/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 diff --git a/tests/test_builtin_hook_impl.py b/tests/test_builtin_hook_impl.py index 7747a3f1..031130ef 100644 --- a/tests/test_builtin_hook_impl.py +++ b/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 diff --git a/tests/test_subagent_tool.py b/tests/test_subagent_tool.py index 1826de24..d77f5930 100644 --- a/tests/test_subagent_tool.py +++ b/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: diff --git a/tests/test_tool_providers.py b/tests/test_tool_providers.py index 18329053..a5e5ea4b 100644 --- a/tests/test_tool_providers.py +++ b/tests/test_tool_providers.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +from dataclasses import replace from pathlib import Path from typing import Any @@ -8,14 +9,17 @@ 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") @@ -23,6 +27,7 @@ async def test_provider_prompt_reaches_the_model_and_registered_tools_remain_cal 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" ) diff --git a/website/src/content/docs/docs/build/tools.mdx b/website/src/content/docs/docs/build/tools.mdx index de281adf..1b73791c 100644 --- a/website/src/content/docs/docs/build/tools.mdx +++ b/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 diff --git a/website/src/content/docs/zh-cn/docs/build/tools.mdx b/website/src/content/docs/zh-cn/docs/build/tools.mdx index a819fb5b..6475f87f 100644 --- a/website/src/content/docs/zh-cn/docs/build/tools.mdx +++ b/website/src/content/docs/zh-cn/docs/build/tools.mdx @@ -65,12 +65,13 @@ Code mode 让模型在 Python 中调用工具，而不必一次只发起一个 `run_code` 会把代码和一个 `call_tool` 回调一起交给当前 session [执行环境](#在执行环境中运行工具)的 `run_code` 执行。工具调用始终经这个回调在宿主上执行，所以 hook 仍能看到每一次调用。超过 `timeout_seconds` 时会取消执行环境里的执行。builtin 的 `LocalEnvironment` 每次调用都会启动一个新的 Python 进程（使用 `sys.executable`，工作目录为 workspace），工具调用以 JSON 行的形式经进程的 stdin/stdout 转发，因此参数和结果都必须能 JSON 序列化。代码执行完、出错或超时后，它会杀掉该进程及其启动的所有子进程。这个进程直接跑在宿主上，权限与 `bash` 相同，因此只应在允许使用 `bash` 的场景下开启 code mode。stub 中的结果类型来自工具函数的返回注解（手动构造的工具则来自 `Tool.output_schema`）。 -### 每次请求的工具呈现 +### 工具发现与请求准备 -`Agent.tool_providers` 接受异步回调 `(tools, tape) -> (tools, tool_prompt)`。Bub 先应用当前工具范围、准备 code mode，再按注册顺序执行这些回调。输入使用运行时名称，之后才转换模型别名；返回的工具及提示片段用于本次请求。 +直接可用的工具注册在 `Agent.tools`。独立的发现清单通过 `agent.tool_sources[source]` 登记，值为运行时名称到 `Tool` 实例的映射。`source` 是插件管理的键，用于更新或移除自己的清单。`Agent.known_tools` 合并直接工具和各来源清单，供命令查找及 `allowed_tools` 使用；后注册的来源优先处理重名，与 provider 顺序无关。更新已有来源不改变优先级。 -provider 准备已注册的工具，工具注册与连接生命周期仍由插件管理。插件绑定自己的发现 helper，并将查询限制在传入范围内。完整原生定义通过工具 schema 发送，系统提示词不再重复列出。 +向 `agent.tool_providers` 追加请求准备函数。`bub.tools` 中的 `ToolProvider` 契约是 `async (tools, tape) -> (tools, prompt_fragment)`。Provider 接收过滤后的工具集合，按提供的顺序执行，并将结果交给下一项。它们可以选择完整定义并贡献发现提示，但不得重新加入被 scope 排除的工具。Builtin preparation 最后执行，因此 code mode 使用最终选中的工具集合。 +请求准备不会修改 `Agent.tools`。通过 scope 的已知工具都会交给 provider，因此需要延迟暴露的发现工具必须由 provider 自己从本次请求中移除。provider 处理后剩下的工具构成本次请求的工具集合：只有这些调用会执行，调用其他名称会收到引导结果。发现状态、来源移除和连接清理由插件管理。完整原生定义通过工具 schema 发送，系统提示词不再重复列出。保持发现摘要稳定，并按当前会话的首次发现顺序追加原生定义；此前前缀有多少可以命中缓存，取决于请求序列化方式。 ## 在执行环境中运行工具 @@ -93,7 +94,7 @@ provider 准备已注册的工具，工具注册与连接生命周期仍由插 REGISTRY: dict[str, Tool] = {} ``` -每一次 `@tool` 调用都会在**导入时**修改这个字典。Bub 内置代理在为模型组装工具列表时从 `REGISTRY` 读取。没有独立的注册步骤。 +每一次 `@tool` 调用都会在**导入时**修改这个字典。`Agent(tools=None)` 在创建时取得其快照，之后导入的工具不会自动更新已有 Agent。可直接向实例的 `Agent.tools` 添加工具，或分别通过 `agent.tool_sources` 登记发现清单、通过 `agent.tool_providers` 登记请求准备函数。 ## 在插件里导入工具模块