From ae14142a8757e05182b9a14c10679fc47dd752a1 Mon Sep 17 00:00:00 2001 From: hanyuanling Date: Mon, 27 Apr 2026 17:51:03 +0800 Subject: [PATCH] fix(codex): stream progress deltas to channels --- nanobot/agent/hook.py | 1 + nanobot/agent/loop.py | 2 +- nanobot/agent/runner.py | 31 +++++++- nanobot/providers/base.py | 2 + nanobot/providers/openai_codex_provider.py | 2 + tests/agent/test_loop_progress.py | 86 ++++++++++++++++++++++ tests/providers/test_providers_init.py | 6 ++ 7 files changed, 128 insertions(+), 2 deletions(-) diff --git a/nanobot/agent/hook.py b/nanobot/agent/hook.py index ad6b31ab..52daf604 100644 --- a/nanobot/agent/hook.py +++ b/nanobot/agent/hook.py @@ -21,6 +21,7 @@ class AgentHookContext: tool_calls: list[ToolCallRequest] = field(default_factory=list) tool_results: list[Any] = field(default_factory=list) tool_events: list[dict[str, str]] = field(default_factory=list) + streamed_content: bool = False final_content: str | None = None stop_reason: str | None = None error: str | None = None diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index b66ca0a1..2e6185e7 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -114,7 +114,7 @@ class _LoopHook(AgentHook): async def before_execute_tools(self, context: AgentHookContext) -> None: if self._on_progress: - if not self._on_stream: + if not self._on_stream and not context.streamed_content: thought = self._loop._strip_think( context.response.content if context.response else None ) diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index be71f649..373d2181 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -21,6 +21,7 @@ from nanobot.utils.helpers import ( estimate_prompt_tokens_chain, find_legal_message_start, maybe_persist_tool_result, + strip_think, truncate_text, ) from nanobot.utils.prompt_templates import render_template @@ -607,14 +608,42 @@ class AgentRunner: messages, tools=spec.tools.get_definitions(), ) - if hook.wants_streaming(): + wants_streaming = hook.wants_streaming() + wants_progress_streaming = ( + not wants_streaming + and spec.progress_callback is not None + and getattr(self.provider, "stream_progress_via_chat_stream", False) is True + ) + + if wants_streaming: async def _stream(delta: str) -> None: + if delta: + context.streamed_content = True await hook.on_stream(context, delta) coro = self.provider.chat_stream_with_retry( **kwargs, on_content_delta=_stream, ) + elif wants_progress_streaming: + stream_buf = "" + + async def _stream_progress(delta: str) -> None: + nonlocal stream_buf + if not delta: + return + prev_clean = strip_think(stream_buf) + stream_buf += delta + new_clean = strip_think(stream_buf) + incremental = new_clean[len(prev_clean):] + if incremental: + context.streamed_content = True + await spec.progress_callback(incremental) + + coro = self.provider.chat_stream_with_retry( + **kwargs, + on_content_delta=_stream_progress, + ) else: coro = self.provider.chat_with_retry(**kwargs) diff --git a/nanobot/providers/base.py b/nanobot/providers/base.py index d7a59bab..40072b82 100644 --- a/nanobot/providers/base.py +++ b/nanobot/providers/base.py @@ -91,6 +91,8 @@ _SYNTHETIC_USER_CONTENT = "(conversation continued)" class LLMProvider(ABC): """Base class for LLM providers.""" + stream_progress_via_chat_stream = False + _CHAT_RETRY_DELAYS = (1, 2, 4) _PERSISTENT_MAX_DELAY = 60 _PERSISTENT_IDENTICAL_ERROR_LIMIT = 10 diff --git a/nanobot/providers/openai_codex_provider.py b/nanobot/providers/openai_codex_provider.py index 44cb2478..103da9a7 100644 --- a/nanobot/providers/openai_codex_provider.py +++ b/nanobot/providers/openai_codex_provider.py @@ -26,6 +26,8 @@ DEFAULT_ORIGINATOR = "nanobot" class OpenAICodexProvider(LLMProvider): """Use Codex OAuth to call the Responses API.""" + stream_progress_via_chat_stream = True + def __init__(self, default_model: str = "openai-codex/gpt-5.1-codex"): super().__init__(api_key=None, api_base=None) self.default_model = default_model diff --git a/tests/agent/test_loop_progress.py b/tests/agent/test_loop_progress.py index 8151cddf..d4662eda 100644 --- a/tests/agent/test_loop_progress.py +++ b/tests/agent/test_loop_progress.py @@ -128,3 +128,89 @@ class TestToolEventProgress: finish = finish_msgs[0].metadata["_tool_events"][0] assert finish["phase"] == "end" assert finish["result"] == "file.txt" + + @pytest.mark.asyncio + async def test_bus_progress_streams_provider_deltas_for_codex_style_provider( + self, + tmp_path: Path, + ) -> None: + """Providers that opt in can stream content deltas through _progress messages.""" + bus = MessageBus() + provider = MagicMock() + provider.stream_progress_via_chat_stream = True + provider.get_default_model.return_value = "openai-codex/gpt-5.5" + + async def chat_stream_with_retry(*, on_content_delta, **kwargs): + await on_content_delta("Hel") + await on_content_delta("lo") + return LLMResponse(content="Hello", tool_calls=[]) + + provider.chat_stream_with_retry = chat_stream_with_retry + provider.chat_with_retry = AsyncMock() + loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="openai-codex/gpt-5.5") + loop.tools.get_definitions = MagicMock(return_value=[]) + + await loop._dispatch(InboundMessage( + channel="websocket", + sender_id="u1", + chat_id="chat1", + content="say hello", + )) + + outbound = [] + while bus.outbound_size > 0: + outbound.append(await bus.consume_outbound()) + + progress = [m for m in outbound if m.metadata.get("_progress")] + final = [m for m in outbound if not m.metadata.get("_progress")] + + assert [m.content for m in progress] == ["Hel", "lo"] + assert final[-1].content == "Hello" + provider.chat_with_retry.assert_not_awaited() + + @pytest.mark.asyncio + async def test_streamed_progress_is_not_repeated_before_tool_execution( + self, + tmp_path: Path, + ) -> None: + """If content was already streamed as progress, tool setup should not repeat it.""" + loop = _make_loop(tmp_path) + loop.provider.stream_progress_via_chat_stream = True + tool_call = ToolCallRequest(id="call1", name="custom_tool", arguments={"path": "foo.txt"}) + calls = iter([ + LLMResponse(content="I will inspect it.", tool_calls=[tool_call]), + LLMResponse(content="Done", tool_calls=[]), + ]) + + async def chat_stream_with_retry(*, on_content_delta, **kwargs): + response = next(calls) + if response.tool_calls: + await on_content_delta("I will ") + await on_content_delta("inspect it.") + return response + + loop.provider.chat_stream_with_retry = chat_stream_with_retry + loop.provider.chat_with_retry = AsyncMock() + loop.tools.get_definitions = MagicMock(return_value=[]) + loop.tools.prepare_call = MagicMock(return_value=(None, {"path": "foo.txt"}, None)) + loop.tools.execute = AsyncMock(return_value="ok") + + progress: list[tuple[str, bool, list[dict] | None]] = [] + + async def on_progress( + content: str, + *, + tool_hint: bool = False, + tool_events: list[dict] | None = None, + ) -> None: + progress.append((content, tool_hint, tool_events)) + + final_content, _, _, _, _ = await loop._run_agent_loop([], on_progress=on_progress) + + assert final_content == "Done" + assert [item[0] for item in progress[:3]] == [ + "I will", + " inspect it.", + 'custom_tool("foo.txt")', + ] + assert all(item[0] != "I will inspect it." for item in progress) diff --git a/tests/providers/test_providers_init.py b/tests/providers/test_providers_init.py index d6912b43..97c869de 100644 --- a/tests/providers/test_providers_init.py +++ b/tests/providers/test_providers_init.py @@ -41,3 +41,9 @@ def test_explicit_provider_import_still_works(monkeypatch) -> None: assert namespace["AnthropicProvider"].__name__ == "AnthropicProvider" assert "nanobot.providers.anthropic_provider" in sys.modules + + +def test_openai_codex_opts_into_progress_streaming() -> None: + from nanobot.providers.openai_codex_provider import OpenAICodexProvider + + assert OpenAICodexProvider.stream_progress_via_chat_stream is True