diff --git a/src/bub/builtin/agent.py b/src/bub/builtin/agent.py index da92c822..1fa78d1f 100644 --- a/src/bub/builtin/agent.py +++ b/src/bub/builtin/agent.py @@ -37,7 +37,6 @@ from bub.turn import TurnState from bub.utils import workspace_from_state -CONTINUE_PROMPT = "Continue the task until all targets are completed." HINT_RE = re.compile(r"\$([A-Za-z0-9_.-]+)") MAX_AUTO_HANDOFF_RETRIES = 1 @@ -291,7 +290,7 @@ async def _stream_events_with_auto_handoff( ) return - next_prompt = self._continue_prompt(tape) + next_prompt = await self.framework.continue_prompt(tape=tape, state=state) await tape.append_event( "loop.step", { @@ -398,11 +397,6 @@ def _system_prompt( blocks.append(skills_prompt) return "\n\n".join(blocks) - def _continue_prompt(self, tape: Tape) -> str: - if "context" in tape.context.state: - return f"{CONTINUE_PROMPT} [context: {tape.context.state['context']}]" - return CONTINUE_PROMPT - def _has_steering_messages(self, state: TurnState) -> bool: steering_inbox = self.framework.get_steering_inbox() return bool(steering_inbox and steering_inbox.message_count(state) > 0) diff --git a/src/bub/builtin/hook_impl.py b/src/bub/builtin/hook_impl.py index 4df649dd..bd6fc4f9 100644 --- a/src/bub/builtin/hook_impl.py +++ b/src/bub/builtin/hook_impl.py @@ -22,8 +22,8 @@ from bub.hooks.interception import ToolCall, ToolCallDecision from bub.model_selection import ModelChoice, ModelOptions from bub.store import TapeStore -from bub.streaming import AsyncStreamEvents -from bub.tape import TapeContext +from bub.streaming import AsyncStreamEvents, StreamState +from bub.tape import Tape, TapeContext from bub.turn import TurnState AGENTS_FILE_NAME = "AGENTS.md" @@ -60,6 +60,7 @@ Excessively long context may cause model call failures. In this case, you MAY use tape.info to retrieve the token usage and you SHOULD use tape.handoff tool to shorten the retrieved history. """ +DEFAULT_CONTINUE_PROMPT = "Continue the task until all targets are completed." class BuiltinImpl: @@ -217,6 +218,13 @@ async def run_model_stream(self, prompt: str | list[dict], session_id: str, stat model=state.get("model"), ) + @hookimpl + def continue_prompt(self, tape: Tape, state: StreamState) -> str: + del state + if "context" in tape.context.state: + return f"{DEFAULT_CONTINUE_PROMPT} [context: {tape.context.state['context']}]" + return DEFAULT_CONTINUE_PROMPT + @hookimpl def register_cli_commands(self, app: typer.Typer) -> None: from bub.builtin import cli diff --git a/src/bub/framework.py b/src/bub/framework.py index b11f1cb9..5d94660c 100644 --- a/src/bub/framework.py +++ b/src/bub/framework.py @@ -23,7 +23,8 @@ from bub.hooks.specs import BUB_HOOK_NAMESPACE, BubHookSpecs from bub.model_selection import ModelOptions from bub.store import AsyncTapeStore, TapeStore -from bub.tape import TapeContext +from bub.streaming import StreamState +from bub.tape import Tape, TapeContext from bub.turn import TurnResult, TurnState from bub.utils import maybe_context_manager @@ -123,6 +124,13 @@ async def build_prompt( prompt = content_of(message) return cast("str | list[dict[str, Any]]", prompt) + async def continue_prompt(self, tape: Tape, state: StreamState) -> str: + """Build the prompt for the next step of an agent loop.""" + prompt = await self._hook_runtime.call_first("continue_prompt", tape=tape, state=state) + if isinstance(prompt, str): + return prompt + raise TypeError("hook.continue_prompt must return str") + async def build_state(self, message: Envelope, session_id: str) -> TurnState: state = {"_runtime_workspace": str(self.workspace), "_runtime_steering_inbox": self.get_steering_inbox()} for hook_state in reversed( diff --git a/src/bub/hooks/specs.py b/src/bub/hooks/specs.py index 6ddf93f3..e85e7e67 100644 --- a/src/bub/hooks/specs.py +++ b/src/bub/hooks/specs.py @@ -20,8 +20,8 @@ ) from bub.model_selection import ModelOptions from bub.store import AsyncTapeStore, TapeStore -from bub.streaming import AsyncStreamEvents -from bub.tape import TapeContext +from bub.streaming import AsyncStreamEvents, StreamState +from bub.tape import Tape, TapeContext from bub.turn import TurnState if TYPE_CHECKING: @@ -65,6 +65,15 @@ def run_model_stream(self, prompt: str | list[dict], session_id: str, state: Tur """ raise NotImplementedError + @hookspec(firstresult=True) + def continue_prompt(self, tape: Tape, state: StreamState) -> str: + """Build the prompt used to continue an agent loop. + + Implementations may be synchronous or asynchronous. The first + non-``None`` result in hook priority order is used. + """ + raise NotImplementedError + @hookspec def load_state(self, message: Envelope, session_id: str) -> TurnState: """Load state snapshot for one session.""" diff --git a/tests/test_builtin_agent.py b/tests/test_builtin_agent.py index 7d8a01e0..f08e1f50 100644 --- a/tests/test_builtin_agent.py +++ b/tests/test_builtin_agent.py @@ -14,6 +14,7 @@ from bub.builtin.settings import AgentSettings from bub.builtin.steering import InMemorySteeringInbox from bub.errors import BubError +from bub.streaming import AsyncStreamEvents, StreamEvent, StreamState from bub.tape import TapeContext from bub.tools import REGISTRY, tool @@ -251,6 +252,43 @@ async def test_agent_run_model_defaults_to_none() -> None: assert completion_kwargs["model"] == "test:model" +@pytest.mark.asyncio +async def test_agent_loop_awaits_continue_prompt_hook_with_stream_state() -> None: + agent = _make_agent() + tape = _FakeTape(_ForkCapture()) + prompts: list[str | list[dict]] = [] + observed_usage: list[dict[str, Any] | None] = [] + + async def run_once(**kwargs: Any) -> AsyncStreamEvents: + prompts.append(kwargs["prompt"]) + should_continue = len(prompts) == 1 + + async def iterator() -> AsyncIterator[StreamEvent]: + yield StreamEvent("final", {"tool_calls": ["call"] if should_continue else []}) + + return AsyncStreamEvents(iterator(), state=StreamState(usage={"step": len(prompts)})) + + async def continue_prompt(*, tape: _FakeTape, state: StreamState) -> str: + observed_usage.append(state.usage) + return "custom continuation" + + agent._run_once = run_once # type: ignore[method-assign] + agent.framework.continue_prompt = continue_prompt + + events = [ + event + async for event in agent._stream_events_with_auto_handoff( + tape=tape, # type: ignore[arg-type] + prompt="initial prompt", + state=StreamState(), + ) + ] + + assert [event.kind for event in events] == ["final", "final"] + assert prompts == ["initial prompt", "custom continuation"] + assert observed_usage == [{"step": 1}] + + @pytest.mark.asyncio async def test_agent_run_model_override_does_not_mutate_default() -> None: """A per-call model override must not leak into the agent's configured model. diff --git a/tests/test_builtin_hook_impl.py b/tests/test_builtin_hook_impl.py index 2d8d95f4..f026ee5c 100644 --- a/tests/test_builtin_hook_impl.py +++ b/tests/test_builtin_hook_impl.py @@ -7,11 +7,11 @@ import pytest -from bub.builtin.hook_impl import AGENTS_FILE_NAME, DEFAULT_SYSTEM_PROMPT, BuiltinImpl +from bub.builtin.hook_impl import AGENTS_FILE_NAME, DEFAULT_CONTINUE_PROMPT, DEFAULT_SYSTEM_PROMPT, BuiltinImpl from bub.channels.message import ChannelMessage from bub.framework import BubFramework from bub.store import AsyncTapeStoreAdapter, FileTapeStore, InMemoryTapeStore -from bub.streaming import AsyncStreamEvents, StreamEvent +from bub.streaming import AsyncStreamEvents, StreamEvent, StreamState from bub.tape import Tape, TapeContext @@ -92,6 +92,15 @@ def test_resolve_session_falls_back_to_channel_and_chat_id(tmp_path: Path) -> No assert impl.resolve_session(message) == "telegram:42" +def test_continue_prompt_includes_tape_context(tmp_path: Path) -> None: + _, impl, _ = _build_impl(tmp_path) + tape = _fake_tape(tmp_path).with_context(TapeContext(state={"context": "telegram metadata"})) + + prompt = impl.continue_prompt(tape=tape, state=StreamState()) + + assert prompt == f"{DEFAULT_CONTINUE_PROMPT} [context: telegram metadata]" + + @pytest.mark.asyncio async def test_load_state_and_save_state_manage_lifespan_and_context(tmp_path: Path) -> None: _, impl, agent = _build_impl(tmp_path) diff --git a/tests/test_framework.py b/tests/test_framework.py index 44ca77c3..2bb00601 100644 --- a/tests/test_framework.py +++ b/tests/test_framework.py @@ -118,6 +118,35 @@ def system_prompt(self, prompt: str, state: dict[str, str]) -> str | None: assert prompt == "low\n\nhigh" +@pytest.mark.asyncio +async def test_continue_prompt_awaits_high_priority_async_hook() -> None: + framework = BubFramework() + tape = cast(Any, SimpleNamespace(context=SimpleNamespace(state={}))) + state = StreamState(usage={"total_tokens": 42}) + called: list[str] = [] + + class SyncPlugin: + @hookimpl + def continue_prompt(self, tape, state): + called.append("sync") + return "sync prompt" + + class AsyncPlugin: + @hookimpl + async def continue_prompt(self, tape: Any, state: StreamState) -> str: + called.append("async") + assert state.usage == {"total_tokens": 42} + return "async prompt" + + framework._plugin_manager.register(SyncPlugin(), name="sync") + framework._plugin_manager.register(AsyncPlugin(), name="async") + + prompt = await framework.continue_prompt(tape=tape, state=state) + + assert prompt == "async prompt" + assert called == ["async"] + + @pytest.mark.asyncio async def test_running_enters_tape_store_once_and_reuses_it() -> None: framework = BubFramework() diff --git a/website/src/content/docs/docs/reference/hooks.mdx b/website/src/content/docs/docs/reference/hooks.mdx index b083cdd1..781398a2 100644 --- a/website/src/content/docs/docs/reference/hooks.mdx +++ b/website/src/content/docs/docs/reference/hooks.mdx @@ -20,6 +20,7 @@ For the *why* and *how* of each stage see [Turn pipeline](/docs/concepts/turn-pi | `build_prompt` | firstresult | `(message: Envelope, session_id: str, state: TurnState) -> str \| list[dict]` | prompt or content parts | `BubFramework.process_inbound` | Falls back to `content_of(message)` when no impl returns a non-`None` value, or when the selected firstresult is falsy. A falsy selected value does not cause lower-priority impls to run. | | `run_model` | firstresult | `(prompt, session_id, state) -> str` | model text | `HookRuntime.run_model` | Legacy text path. Implement either `run_model` or `run_model_stream`, not both. | | `run_model_stream` | firstresult | `(prompt, session_id, state) -> AsyncStreamEvents` | async stream | `HookRuntime.run_model_stream` | Preferred. Falls back to wrapping a `run_model` result in a one-chunk stream when no streaming impl exists. | +| `continue_prompt` | firstresult | `(tape: Tape, state: StreamState) -> str` | next-step prompt | `Agent._stream_events_with_auto_handoff` | Runs before each continued agent-loop step. Implementations may be sync or async; the builtin preserves the default prompt and optional tape context. | | `save_state` | broadcast | `(session_id, state, message, model_output) -> None` | none | `process_inbound` model-stage finally block | Runs after prompt resolution for model-stage success or failure; failures before prompt/model execution skip it. | | `render_outbound` | broadcast | `(message, session_id, state, model_output) -> list[Envelope]` | outbound batch | `BubFramework._collect_outbounds` | All batches are concatenated via `unpack_batch`. Empty results trigger a default echo envelope. | | `dispatch_outbound` | broadcast | `(message: Envelope) -> bool` | sent flag | `process_inbound` per outbound | Each outbound is fanned out to every impl. | diff --git a/website/src/content/docs/zh-cn/docs/reference/hooks.mdx b/website/src/content/docs/zh-cn/docs/reference/hooks.mdx index 1cd6a54f..52e84ba3 100644 --- a/website/src/content/docs/zh-cn/docs/reference/hooks.mdx +++ b/website/src/content/docs/zh-cn/docs/reference/hooks.mdx @@ -20,6 +20,7 @@ description: BubHookSpecs 中每个钩子的类型、签名、返回值与调用 | `build_prompt` | firstresult | `(message: Envelope, session_id: str, state: TurnState) -> str \| list[dict]` | prompt or content parts | `BubFramework.process_inbound` | 没有实现返回非 `None` 值,或选中的 firstresult 为 falsy 时,回退到 `content_of(message)`。已选中的 falsy 值不会触发继续尝试低优先级实现。 | | `run_model` | firstresult | `(prompt, session_id, state) -> str` | model text | `HookRuntime.run_model` | 旧版纯文本路径。`run_model` 与 `run_model_stream` 选其一实现,不要同时实现。 | | `run_model_stream` | firstresult | `(prompt, session_id, state) -> AsyncStreamEvents` | async stream | `HookRuntime.run_model_stream` | 推荐路径。无流式实现时框架将 `run_model` 结果包装为单个 chunk 的流。 | +| `continue_prompt` | firstresult | `(tape: Tape, state: StreamState) -> str` | 下一步 prompt | `Agent._stream_events_with_auto_handoff` | 每次 agent loop 继续前调用。实现可同步或异步;builtin 保留默认 prompt 与可选 tape context。 | | `save_state` | broadcast | `(session_id, state, message, model_output) -> None` | none | `process_inbound` model-stage finally block | prompt 解析完成后,在模型阶段成功或失败时执行;若失败发生在 prompt/model 执行前则跳过。 | | `render_outbound` | broadcast | `(message, session_id, state, model_output) -> list[Envelope]` | outbound batch | `BubFramework._collect_outbounds` | 所有批次通过 `unpack_batch` 拼接;结果为空时使用默认回声 envelope。 | | `dispatch_outbound` | broadcast | `(message: Envelope) -> bool` | sent flag | `process_inbound` per outbound | 每个 outbound 都会广播给所有实现。 |