feat: add continue prompt hook

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)

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.
 </context_contract>
 """
+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

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(

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."""

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.

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)

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()

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. |

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 都会广播给所有实现。 |