From d576804f23fc06968cccec4f991084aedd211df8 Mon Sep 17 00:00:00 2001 From: chengyongru <2755839590@qq.com> Date: Sun, 26 Jul 2026 20:15:29 +0800 Subject: [PATCH 01/48] feat(channels): enable tool hints by default --- docs/channel-package-guide.md | 4 ++-- docs/configuration.md | 9 +++++---- nanobot/channels/base.py | 2 +- nanobot/channels/mattermost/runtime.py | 2 +- .../mattermost/tests/test_mattermost_channel.py | 4 ++++ nanobot/config/schema.py | 2 +- .../test_channel_manager_delta_coalescing.py | 15 +++++++++------ tests/channels/test_channel_plugins.py | 5 ++++- 8 files changed, 27 insertions(+), 16 deletions(-) diff --git a/docs/channel-package-guide.md b/docs/channel-package-guide.md index d53904b39..1ea4cb975 100644 --- a/docs/channel-package-guide.md +++ b/docs/channel-package-guide.md @@ -618,7 +618,7 @@ async def send(self, msg: OutboundMessage) -> None: await self._send_message(msg.chat_id, msg.content, media=msg.media) ``` -Tool hints are off by default for most channels. Users can enable them globally or per channel: +Tool hints are on by default. Users can disable them globally or per channel: ```json { @@ -626,7 +626,7 @@ Tool hints are off by default for most channels. Users can enable them globally "sendToolHints": true, "webhook": { "enabled": true, - "sendToolHints": true + "sendToolHints": false } } } diff --git a/docs/configuration.md b/docs/configuration.md index 3f461a07e..ed80cf62b 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -1555,7 +1555,7 @@ Global settings that apply to all channels. Configure under the `channels` secti { "channels": { "sendProgress": true, - "sendToolHints": false, + "sendToolHints": true, "extractDocumentText": true, "sendMaxRetries": 3, "telegram": { @@ -1568,7 +1568,7 @@ Global settings that apply to all channels. Configure under the `channels` secti | Setting | Default | Description | |---------|---------|-------------| | `sendProgress` | `true` | Stream agent's text progress to the channel | -| `sendToolHints` | `false` | Stream tool-call hints (e.g. `read_file("…")`) | +| `sendToolHints` | `true` | Stream tool-call hints (e.g. `read_file("…")`) | | `showReasoning` | `true` | Allow channels to surface model reasoning/thinking content (DeepSeek-R1 `reasoning_content`, Anthropic `thinking_blocks`, inline `` tags). Reasoning flows as a dedicated stream with `_reasoning_delta` / `_reasoning_end` markers — channels override `send_reasoning_delta` / `send_reasoning_end` to render in-place updates. Even with `true`, channels without those overrides stay no-op silently. Currently surfaced on CLI and WebSocket/WebUI (italic shimmer header, auto-collapses after the stream ends); Telegram / Slack / Discord / Feishu / WeChat / Matrix / Mattermost keep the base no-op until their bubble UI is adapted. Independent of `sendProgress`. | | `extractDocumentText` | `true` | Extract supported document/text attachments into the model prompt. PDF, DOCX, XLSX, and PPTX readers are included in the standard installation. Set to `false` to keep document content out of the prompt and include attachment path references instead. | | `sendMaxRetries` | `3` | Max delivery attempts per outbound message, including the initial send (0-10 configured, minimum 1 actual attempt) | @@ -1581,10 +1581,11 @@ Global settings that apply to all channels. Configure under the `channels` secti { "channels": { "sendProgress": true, - "sendToolHints": false, + "sendToolHints": true, "telegram": { "enabled": true, - "sendProgress": false + "sendProgress": false, + "sendToolHints": false }, "websocket": { "enabled": true, diff --git a/nanobot/channels/base.py b/nanobot/channels/base.py index 4dcbbf96b..c4f51b7ac 100644 --- a/nanobot/channels/base.py +++ b/nanobot/channels/base.py @@ -29,7 +29,7 @@ class BaseChannel(ABC): name: str = "base" display_name: str = "Base" send_progress: bool = True - send_tool_hints: bool = False + send_tool_hints: bool = True show_reasoning: bool = True def __init__(self, config: Any, bus: MessageBus): diff --git a/nanobot/channels/mattermost/runtime.py b/nanobot/channels/mattermost/runtime.py index 1a2297312..756daed73 100644 --- a/nanobot/channels/mattermost/runtime.py +++ b/nanobot/channels/mattermost/runtime.py @@ -56,7 +56,7 @@ class MattermostConfig(Base): react_emoji: str = "eyes" done_emoji: str = "white_check_mark" send_progress: bool = True - send_tool_hints: bool = False + send_tool_hints: bool = True dm: MattermostDMConfig = Field(default_factory=MattermostDMConfig) diff --git a/nanobot/channels/mattermost/tests/test_mattermost_channel.py b/nanobot/channels/mattermost/tests/test_mattermost_channel.py index 2371fef37..3b27ccf1c 100644 --- a/nanobot/channels/mattermost/tests/test_mattermost_channel.py +++ b/nanobot/channels/mattermost/tests/test_mattermost_channel.py @@ -119,6 +119,7 @@ def test_config_defaults(): assert config.token == "" assert config.streaming is True assert config.streaming_max_chars == 16000 + assert config.send_tool_hints is True assert config.dm.enabled is True assert config.dm.policy == "open" assert config.reply_in_thread is True @@ -131,6 +132,7 @@ def test_config_camelcase_aliases(): "allowFromMatchMode": "username", "streamingMaxChars": 8000, "replyInThread": False, + "sendToolHints": False, } config = MattermostConfig.model_validate(raw) assert config.server_url == "https://mm.example.com" @@ -138,11 +140,13 @@ def test_config_camelcase_aliases(): assert config.allow_from_match_mode == "username" assert config.streaming_max_chars == 8000 assert config.reply_in_thread is False + assert config.send_tool_hints is False def test_config_default_config_classmethod(): d = MattermostChannel.default_config() assert d["enabled"] is False + assert d["sendToolHints"] is True assert d["serverUrl"] == "" assert d["token"] == "" diff --git a/nanobot/config/schema.py b/nanobot/config/schema.py index fef6c2610..8d264ce6b 100644 --- a/nanobot/config/schema.py +++ b/nanobot/config/schema.py @@ -30,7 +30,7 @@ class ChannelsConfig(Base): model_config = ConfigDict(extra="allow") send_progress: bool = True # stream agent's text progress to the channel - send_tool_hints: bool = False # stream tool-call hints (e.g. read_file("…")) + send_tool_hints: bool = True # stream tool-call hints (e.g. read_file("…")) show_reasoning: bool = True # surface model reasoning when channel implements it extract_document_text: bool = True # extract text from document attachments before sending to the model send_max_retries: int = Field(default=3, ge=0, le=10) # Max delivery attempts (initial send included) diff --git a/tests/channels/test_channel_manager_delta_coalescing.py b/tests/channels/test_channel_manager_delta_coalescing.py index 6e03df5ef..de60124e2 100644 --- a/tests/channels/test_channel_manager_delta_coalescing.py +++ b/tests/channels/test_channel_manager_delta_coalescing.py @@ -74,7 +74,7 @@ def bus(): @pytest.fixture def manager(config, bus): manager = ChannelManager(config, bus) - manager.channels["mock"] = MockChannel({}, bus) + manager.channels["mock"] = manager._build_channel("mock", MockChannel, {}) return manager @@ -284,14 +284,17 @@ class TestProgressFiltering: def test_progress_visibility_uses_global_defaults(self, manager): assert manager._should_send_progress("mock", tool_hint=False) is True - assert manager._should_send_progress("mock", tool_hint=True) is False + assert manager._should_send_progress("mock", tool_hint=True) is True - def test_progress_visibility_uses_channel_overrides(self, manager): - manager.channels["mock"].send_progress = False - manager.channels["mock"].send_tool_hints = True + def test_progress_visibility_uses_channel_overrides(self, manager, bus): + manager.channels["mock"] = manager._build_channel( + "mock", + MockChannel, + {"sendProgress": False, "sendToolHints": False}, + ) assert manager._should_send_progress("mock", tool_hint=False) is False - assert manager._should_send_progress("mock", tool_hint=True) is True + assert manager._should_send_progress("mock", tool_hint=True) is False def test_progress_visibility_returns_false_for_missing_channel(self, manager): assert manager._should_send_progress("nonexistent", tool_hint=False) is False diff --git a/tests/channels/test_channel_plugins.py b/tests/channels/test_channel_plugins.py index 8399242fb..50ff6905f 100644 --- a/tests/channels/test_channel_plugins.py +++ b/tests/channels/test_channel_plugins.py @@ -269,9 +269,12 @@ def test_channels_config_has_no_per_channel_fields(): cfg = ChannelsConfig() assert not hasattr(cfg, "telegram") assert cfg.send_progress is True - assert cfg.send_tool_hints is False + assert cfg.send_tool_hints is True assert cfg.extract_document_text is True + opted_out = ChannelsConfig.model_validate({"sendToolHints": False}) + assert opted_out.send_tool_hints is False + def test_channels_config_extract_document_text_accepts_camel_alias(): cfg = ChannelsConfig.model_validate({"extractDocumentText": False}) From fb881543778d836302f65e40d0304dedf9673663 Mon Sep 17 00:00:00 2001 From: santhreal <64453045+santhreal@users.noreply.github.com> Date: Sat, 25 Jul 2026 21:49:46 -0700 Subject: [PATCH 02/48] fix(feishu): tolerate null text fields when extracting post content --- nanobot/channels/feishu/runtime.py | 14 ++++++++++--- .../tests/test_feishu_card_extraction.py | 21 +++++++++++++++++++ 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/nanobot/channels/feishu/runtime.py b/nanobot/channels/feishu/runtime.py index 02473e668..b6ddb2254 100644 --- a/nanobot/channels/feishu/runtime.py +++ b/nanobot/channels/feishu/runtime.py @@ -356,7 +356,8 @@ def _extract_post_content(content_json: dict) -> tuple[str, list[str]]: if not isinstance(block, dict) or not isinstance(block.get("content"), list): return None, [] texts, images = [], [] - if title := block.get("title"): + title = block.get("title") + if isinstance(title, str) and title: texts.append(title) for row in block["content"]: if not isinstance(row, list): @@ -366,12 +367,19 @@ def _extract_post_content(content_json: dict) -> tuple[str, list[str]]: continue tag = el.get("tag") if tag in ("text", "a"): - texts.append(el.get("text", "")) + text = el.get("text", "") + if isinstance(text, str): + texts.append(text) elif tag == "at": - texts.append(f"@{el.get('user_name', 'user')}") + user = el.get("user_name", "user") + texts.append(f"@{user if isinstance(user, str) and user else 'user'}") elif tag == "code_block": lang = el.get("language", "") code_text = el.get("text", "") + if not isinstance(lang, str): + lang = "" + if not isinstance(code_text, str): + code_text = "" texts.append(f"\n```{lang}\n{code_text}\n```\n") elif tag == "img" and (key := el.get("image_key")): images.append(key) diff --git a/nanobot/channels/feishu/tests/test_feishu_card_extraction.py b/nanobot/channels/feishu/tests/test_feishu_card_extraction.py index ccbe8d32a..45d1ea3a9 100644 --- a/nanobot/channels/feishu/tests/test_feishu_card_extraction.py +++ b/nanobot/channels/feishu/tests/test_feishu_card_extraction.py @@ -37,3 +37,24 @@ def test_extract_interactive_card_reads_table_rows() -> None: } assert _extract_share_card_content(content, "interactive") == "Name | Score\nAlice | 98" + + +from nanobot.channels.feishu.runtime import _extract_post_content + + +def test_extract_post_content_tolerates_null_text_fields() -> None: + text, images = _extract_post_content( + { + "title": "T", + "content": [ + [ + {"tag": "text", "text": None}, + {"tag": "a", "text": None}, + {"tag": "text", "text": "ok"}, + {"tag": "code_block", "language": None, "text": None}, + ] + ], + } + ) + assert "ok" in text + assert images == [] From a7cac65c7606f28c7def302a166f57bd1b495cd8 Mon Sep 17 00:00:00 2001 From: santhreal <64453045+santhreal@users.noreply.github.com> Date: Sat, 25 Jul 2026 23:29:55 -0700 Subject: [PATCH 03/48] fix(feishu): move post extract test import to module top --- .../channels/feishu/tests/test_feishu_card_extraction.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/nanobot/channels/feishu/tests/test_feishu_card_extraction.py b/nanobot/channels/feishu/tests/test_feishu_card_extraction.py index 45d1ea3a9..c1b55560f 100644 --- a/nanobot/channels/feishu/tests/test_feishu_card_extraction.py +++ b/nanobot/channels/feishu/tests/test_feishu_card_extraction.py @@ -1,6 +1,9 @@ import json -from nanobot.channels.feishu.runtime import _extract_share_card_content +from nanobot.channels.feishu.runtime import ( + _extract_post_content, + _extract_share_card_content, +) def test_extract_interactive_card_reads_user_dsl_body_elements() -> None: @@ -39,9 +42,6 @@ def test_extract_interactive_card_reads_table_rows() -> None: assert _extract_share_card_content(content, "interactive") == "Name | Score\nAlice | 98" -from nanobot.channels.feishu.runtime import _extract_post_content - - def test_extract_post_content_tolerates_null_text_fields() -> None: text, images = _extract_post_content( { From 30750060ce7a033d2637bdded0c745ee1278b57c Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Sun, 26 Jul 2026 22:44:09 +0800 Subject: [PATCH 04/48] test(feishu): cover null post metadata fields --- .../channels/feishu/tests/test_feishu_card_extraction.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/nanobot/channels/feishu/tests/test_feishu_card_extraction.py b/nanobot/channels/feishu/tests/test_feishu_card_extraction.py index c1b55560f..b2f411d9e 100644 --- a/nanobot/channels/feishu/tests/test_feishu_card_extraction.py +++ b/nanobot/channels/feishu/tests/test_feishu_card_extraction.py @@ -42,19 +42,21 @@ def test_extract_interactive_card_reads_table_rows() -> None: assert _extract_share_card_content(content, "interactive") == "Name | Score\nAlice | 98" -def test_extract_post_content_tolerates_null_text_fields() -> None: +def test_extract_post_content_tolerates_null_fields() -> None: text, images = _extract_post_content( { - "title": "T", + "title": None, "content": [ [ {"tag": "text", "text": None}, {"tag": "a", "text": None}, + {"tag": "at", "user_name": None}, {"tag": "text", "text": "ok"}, {"tag": "code_block", "language": None, "text": None}, ] ], } ) + assert "@user" in text assert "ok" in text assert images == [] From 1e505ff405e1717a7913a2320e02d098a4ef0b5a Mon Sep 17 00:00:00 2001 From: santhreal <64453045+santhreal@users.noreply.github.com> Date: Sat, 25 Jul 2026 21:48:53 -0700 Subject: [PATCH 05/48] fix(triggers): coerce string lastRunAtMs when loading local triggers --- nanobot/triggers/local_types.py | 9 +++++++- tests/triggers/test_local_triggers.py | 32 +++++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/nanobot/triggers/local_types.py b/nanobot/triggers/local_types.py index f73439631..d3fa0a284 100644 --- a/nanobot/triggers/local_types.py +++ b/nanobot/triggers/local_types.py @@ -16,6 +16,13 @@ def _int_or_zero(value: Any) -> int: return 0 if value is None or value == "" else int(value) +def _optional_int(value: Any) -> int | None: + """Coerce a stored JSON numeric; null/blank stays None.""" + if value is None or value == "": + return None + return int(value) + + @dataclass class TriggerRunRecord: """A single local trigger delivery record.""" @@ -77,7 +84,7 @@ class LocalTrigger: origin_metadata=dict(_get(data, "originMetadata", "origin_metadata", {}) or {}), created_at_ms=_int_or_zero(_get(data, "createdAtMs", "created_at_ms", 0)), updated_at_ms=_int_or_zero(_get(data, "updatedAtMs", "updated_at_ms", 0)), - last_run_at_ms=_get(data, "lastRunAtMs", "last_run_at_ms"), + last_run_at_ms=_optional_int(_get(data, "lastRunAtMs", "last_run_at_ms")), last_status=_get(data, "lastStatus", "last_status"), # type: ignore[arg-type] last_error=_get(data, "lastError", "last_error"), run_history=history, diff --git a/tests/triggers/test_local_triggers.py b/tests/triggers/test_local_triggers.py index 41fe91d05..d4fd279f6 100644 --- a/tests/triggers/test_local_triggers.py +++ b/tests/triggers/test_local_triggers.py @@ -585,3 +585,35 @@ def test_local_trigger_from_dict_accepts_null_run_at_ms() -> None: ) assert delivery.created_at_ms == 0 assert delivery.attempts == 0 + + +def test_local_trigger_from_dict_coerces_string_last_run_at_ms() -> None: + """String lastRunAtMs must coerce to int like cron store ms fields.""" + trigger = LocalTrigger.from_dict( + { + "id": "t1", + "name": "n", + "enabled": True, + "channel": "websocket", + "chatId": "c1", + "sessionKey": "websocket:c1", + "lastRunAtMs": "1710000000000", + "createdAtMs": 1, + "updatedAtMs": 1, + } + ) + assert trigger.last_run_at_ms == 1710000000000 + assert trigger.last_run_at_ms < 1710000000001 + + trigger_null = LocalTrigger.from_dict( + { + "id": "t2", + "name": "n", + "enabled": True, + "sessionKey": "websocket:c1", + "lastRunAtMs": None, + "createdAtMs": 1, + "updatedAtMs": 1, + } + ) + assert trigger_null.last_run_at_ms is None From aaf2eef5688e1fb6d3c34e07e53ffdb291d42896 Mon Sep 17 00:00:00 2001 From: santhreal <64453045+santhreal@users.noreply.github.com> Date: Sat, 25 Jul 2026 21:47:05 -0700 Subject: [PATCH 06/48] fix(feishu): tolerate null multi_url and list fields in card extract --- nanobot/channels/feishu/runtime.py | 19 +++++++++----- .../tests/test_feishu_card_extraction.py | 26 +++++++++++++++++++ 2 files changed, 38 insertions(+), 7 deletions(-) diff --git a/nanobot/channels/feishu/runtime.py b/nanobot/channels/feishu/runtime.py index b6ddb2254..51f6d8b56 100644 --- a/nanobot/channels/feishu/runtime.py +++ b/nanobot/channels/feishu/runtime.py @@ -269,7 +269,7 @@ def _extract_element_content(element: dict) -> list[str]: parts.append(text_content) elif isinstance(text, str): parts.append(text) - for field in element.get("fields", []): + for field in element.get("fields") or []: if isinstance(field, dict): field_text = field.get("text", {}) if isinstance(field_text, dict): @@ -291,7 +291,10 @@ def _extract_element_content(element: dict) -> list[str]: c = text.get("content", "") if c: parts.append(c) - url = element.get("url", "") or element.get("multi_url", {}).get("url", "") + multi_url = element.get("multi_url") or {} + url = element.get("url", "") or ( + multi_url.get("url", "") if isinstance(multi_url, dict) else "" + ) if url: parts.append(f"link: {url}") @@ -300,12 +303,14 @@ def _extract_element_content(element: dict) -> list[str]: parts.append(alt.get("content", "[image]") if isinstance(alt, dict) else "[image]") elif tag == "note": - for ne in element.get("elements", []): + for ne in element.get("elements") or []: parts.extend(_extract_element_content(ne)) elif tag == "column_set": - for col in element.get("columns", []): - for ce in col.get("elements", []): + for col in element.get("columns") or []: + if not isinstance(col, dict): + continue + for ce in col.get("elements") or []: parts.extend(_extract_element_content(ce)) elif tag == "plain_text": @@ -319,7 +324,7 @@ def _extract_element_content(element: dict) -> list[str]: for column in (element.get("columns") or []) if isinstance(column, dict) and column.get("name") ] - rows = element.get("rows", []) + rows = element.get("rows") or [] if columns: parts.append(" | ".join(header for _, header in columns)) if isinstance(rows, list): @@ -337,7 +342,7 @@ def _extract_element_content(element: dict) -> list[str]: parts.append(row_text) else: - for ne in element.get("elements", []): + for ne in element.get("elements") or []: parts.extend(_extract_element_content(ne)) return parts diff --git a/nanobot/channels/feishu/tests/test_feishu_card_extraction.py b/nanobot/channels/feishu/tests/test_feishu_card_extraction.py index b2f411d9e..c403857a4 100644 --- a/nanobot/channels/feishu/tests/test_feishu_card_extraction.py +++ b/nanobot/channels/feishu/tests/test_feishu_card_extraction.py @@ -1,6 +1,7 @@ import json from nanobot.channels.feishu.runtime import ( + _extract_element_content, _extract_post_content, _extract_share_card_content, ) @@ -60,3 +61,28 @@ def test_extract_post_content_tolerates_null_fields() -> None: assert "@user" in text assert "ok" in text assert images == [] + + +def test_extract_button_tolerates_null_multi_url() -> None: + element = {"tag": "button", "text": {"content": "Go"}, "multi_url": None} + assert _extract_element_content(element) == ["Go"] + + +def test_extract_column_set_tolerates_null_columns_and_elements() -> None: + assert _extract_element_content({"tag": "column_set", "columns": None}) == [] + assert _extract_element_content( + {"tag": "column_set", "columns": [{"elements": None}]} + ) == [] + + +def test_extract_div_tolerates_null_fields() -> None: + assert _extract_element_content( + {"tag": "div", "text": {"content": "hi"}, "fields": None} + ) == ["hi"] + + +def test_interactive_card_button_null_multi_url() -> None: + content = { + "elements": [{"tag": "button", "text": {"content": "Go"}, "multi_url": None}] + } + assert _extract_share_card_content(content, "interactive") == "Go" From 07c3e02d5c4230da9726e9e280dbba2381d142b9 Mon Sep 17 00:00:00 2001 From: santhreal <64453045+santhreal@users.noreply.github.com> Date: Sat, 25 Jul 2026 21:47:05 -0700 Subject: [PATCH 07/48] fix(triggers): treat null runHistory as empty when loading triggers --- nanobot/triggers/local_types.py | 3 ++- tests/triggers/test_local_triggers.py | 18 ++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/nanobot/triggers/local_types.py b/nanobot/triggers/local_types.py index d3fa0a284..9dc890ad6 100644 --- a/nanobot/triggers/local_types.py +++ b/nanobot/triggers/local_types.py @@ -68,9 +68,10 @@ class LocalTrigger: @classmethod def from_dict(cls, data: dict[str, Any]) -> "LocalTrigger": + raw_history = data.get("runHistory", data.get("run_history", [])) or [] history = [ record if isinstance(record, TriggerRunRecord) else TriggerRunRecord.from_dict(record) - for record in data.get("runHistory", data.get("run_history", [])) + for record in raw_history if isinstance(record, (dict, TriggerRunRecord)) ] return cls( diff --git a/tests/triggers/test_local_triggers.py b/tests/triggers/test_local_triggers.py index d4fd279f6..b4d6304c9 100644 --- a/tests/triggers/test_local_triggers.py +++ b/tests/triggers/test_local_triggers.py @@ -617,3 +617,21 @@ def test_local_trigger_from_dict_coerces_string_last_run_at_ms() -> None: } ) assert trigger_null.last_run_at_ms is None + + +def test_local_trigger_from_dict_accepts_null_run_history() -> None: + """Null runHistory must load as empty, matching CronJobState.from_store_dict.""" + trigger = LocalTrigger.from_dict( + { + "id": "t1", + "name": "n", + "enabled": True, + "channel": "websocket", + "chatId": "c1", + "sessionKey": "websocket:c1", + "runHistory": None, + "createdAtMs": 1, + "updatedAtMs": 1, + } + ) + assert trigger.run_history == [] From eb93060f95d1148d0d0acb9d44aeb3736d92bd1d Mon Sep 17 00:00:00 2001 From: yu-xin-c <2182712990@qq.com> Date: Sat, 25 Jul 2026 21:28:57 +0800 Subject: [PATCH 08/48] fix(agent): preserve pending runtime context --- nanobot/agent/loop.py | 50 +++++++++++--- tests/agent/test_runner_injections.py | 94 +++++++++++++++++++++++++++ 2 files changed, 136 insertions(+), 8 deletions(-) diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index fba15e00b..7ac786d18 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -745,14 +745,23 @@ class AgentLoop: self, ctx: TurnContext, ) -> list[RuntimeContextBlock]: - tools = ctx.tools or self.tools + assert ctx.request_context is not None + return await self._resolve_runtime_context_for_request( + ctx.request_context, + ctx.tools or self.tools, + ) + + async def _resolve_runtime_context_for_request( + self, + request: RequestContext, + tools: ToolRegistry, + ) -> list[RuntimeContextBlock]: providers = [ *tools.get_runtime_context_providers(), *self._runtime_context_providers, ] - assert ctx.request_context is not None - blocks = runtime_context_blocks_from_metadata(ctx.request_context.metadata) - blocks.extend(await resolve_runtime_context(providers, ctx.request_context)) + blocks = runtime_context_blocks_from_metadata(request.metadata) + blocks.extend(await resolve_runtime_context(providers, request)) return blocks async def _dispatch_command_inline( @@ -855,7 +864,7 @@ class AgentLoop: if pending_queue is None: return [] - def _to_user_message(pending_msg: InboundMessage) -> dict[str, Any]: + async def _to_user_message(pending_msg: InboundMessage) -> dict[str, Any]: content = pending_msg.content media = pending_msg.media if pending_msg.media else None if media: @@ -864,6 +873,31 @@ class AgentLoop: user_content = self.context._build_user_content(content, media) row: dict[str, Any] = {"role": "user", "content": user_content} metadata = pending_msg.metadata if isinstance(pending_msg.metadata, dict) else {} + if pending_msg.channel != "system": + scope = self.workspace_scopes.for_turn( + channel=pending_msg.channel, + message_metadata=metadata, + session_metadata=session.metadata if session is not None else None, + ) + pending_request = RequestContext( + channel=pending_msg.channel, + chat_id=pending_msg.chat_id, + message_id=metadata.get("message_id"), + session_key=active_session_key, + original_user_text=pending_msg.content, + runtime=runtime, + metadata=dict(metadata), + sender_id=pending_msg.sender_id, + turn_id=request_ctx.turn_id, + workspace=scope.project_path, + ) + blocks = await self._resolve_runtime_context_for_request( + pending_request, + effective_tools, + ) + row["content"], marker = append_runtime_context(user_content, blocks) + if marker is not None: + row["_meta"] = {RUNTIME_CONTEXT_MESSAGE_META: marker} if ( pending_msg.sender_id == "subagent" and metadata.get("injected_event") == "subagent_result" @@ -880,7 +914,7 @@ class AgentLoop: items: list[dict[str, Any]] = [] while len(items) < limit: try: - items.append(_to_user_message(pending_queue.get_nowait())) + items.append(await _to_user_message(pending_queue.get_nowait())) except asyncio.QueueEmpty: break @@ -898,10 +932,10 @@ class AgentLoop: session.key, ) return items - items.append(_to_user_message(msg)) + items.append(await _to_user_message(msg)) while len(items) < limit: try: - items.append(_to_user_message(pending_queue.get_nowait())) + items.append(await _to_user_message(pending_queue.get_nowait())) except asyncio.QueueEmpty: break diff --git a/tests/agent/test_runner_injections.py b/tests/agent/test_runner_injections.py index b3c39fc60..9270ba1e6 100644 --- a/tests/agent/test_runner_injections.py +++ b/tests/agent/test_runner_injections.py @@ -468,6 +468,100 @@ async def test_loop_injected_followup_preserves_image_media(tmp_path): ) +@pytest.mark.asyncio +async def test_pending_injection_resolves_its_own_runtime_context(tmp_path): + from nanobot.agent.loop import AgentLoop + from nanobot.bus.events import InboundMessage + from nanobot.bus.queue import MessageBus + from nanobot.runtime_context import ( + RUNTIME_CONTEXT_MESSAGE_META, + RuntimeContextBlock, + public_history_message, + wrap_runtime_context_lines, + ) + + provider = MagicMock() + provider.get_default_model.return_value = "test-model" + provider.chat_with_retry = AsyncMock(side_effect=[ + LLMResponse(content="first answer", tool_calls=[], usage={}), + LLMResponse(content="second answer", tool_calls=[], usage={}), + ]) + loop = AgentLoop( + bus=MessageBus(), + provider=provider, + workspace=tmp_path, + model="test-model", + ) + loop.tools.get_definitions = MagicMock(return_value=[]) + seen_contexts = [] + + async def provide_identity(request): + seen_contexts.append(( + request.channel, + request.chat_id, + request.sender_id, + request.message_id, + request.session_key, + request.original_user_text, + request.metadata["sender_name"], + request.metadata["thread_id"], + )) + return RuntimeContextBlock( + source="identity", + content=wrap_runtime_context_lines([ + " | ".join(str(value) for value in seen_contexts[-1]), + ]), + ) + + loop.register_runtime_context_provider(provide_identity) + session = loop.sessions.get_or_create("telegram:group-1") + pending_queue = asyncio.Queue() + await pending_queue.put(InboundMessage( + channel="telegram", + sender_id="user-b", + chat_id="group-1", + content="follow-up from the second speaker", + metadata={ + "message_id": "message-2", + "sender_name": "Bob", + "thread_id": "topic-7", + }, + )) + + _, _, all_messages, _, _ = await loop._run_agent_loop( + [{"role": "user", "content": "initial message from user A"}], + runtime=loop.llm_runtime(), + session=session, + channel="telegram", + chat_id="group-1", + session_key=session.key, + pending_queue=pending_queue, + ) + + assert seen_contexts == [( + "telegram", + "group-1", + "user-b", + "message-2", + session.key, + "follow-up from the second speaker", + "Bob", + "topic-7", + )] + + injected = [message for message in all_messages if message.get("role") == "user"][-1] + assert "follow-up from the second speaker" in str(injected["content"]) + model_messages = provider.chat_with_retry.await_args_list[-1].kwargs["messages"] + assert "telegram | group-1 | user-b | message-2" in str(model_messages) + assert "Bob | topic-7" in str(model_messages) + assert injected["_meta"][RUNTIME_CONTEXT_MESSAGE_META]["sources"] == ["identity"] + + loop._save_turn(session, all_messages, skip=1) + persisted = [message for message in session.messages if message.get("role") == "user"][-1] + assert "telegram | group-1 | user-b | message-2" in str(persisted["content"]) + assert public_history_message(persisted)["content"] == "follow-up from the second speaker" + + @pytest.mark.asyncio async def test_subagent_pending_injection_is_hidden_history_and_not_merged(tmp_path): from nanobot.agent.loop import AgentLoop From ff379b91cf46929918e94bf9780c3a7e4ec008f2 Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Sun, 26 Jul 2026 17:15:41 +0800 Subject: [PATCH 09/48] fix(agent): preserve merged runtime context markers --- nanobot/agent/runner.py | 52 ++++++++++++- nanobot/runtime_context.py | 64 +++++++++++++++ tests/agent/test_runner_injections.py | 107 +++++++++++++++++++++++--- 3 files changed, 208 insertions(+), 15 deletions(-) diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index 060c0cef4..9071a34ab 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -19,6 +19,11 @@ from nanobot.agent.context_governance import ( from nanobot.agent.hook import AgentHook, AgentHookContext, AgentRunHookContext from nanobot.agent.tools.registry import ToolRegistry, is_tool_error_result from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest +from nanobot.runtime_context import ( + RUNTIME_CONTEXT_MESSAGE_META, + detach_runtime_context, + reattach_runtime_context, +) from nanobot.session.history_visibility import is_hidden_history_message from nanobot.utils.helpers import ( IncrementalThinkExtractor, @@ -137,10 +142,51 @@ class AgentRunner: and not is_hidden_history_message(messages[-1]) ): merged = dict(messages[-1]) - merged["content"] = cls._merge_message_content( - merged.get("content"), - injection.get("content"), + left_meta = merged.get("_meta") + right_meta = injection.get("_meta") + left_marker = ( + left_meta.get(RUNTIME_CONTEXT_MESSAGE_META) + if isinstance(left_meta, dict) + else None ) + right_marker = ( + right_meta.get(RUNTIME_CONTEXT_MESSAGE_META) + if isinstance(right_meta, dict) + else None + ) + detached_left = ( + detach_runtime_context(merged.get("content"), left_marker) + if isinstance(left_marker, dict) + else (merged.get("content"), [], []) + ) + detached_right = ( + detach_runtime_context(injection.get("content"), right_marker) + if isinstance(right_marker, dict) + else (injection.get("content"), [], []) + ) + if detached_left is not None and detached_right is not None: + left_content, left_sources, left_blocks = detached_left + right_content, right_sources, right_blocks = detached_right + merged_content = cls._merge_message_content(left_content, right_content) + context_blocks = [*left_blocks, *right_blocks] + if context_blocks: + merged_content, marker = reattach_runtime_context( + merged_content, + [*left_sources, *right_sources], + context_blocks, + ) + internal_meta = dict(left_meta) if isinstance(left_meta, dict) else {} + if isinstance(right_meta, dict): + for key, value in right_meta.items(): + internal_meta.setdefault(key, value) + internal_meta[RUNTIME_CONTEXT_MESSAGE_META] = marker + merged["_meta"] = internal_meta + merged["content"] = merged_content + else: + merged["content"] = cls._merge_message_content( + merged.get("content"), + injection.get("content"), + ) messages[-1] = merged continue messages.append(injection) diff --git a/nanobot/runtime_context.py b/nanobot/runtime_context.py index bb45e68ed..29d9f6c0c 100644 --- a/nanobot/runtime_context.py +++ b/nanobot/runtime_context.py @@ -139,6 +139,70 @@ def append_runtime_context( } +def detach_runtime_context( + content: Any, + marker: Mapping[str, Any], +) -> tuple[Any, list[str], list[dict[str, Any]]] | None: + """Detach one validated runtime-context suffix for safe message merging.""" + if marker.get("version") != 1: + return None + raw_sources = marker.get("sources") + sources = [ + source + for source in raw_sources + if isinstance(source, str) and source + ] if isinstance(raw_sources, list) else [] + + suffix = marker.get("suffix") + if isinstance(content, str) and isinstance(suffix, str) and suffix: + if content == suffix: + clean_content = "" + elif content.endswith("\n\n" + suffix): + clean_content = content[: -(len(suffix) + 2)] + else: + return None + return clean_content, sources, [{"type": "text", "text": suffix}] + + expected = marker.get("blocks") + if isinstance(content, list) and isinstance(expected, list) and expected: + count = len(expected) + if content[-count:] != expected: + return None + return content[:-count], sources, deepcopy(expected) + return None + + +def reattach_runtime_context( + content: Any, + sources: Sequence[str], + blocks: Sequence[Mapping[str, Any]], +) -> tuple[Any, dict[str, Any]]: + """Append detached runtime-context blocks after visible messages are merged.""" + context_blocks = [deepcopy(dict(block)) for block in blocks] + if isinstance(content, str) and all( + block.get("type") == "text" and isinstance(block.get("text"), str) + for block in context_blocks + ): + suffix = "\n\n".join(block["text"] for block in context_blocks) + merged = f"{content}\n\n{suffix}" if content else suffix + return merged, { + "version": 1, + "sources": list(sources), + "suffix": suffix, + } + + visible_blocks = ( + [*content] + if isinstance(content, list) + else ([] if content is None else [{"type": "text", "text": str(content)}]) + ) + return [*visible_blocks, *context_blocks], { + "version": 1, + "sources": list(sources), + "blocks": context_blocks, + } + + def public_history_message(message: Mapping[str, Any]) -> dict[str, Any]: """Return a user-visible copy with trusted runtime context removed exactly.""" cleaned = deepcopy(dict(message)) diff --git a/tests/agent/test_runner_injections.py b/tests/agent/test_runner_injections.py index 9270ba1e6..9b96b4a8c 100644 --- a/tests/agent/test_runner_injections.py +++ b/tests/agent/test_runner_injections.py @@ -527,6 +527,17 @@ async def test_pending_injection_resolves_its_own_runtime_context(tmp_path): "thread_id": "topic-7", }, )) + await pending_queue.put(InboundMessage( + channel="telegram", + sender_id="user-c", + chat_id="group-1", + content="another follow-up", + metadata={ + "message_id": "message-3", + "sender_name": "Carol", + "thread_id": "topic-7", + }, + )) _, _, all_messages, _, _ = await loop._run_agent_loop( [{"role": "user", "content": "initial message from user A"}], @@ -538,28 +549,48 @@ async def test_pending_injection_resolves_its_own_runtime_context(tmp_path): pending_queue=pending_queue, ) - assert seen_contexts == [( - "telegram", - "group-1", - "user-b", - "message-2", - session.key, - "follow-up from the second speaker", - "Bob", - "topic-7", - )] + assert seen_contexts == [ + ( + "telegram", + "group-1", + "user-b", + "message-2", + session.key, + "follow-up from the second speaker", + "Bob", + "topic-7", + ), + ( + "telegram", + "group-1", + "user-c", + "message-3", + session.key, + "another follow-up", + "Carol", + "topic-7", + ), + ] injected = [message for message in all_messages if message.get("role") == "user"][-1] assert "follow-up from the second speaker" in str(injected["content"]) model_messages = provider.chat_with_retry.await_args_list[-1].kwargs["messages"] assert "telegram | group-1 | user-b | message-2" in str(model_messages) assert "Bob | topic-7" in str(model_messages) - assert injected["_meta"][RUNTIME_CONTEXT_MESSAGE_META]["sources"] == ["identity"] + assert "telegram | group-1 | user-c | message-3" in str(model_messages) + assert "Carol | topic-7" in str(model_messages) + assert injected["_meta"][RUNTIME_CONTEXT_MESSAGE_META]["sources"] == [ + "identity", + "identity", + ] loop._save_turn(session, all_messages, skip=1) persisted = [message for message in session.messages if message.get("role") == "user"][-1] assert "telegram | group-1 | user-b | message-2" in str(persisted["content"]) - assert public_history_message(persisted)["content"] == "follow-up from the second speaker" + assert "telegram | group-1 | user-c | message-3" in str(persisted["content"]) + assert public_history_message(persisted)["content"] == ( + "follow-up from the second speaker\n\nanother follow-up" + ) @pytest.mark.asyncio @@ -688,6 +719,58 @@ async def test_runner_merges_multiple_injected_user_messages_without_losing_medi ) +def test_runner_merge_preserves_runtime_markers_with_media() -> None: + from nanobot.agent.runner import AgentRunner + from nanobot.runtime_context import ( + RUNTIME_CONTEXT_HISTORY_META, + RUNTIME_CONTEXT_MESSAGE_META, + RuntimeContextBlock, + append_runtime_context, + public_history_message, + ) + + first_visible = [ + {"type": "text", "text": "first"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AA=="}}, + ] + first_content, first_marker = append_runtime_context( + first_visible, + [RuntimeContextBlock(source="first", content="private first")], + ) + second_content, second_marker = append_runtime_context( + "second", + [RuntimeContextBlock(source="second", content="private second")], + ) + messages: list[dict] = [] + + AgentRunner._append_injected_messages(messages, [ + { + "role": "user", + "content": first_content, + "_meta": {RUNTIME_CONTEXT_MESSAGE_META: first_marker}, + }, + { + "role": "user", + "content": second_content, + "_meta": {RUNTIME_CONTEXT_MESSAGE_META: second_marker}, + }, + ]) + + assert len(messages) == 1 + merged = messages[0] + assert "private first" in str(merged["content"]) + assert "private second" in str(merged["content"]) + persisted = { + "role": "user", + "content": merged["content"], + RUNTIME_CONTEXT_HISTORY_META: merged["_meta"][RUNTIME_CONTEXT_MESSAGE_META], + } + assert public_history_message(persisted)["content"] == [ + *first_visible, + {"type": "text", "text": "second"}, + ] + + @pytest.mark.asyncio async def test_injection_cycles_capped_at_max(): """Injection cycles should be capped at _MAX_INJECTION_CYCLES.""" From be43a5457067d15278f8bbd70aab186d2a7c7733 Mon Sep 17 00:00:00 2001 From: chengyongru <2755839590@qq.com> Date: Sun, 26 Jul 2026 23:37:57 +0800 Subject: [PATCH 10/48] fix(webui): prevent mobile thread overflow --- .../src/components/thread/ThreadViewport.tsx | 2 +- webui/src/tests/thread-viewport.test.tsx | 22 +++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/webui/src/components/thread/ThreadViewport.tsx b/webui/src/components/thread/ThreadViewport.tsx index b1bce06eb..7f47c7df6 100644 --- a/webui/src/components/thread/ThreadViewport.tsx +++ b/webui/src/components/thread/ThreadViewport.tsx @@ -638,7 +638,7 @@ export const ThreadViewport = forwardRef
{ expect(messageRegion.className).not.toContain("5rem"); }); + it("allows long messages to shrink within the shared mobile grid column", () => { + render( + composer
} + />, + ); + + expect(screen.getByTestId("thread-message-region")).toHaveClass("min-w-0"); + }); + it("top-aligns a short active turn while the agent is responding", () => { render( Date: Tue, 14 Jul 2026 21:52:26 +0800 Subject: [PATCH 11/48] fix(heartbeat): route unified sessions to last channel --- nanobot/agent/loop.py | 28 ++++++++++++++++++++++- nanobot/cli/commands.py | 18 +++++++++++++++ nanobot/session/keys.py | 30 +++++++++++++++++++++++++ tests/agent/test_loop_save_turn.py | 26 +++++++++++++++++++++ tests/cli/test_commands.py | 36 ++++++++++++++++++++++++++++++ 5 files changed, 137 insertions(+), 1 deletion(-) diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 7ac786d18..0a851b673 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -73,7 +73,7 @@ from nanobot.session.goal_state import ( sustained_goal_active, ) from nanobot.session.history_visibility import HIDDEN_HISTORY_META -from nanobot.session.keys import UNIFIED_SESSION_KEY +from nanobot.session.keys import UNIFIED_SESSION_KEY, remember_last_channel from nanobot.session.manager import ( Session, SessionManager, @@ -798,6 +798,27 @@ class AgentLoop: return UNIFIED_SESSION_KEY return msg.session_key + def _remember_unified_session_route( + self, + session: Session, + msg: InboundMessage, + *, + is_user_turn: bool, + ) -> None: + """Remember the latest user-facing route for unified-session delivery.""" + if ( + not self._unified_session + or session.key != UNIFIED_SESSION_KEY + or not is_user_turn + or msg.channel in {"cli", "system"} + or msg.sender_id == "subagent" + ): + return + _, automation_metadata = automation_history_overrides(msg.metadata) + if automation_metadata: + return + remember_last_channel(session.metadata, msg.channel, msg.chat_id) + @staticmethod def _replay_token_budget(runtime: LLMRuntime) -> int: """Derive a token budget for session history replay from the context window.""" @@ -1490,6 +1511,11 @@ class AgentLoop: # ensure it exists in case this handler is invoked independently. if ctx.session is None: ctx.session = self.sessions.get_or_create(ctx.session_key) + self._remember_unified_session_route( + ctx.session, + msg, + is_user_turn=ctx.kind is TurnKind.USER, + ) await ctx.delivery.started() if ctx.kind is TurnKind.USER: self.workspace_scopes.persist_message_scope(ctx.session, msg) diff --git a/nanobot/cli/commands.py b/nanobot/cli/commands.py index 71dd978dc..a91f408da 100644 --- a/nanobot/cli/commands.py +++ b/nanobot/cli/commands.py @@ -77,6 +77,10 @@ from nanobot.cli.stream import StreamRenderer, ThinkingSpinner # noqa: E402 from nanobot.config.paths import get_workspace_path, is_default_workspace # noqa: E402 from nanobot.config.schema import Config # noqa: E402 from nanobot.security.network import is_loopback_host # noqa: E402 +from nanobot.session.keys import ( # noqa: E402 + UNIFIED_SESSION_KEY, + last_channel_from_metadata, +) from nanobot.utils.evaluator import evaluate_response, resolve_evaluator_prompt # noqa: E402 from nanobot.utils.helpers import ( # noqa: E402 sanitize_surrogates as _sanitize_surrogates, @@ -264,6 +268,7 @@ def _pick_heartbeat_target_from_sessions( enabled_channels: Iterable[str], sessions: Iterable[dict[str, Any]], archived_keys: Iterable[str], + unified_session_metadata: dict[str, Any] | None = None, ) -> tuple[str, str]: enabled = set(enabled_channels) archived = set(archived_keys) @@ -271,6 +276,13 @@ def _pick_heartbeat_target_from_sessions( key = item.get("key") or "" if key in archived: continue + if key == UNIFIED_SESSION_KEY: + route = last_channel_from_metadata(unified_session_metadata) + if route is not None: + channel, chat_id = route + if channel not in {"cli", "system"} and channel in enabled: + return channel, chat_id + continue if ":" not in key: continue channel, chat_id = key.split(":", 1) @@ -1984,10 +1996,16 @@ def _run_gateway( def _pick_heartbeat_target() -> tuple[str, str]: """Pick a routable channel/chat target for heartbeat-triggered messages.""" sidebar_state = read_webui_sidebar_state() + unified_metadata = None + if config.agents.defaults.unified_session: + record = session_manager.read_session_metadata(UNIFIED_SESSION_KEY) + if isinstance(record, dict) and isinstance(record.get("metadata"), dict): + unified_metadata = record["metadata"] return _pick_heartbeat_target_from_sessions( enabled_channels=channels.enabled_channels, sessions=session_manager.list_sessions(), archived_keys=sidebar_state.get("archived_keys", []), + unified_session_metadata=unified_metadata, ) if channels.enabled_channels: diff --git a/nanobot/session/keys.py b/nanobot/session/keys.py index ce581bdc8..45f6d1cf7 100644 --- a/nanobot/session/keys.py +++ b/nanobot/session/keys.py @@ -2,7 +2,11 @@ from __future__ import annotations +from collections.abc import Mapping, MutableMapping +from typing import Any + UNIFIED_SESSION_KEY = "unified:default" +LAST_CHANNEL_METADATA_KEY = "last_channel" def session_key_for_channel(channel: str, chat_id: str, *, unified_session: bool = False) -> str: @@ -10,3 +14,29 @@ def session_key_for_channel(channel: str, chat_id: str, *, unified_session: bool if unified_session: return UNIFIED_SESSION_KEY return f"{channel}:{chat_id}" + + +def remember_last_channel( + metadata: MutableMapping[str, Any], + channel: str, + chat_id: str, +) -> None: + """Persist the latest concrete delivery route in session metadata.""" + if not channel or not chat_id: + return + metadata[LAST_CHANNEL_METADATA_KEY] = f"{channel}:{chat_id}" + + +def last_channel_from_metadata( + metadata: Mapping[str, Any] | None, +) -> tuple[str, str] | None: + """Return a concrete delivery route from persisted session metadata.""" + if not isinstance(metadata, Mapping): + return None + route = metadata.get(LAST_CHANNEL_METADATA_KEY) + if not isinstance(route, str) or ":" not in route: + return None + channel, chat_id = route.split(":", 1) + if not channel or not chat_id: + return None + return channel, chat_id diff --git a/tests/agent/test_loop_save_turn.py b/tests/agent/test_loop_save_turn.py index a87600a58..ed53b2c2e 100644 --- a/tests/agent/test_loop_save_turn.py +++ b/tests/agent/test_loop_save_turn.py @@ -30,6 +30,10 @@ from nanobot.runtime_context import ( ) from nanobot.session.automation_turns import AUTOMATION_HISTORY_META from nanobot.session.goal_state import GOAL_STATE_KEY +from nanobot.session.keys import ( + LAST_CHANNEL_METADATA_KEY, + UNIFIED_SESSION_KEY, +) from nanobot.session.manager import Session, SessionManager from nanobot.session.turn_continuation import ( INTERNAL_CONTINUATION_META, @@ -682,6 +686,28 @@ async def test_process_message_persists_user_message_before_turn_completes(tmp_p assert persisted.updated_at >= persisted.created_at +@pytest.mark.asyncio +async def test_process_message_persists_unified_session_delivery_route(tmp_path: Path) -> None: + loop = _make_full_loop(tmp_path) + loop._unified_session = True + loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] + loop._run_agent_loop = AsyncMock(side_effect=RuntimeError("boom")) # type: ignore[method-assign] + + msg = InboundMessage( + channel="feishu", + sender_id="u1", + chat_id="oc_123", + content="persist my route", + session_key_override=UNIFIED_SESSION_KEY, + ) + with pytest.raises(RuntimeError, match="boom"): + await loop._process_message(msg) + + loop.sessions.invalidate(UNIFIED_SESSION_KEY) + persisted = loop.sessions.get_or_create(UNIFIED_SESSION_KEY) + assert persisted.metadata[LAST_CHANNEL_METADATA_KEY] == "feishu:oc_123" + + # 1x1 PNG used by the media-persistence tests. ``extract_documents`` runs # at the top of ``_process_message`` and filters ``msg.media`` down to # paths that magic-byte-sniff as images, so the test fixture needs real diff --git a/tests/cli/test_commands.py b/tests/cli/test_commands.py index 9f94dbba7..34c9c674f 100644 --- a/tests/cli/test_commands.py +++ b/tests/cli/test_commands.py @@ -1787,6 +1787,42 @@ def test_heartbeat_target_skips_archived_webui_sessions(): assert target == ("websocket", "active") +def test_heartbeat_target_uses_last_channel_for_unified_session(): + from nanobot.cli.commands import _pick_heartbeat_target_from_sessions + from nanobot.session.keys import LAST_CHANNEL_METADATA_KEY, UNIFIED_SESSION_KEY + + target = _pick_heartbeat_target_from_sessions( + enabled_channels=["telegram", "discord"], + archived_keys=[], + sessions=[{"key": UNIFIED_SESSION_KEY}], + unified_session_metadata={LAST_CHANNEL_METADATA_KEY: "discord:chat-42"}, + ) + + assert target == ("discord", "chat-42") + + +@pytest.mark.parametrize( + "metadata", + [ + {"last_channel": "telegram:chat-42"}, + {"last_channel": "cli:direct"}, + {"last_channel": "invalid"}, + ], +) +def test_heartbeat_target_rejects_unroutable_unified_metadata(metadata): + from nanobot.cli.commands import _pick_heartbeat_target_from_sessions + from nanobot.session.keys import UNIFIED_SESSION_KEY + + target = _pick_heartbeat_target_from_sessions( + enabled_channels=["discord"], + archived_keys=[], + sessions=[{"key": UNIFIED_SESSION_KEY}], + unified_session_metadata=metadata, + ) + + assert target == ("cli", "direct") + + def _write_instance_config(tmp_path: Path) -> Path: config_file = tmp_path / "instance" / "config.json" config_file.parent.mkdir(parents=True) From 5d8046deefdda8c0a1d8a6f79c195add552335a3 Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Sun, 26 Jul 2026 23:59:00 +0800 Subject: [PATCH 12/48] test(heartbeat): cover ignored unified routes --- nanobot/agent/loop.py | 2 +- tests/agent/test_loop_save_turn.py | 57 ++++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 0a851b673..f079867bb 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -1514,7 +1514,7 @@ class AgentLoop: self._remember_unified_session_route( ctx.session, msg, - is_user_turn=ctx.kind is TurnKind.USER, + is_user_turn=ctx.original_user_text is not None, ) await ctx.delivery.started() if ctx.kind is TurnKind.USER: diff --git a/tests/agent/test_loop_save_turn.py b/tests/agent/test_loop_save_turn.py index ed53b2c2e..59946f915 100644 --- a/tests/agent/test_loop_save_turn.py +++ b/tests/agent/test_loop_save_turn.py @@ -708,6 +708,63 @@ async def test_process_message_persists_unified_session_delivery_route(tmp_path: assert persisted.metadata[LAST_CHANNEL_METADATA_KEY] == "feishu:oc_123" +@pytest.mark.parametrize( + ("msg", "is_user_turn"), + [ + ( + InboundMessage( + channel="cli", + sender_id="u1", + chat_id="direct", + content="cli input", + ), + True, + ), + ( + InboundMessage( + channel="system", + sender_id="system", + chat_id="discord:automation", + content="system event", + ), + False, + ), + ( + InboundMessage( + channel="discord", + sender_id="subagent", + chat_id="subagent-result", + content="subagent result", + ), + True, + ), + ( + InboundMessage( + channel="discord", + sender_id="u1", + chat_id="automation", + content="scheduled turn", + metadata={CRON_TRIGGER_META: {"job_id": "job-1"}}, + ), + True, + ), + ], +) +def test_unified_session_route_ignores_non_user_destinations( + tmp_path: Path, + msg: InboundMessage, + is_user_turn: bool, +) -> None: + loop = _make_full_loop(tmp_path) + loop._unified_session = True + session = loop.sessions.get_or_create(UNIFIED_SESSION_KEY) + session.metadata[LAST_CHANNEL_METADATA_KEY] = "telegram:existing" + + loop._remember_unified_session_route(session, msg, is_user_turn=is_user_turn) + + assert session.metadata[LAST_CHANNEL_METADATA_KEY] == "telegram:existing" + + # 1x1 PNG used by the media-persistence tests. ``extract_documents`` runs # at the top of ``_process_message`` and filters ``msg.media`` down to # paths that magic-byte-sniff as images, so the test fixture needs real From 01a11b39803cd106e40a747d5be03ca2987defab Mon Sep 17 00:00:00 2001 From: yu-xin-c <2182712990@qq.com> Date: Thu, 18 Jun 2026 23:19:00 +0800 Subject: [PATCH 13/48] feat(exec): allow extra bwrap bind roots --- docs/configuration.md | 2 + nanobot/agent/tools/sandbox.py | 54 ++++++++++++++++++++-- nanobot/agent/tools/shell.py | 45 +++++++++++++++++- tests/tools/test_exec_platform.py | 30 ++++++++++++ tests/tools/test_exec_security.py | 71 +++++++++++++++++++++++++++++ tests/tools/test_sandbox.py | 51 +++++++++++++++++++++ tests/tools/test_tool_validation.py | 16 +++++++ 7 files changed, 263 insertions(+), 6 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index ed80cf62b..3d4be10d9 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -1995,6 +1995,8 @@ For API keys, tokens, and other secrets, see [Environment Variables for Secrets] | `tools.exec.timeout` | `60` | Default hard timeout in seconds for shell commands. Config values may exceed the per-call tool cap; set `0` to disable the hard timeout for trusted long-running commands. | | `tools.exec.pathPrepend` | `""` | Extra directories to prepend to `PATH` when running shell commands. Use this when configured tools should win executable lookup precedence, such as a Python virtual environment's `bin` or `Scripts` directory. | | `tools.exec.pathAppend` | `""` | Extra directories to append to `PATH` when running shell commands (e.g. `/usr/sbin` for `ufw`). | +| `tools.exec.sandboxRoBinds` | `[]` | Extra absolute paths to read-only bind into the `"bwrap"` sandbox with `--ro-bind-try`, such as `/home/user/.local/bin` or `/home/user/.cargo/bin` when those paths are also in `pathPrepend`/`pathAppend`. These roots are also accepted by the shell absolute-path guard only while bwrap is active. | +| `tools.exec.sandboxRwBinds` | `[]` | Extra absolute paths to read-write bind into the `"bwrap"` sandbox with `--bind-try`, for trusted tool caches or scratch directories. Use sparingly: paths listed here are intentionally writable by shell commands inside the sandbox. | | `tools.webuiAllowRemotePackageInstall` | `false` | When `false`, the WebUI can install missing optional packages only from a browser opened on the same machine as nanobot. Set to `true` only when a trusted remote admin is allowed to install Python packages into this nanobot environment. | | `tools.ssrfWhitelist` | `[]` | CIDR ranges exempted from the shared SSRF guard used by web fetches and HTTP/SSE MCP connections. Prefer exact host CIDRs such as `192.168.1.50/32`; broad ranges increase SSRF exposure. | | `channels.*.allowFrom` | omitted | Access control per channel. Omit to use pairing-only mode; set `["*"]` to allow everyone; or list specific user IDs. See [Pairing](#pairing) for details. | diff --git a/nanobot/agent/tools/sandbox.py b/nanobot/agent/tools/sandbox.py index d1f771b27..0d602abfb 100644 --- a/nanobot/agent/tools/sandbox.py +++ b/nanobot/agent/tools/sandbox.py @@ -5,13 +5,40 @@ To add a new backend, implement a function with the signature: and register it in _BACKENDS below. """ +import os import shlex from pathlib import Path +from typing import Iterable from nanobot.config.paths import get_media_dir -def _bwrap(command: str, workspace: str, cwd: str) -> str: +def _normalize_bind_paths(paths: Iterable[str] | None) -> list[str]: + out: list[str] = [] + seen: set[str] = set() + for raw in paths or []: + value = str(raw).strip() + if not value: + continue + path = Path(os.path.expandvars(value)).expanduser() + if not path.is_absolute(): + continue + resolved = str(path.resolve(strict=False)) + if resolved in seen: + continue + seen.add(resolved) + out.append(resolved) + return out + + +def _bwrap( + command: str, + workspace: str, + cwd: str, + *, + sandbox_ro_binds: Iterable[str] | None = None, + sandbox_rw_binds: Iterable[str] | None = None, +) -> str: """Wrap command in a bubblewrap sandbox (requires bwrap in container). Only the workspace is bind-mounted read-write; its parent dir (which holds @@ -51,17 +78,34 @@ def _bwrap(command: str, workspace: str, cwd: str) -> str: "--dir", str(ws), # recreate workspace mount point "--bind", str(ws), str(ws), "--ro-bind-try", str(media), str(media), # read-only access to media - "--chdir", sandbox_cwd, - "--", "sh", "-c", command, ] + for p in _normalize_bind_paths(sandbox_ro_binds): + args += ["--ro-bind-try", p, p] + for p in _normalize_bind_paths(sandbox_rw_binds): + args += ["--bind-try", p, p] + args += ["--chdir", sandbox_cwd, "--", "sh", "-c", command] return shlex.join(args) _BACKENDS = {"bwrap": _bwrap} -def wrap_command(sandbox: str, command: str, workspace: str, cwd: str) -> str: +def wrap_command( + sandbox: str, + command: str, + workspace: str, + cwd: str, + *, + sandbox_ro_binds: Iterable[str] | None = None, + sandbox_rw_binds: Iterable[str] | None = None, +) -> str: """Wrap *command* using the named sandbox backend.""" if backend := _BACKENDS.get(sandbox): - return backend(command, workspace, cwd) + return backend( + command, + workspace, + cwd, + sandbox_ro_binds=sandbox_ro_binds, + sandbox_rw_binds=sandbox_rw_binds, + ) raise ValueError(f"Unknown sandbox backend {sandbox!r}. Available: {list(_BACKENDS)}") diff --git a/nanobot/agent/tools/shell.py b/nanobot/agent/tools/shell.py index 9fa639a79..be3e56056 100644 --- a/nanobot/agent/tools/shell.py +++ b/nanobot/agent/tools/shell.py @@ -84,6 +84,8 @@ class ExecToolConfig(Base): path_prepend: str = "" path_append: str = "" sandbox: str = "" + sandbox_ro_binds: list[str] = Field(default_factory=list) + sandbox_rw_binds: list[str] = Field(default_factory=list) allowed_env_keys: list[str] = Field(default_factory=list) allow_patterns: list[str] = Field(default_factory=list) deny_patterns: list[str] = Field(default_factory=list) @@ -187,6 +189,8 @@ class ExecTool(Tool): sandbox=cfg.sandbox, path_prepend=cfg.path_prepend, path_append=cfg.path_append, + sandbox_ro_binds=cfg.sandbox_ro_binds, + sandbox_rw_binds=cfg.sandbox_rw_binds, allowed_env_keys=cfg.allowed_env_keys, allow_patterns=cfg.allow_patterns, deny_patterns=cfg.deny_patterns, @@ -205,6 +209,8 @@ class ExecTool(Tool): sandbox: str = "", path_prepend: str = "", path_append: str = "", + sandbox_ro_binds: list[str] | None = None, + sandbox_rw_binds: list[str] | None = None, allowed_env_keys: list[str] | None = None, session_manager: Any | None = None, ): @@ -237,6 +243,8 @@ class ExecTool(Tool): self.webui_allow_local_service_access = webui_allow_local_service_access self.path_prepend = path_prepend self.path_append = path_append + self.sandbox_ro_binds = self._normalize_bind_roots(sandbox_ro_binds) + self.sandbox_rw_binds = self._normalize_bind_roots(sandbox_rw_binds) self.allowed_env_keys = allowed_env_keys or [] self._session_manager = session_manager or DEFAULT_EXEC_SESSION_MANAGER @@ -464,7 +472,14 @@ class ExecTool(Tool): ) else: workspace = workspace_root or cwd - command = wrap_command(self.sandbox, command, workspace, cwd) + command = wrap_command( + self.sandbox, + command, + workspace, + cwd, + sandbox_ro_binds=[str(p) for p in self.sandbox_ro_binds], + sandbox_rw_binds=[str(p) for p in self.sandbox_rw_binds], + ) cwd = str(Path(workspace).resolve()) effective_timeout = self._resolve_timeout(timeout) @@ -794,6 +809,7 @@ class ExecTool(Tool): if workspace_root else None ) + sandbox_bind_roots = self._active_sandbox_bind_roots() for raw in self._extract_absolute_paths(cmd): try: @@ -817,6 +833,8 @@ class ExecTool(Tool): ) if not allowed and resolved_workspace is not None: allowed = is_path_within(p, resolved_workspace) + if not allowed and sandbox_bind_roots: + allowed = any(is_path_within(p, root) for root in sandbox_bind_roots) if p.is_absolute() and not allowed: return ToolResult.error( "Error: Command blocked by safety guard (path outside working dir)" @@ -921,3 +939,28 @@ class ExecTool(Tool): posix_paths = re.findall(r"(?:^|[\s|>='\"])(/[^\s\"'>;|<]+)", command) # POSIX: /absolute only home_paths = re.findall(r"(?:^|[\s>='\"])(~[/+][^\s\"'>;|<]*)", command) # POSIX/Windows home shortcut: ~/ or ~+ return win_paths + posix_paths + home_paths + + @staticmethod + def _normalize_bind_roots(paths: list[str] | None) -> list[Path]: + roots: list[Path] = [] + seen: set[str] = set() + for raw in paths or []: + value = str(raw).strip() + if not value: + continue + path = Path(os.path.expandvars(value)).expanduser() + if not path.is_absolute(): + continue + with suppress(OSError, RuntimeError, ValueError): + resolved = path.resolve(strict=False) + key = os.path.normcase(os.fspath(resolved)) + if key in seen: + continue + seen.add(key) + roots.append(resolved) + return roots + + def _active_sandbox_bind_roots(self) -> list[Path]: + if self.sandbox != "bwrap" or _IS_WINDOWS: + return [] + return [*self.sandbox_ro_binds, *self.sandbox_rw_binds] diff --git a/tests/tools/test_exec_platform.py b/tests/tools/test_exec_platform.py index 514ccd956..4a6277b24 100644 --- a/tests/tools/test_exec_platform.py +++ b/tests/tools/test_exec_platform.py @@ -8,6 +8,7 @@ platform-specific binaries (all subprocess calls are mocked). import asyncio import shutil import sys +from pathlib import Path from unittest.mock import AsyncMock, patch import pytest @@ -473,6 +474,35 @@ class TestSandboxPlatform: spawned_cmd = mock_spawn.call_args[0][0] assert "bwrap" in spawned_cmd + @pytest.mark.asyncio + async def test_bwrap_receives_configured_bind_roots(self): + """Configured bwrap bind roots should be forwarded to the sandbox wrapper.""" + mock_proc = AsyncMock() + mock_proc.communicate.return_value = (b"sandboxed", b"") + mock_proc.returncode = 0 + + with ( + patch("nanobot.agent.tools.shell._IS_WINDOWS", False), + patch("nanobot.agent.tools.shell.wrap_command", return_value="bwrap -- sh -c ls") as mock_wrap, + patch.object(ExecTool, "_spawn", return_value=mock_proc), + patch.object(ExecTool, "_guard_command", return_value=None), + ): + tool = ExecTool( + sandbox="bwrap", + working_dir="/workspace", + sandbox_ro_binds=["/home/user/.local/bin"], + sandbox_rw_binds=["/home/user/.cache/uv"], + ) + await tool.execute(command="ls") + + kwargs = mock_wrap.call_args.kwargs + assert kwargs["sandbox_ro_binds"] == [ + str(Path("/home/user/.local/bin").resolve(strict=False)) + ] + assert kwargs["sandbox_rw_binds"] == [ + str(Path("/home/user/.cache/uv").resolve(strict=False)) + ] + # --------------------------------------------------------------------------- # end-to-end (mocked subprocess, full execute path) diff --git a/tests/tools/test_exec_security.py b/tests/tools/test_exec_security.py index 73a58ce2a..6e3a6f8a8 100644 --- a/tests/tools/test_exec_security.py +++ b/tests/tools/test_exec_security.py @@ -314,6 +314,77 @@ def test_exec_still_blocks_real_outside_path_via_redirect(tmp_path): assert "path outside working dir" in blocked +def test_exec_allows_absolute_path_inside_bwrap_ro_bind(tmp_path): + workspace = tmp_path / "workspace" + workspace.mkdir() + tool_bin = tmp_path / "home" / ".local" / "bin" + tool_bin.mkdir(parents=True) + uv = tool_bin / "uv" + uv.write_text("#!/bin/sh\n") + tool = ExecTool( + working_dir=str(workspace), + restrict_to_workspace=True, + sandbox="bwrap", + sandbox_ro_binds=[str(tool_bin)], + ) + + blocked = tool._guard_command( + f"{uv} --version", + str(workspace), + restrict_to_workspace=True, + workspace_root=str(workspace), + ) + + assert blocked is None + + +def test_exec_allows_absolute_path_inside_bwrap_rw_bind(tmp_path): + workspace = tmp_path / "workspace" + workspace.mkdir() + cache_dir = tmp_path / "cache" + cache_dir.mkdir() + tool = ExecTool( + working_dir=str(workspace), + restrict_to_workspace=True, + sandbox="bwrap", + sandbox_rw_binds=[str(cache_dir)], + ) + + blocked = tool._guard_command( + f"touch {cache_dir / 'stamp'}", + str(workspace), + restrict_to_workspace=True, + workspace_root=str(workspace), + ) + + assert blocked is None + + +def test_exec_bind_roots_do_not_widen_guard_without_bwrap(tmp_path): + workspace = tmp_path / "workspace" + workspace.mkdir() + tool_bin = tmp_path / "home" / ".local" / "bin" + tool_bin.mkdir(parents=True) + uv = tool_bin / "uv" + uv.write_text("#!/bin/sh\n") + tool = ExecTool( + working_dir=str(workspace), + restrict_to_workspace=True, + sandbox="", + sandbox_ro_binds=[str(tool_bin)], + ) + + blocked = tool._guard_command( + f"{uv} --version", + str(workspace), + restrict_to_workspace=True, + workspace_root=str(workspace), + ) + + assert blocked is not None + assert "path outside working dir" in blocked + + # --- format command blocking ----------------------------------------------- diff --git a/tests/tools/test_sandbox.py b/tests/tools/test_sandbox.py index 57495a315..752713517 100644 --- a/tests/tools/test_sandbox.py +++ b/tests/tools/test_sandbox.py @@ -150,6 +150,57 @@ class TestBwrapBackend: try_pairs = {(tokens[i + 1], tokens[i + 2]) for i in try_indices} assert (str(fake_media), str(fake_media)) in try_pairs + def test_custom_read_only_binds_use_ro_bind_try(self, tmp_path): + ws = tmp_path / "project" + tool_bin = tmp_path / "home" / ".local" / "bin" + + result = wrap_command( + "bwrap", + "uv --version", + str(ws), + str(ws), + sandbox_ro_binds=[str(tool_bin)], + ) + tokens = _parse(result) + + try_indices = [i for i, t in enumerate(tokens) if t == "--ro-bind-try"] + try_pairs = {(tokens[i + 1], tokens[i + 2]) for i in try_indices} + assert (str(tool_bin.resolve(strict=False)), str(tool_bin.resolve(strict=False))) in try_pairs + + def test_custom_read_write_binds_use_bind_try(self, tmp_path): + ws = tmp_path / "project" + cache_dir = tmp_path / "cache" + + result = wrap_command( + "bwrap", + "touch cache/file", + str(ws), + str(ws), + sandbox_rw_binds=[str(cache_dir)], + ) + tokens = _parse(result) + + bind_try_indices = [i for i, t in enumerate(tokens) if t == "--bind-try"] + bind_try_pairs = {(tokens[i + 1], tokens[i + 2]) for i in bind_try_indices} + resolved = str(cache_dir.resolve(strict=False)) + assert (resolved, resolved) in bind_try_pairs + + def test_custom_relative_bind_paths_are_ignored(self, tmp_path): + ws = tmp_path / "project" + + result = wrap_command( + "bwrap", + "ls", + str(ws), + str(ws), + sandbox_ro_binds=["relative/bin"], + sandbox_rw_binds=["relative/cache"], + ) + tokens = _parse(result) + + assert "relative/bin" not in tokens + assert "relative/cache" not in tokens + class TestUnknownBackend: def test_raises_value_error(self, tmp_path): diff --git a/tests/tools/test_tool_validation.py b/tests/tools/test_tool_validation.py index e05c7de49..8ffd83e47 100644 --- a/tests/tools/test_tool_validation.py +++ b/tests/tools/test_tool_validation.py @@ -714,6 +714,22 @@ def test_exec_config_timeout_uncapped_and_zero() -> None: ExecToolConfig(timeout=-1) +def test_exec_config_accepts_bwrap_bind_aliases() -> None: + cfg = ExecToolConfig.model_validate( + { + "sandboxRoBinds": ["/home/user/.local/bin"], + "sandboxRwBinds": ["/home/user/.cache/uv"], + } + ) + + dumped = cfg.model_dump(by_alias=True) + + assert cfg.sandbox_ro_binds == ["/home/user/.local/bin"] + assert cfg.sandbox_rw_binds == ["/home/user/.cache/uv"] + assert dumped["sandboxRoBinds"] == ["/home/user/.local/bin"] + assert dumped["sandboxRwBinds"] == ["/home/user/.cache/uv"] + + def test_resolve_timeout_config_uncapped_and_unlimited() -> None: """Config timeout drives the hard timeout uncapped; 0 means no limit (#3595).""" assert ExecTool(timeout=3600)._resolve_timeout(None) == 3600 From 22e61003f90c03295cbc86b601d5dfece425dc03 Mon Sep 17 00:00:00 2001 From: yu-xin-c <2182712990@qq.com> Date: Fri, 19 Jun 2026 02:00:15 +0800 Subject: [PATCH 14/48] test(exec): make bwrap bind tests portable --- tests/tools/test_exec_platform.py | 13 +++++++------ tests/tools/test_exec_security.py | 6 ++++-- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/tests/tools/test_exec_platform.py b/tests/tools/test_exec_platform.py index 4a6277b24..9981e2200 100644 --- a/tests/tools/test_exec_platform.py +++ b/tests/tools/test_exec_platform.py @@ -8,7 +8,6 @@ platform-specific binaries (all subprocess calls are mocked). import asyncio import shutil import sys -from pathlib import Path from unittest.mock import AsyncMock, patch import pytest @@ -475,11 +474,13 @@ class TestSandboxPlatform: assert "bwrap" in spawned_cmd @pytest.mark.asyncio - async def test_bwrap_receives_configured_bind_roots(self): + async def test_bwrap_receives_configured_bind_roots(self, tmp_path): """Configured bwrap bind roots should be forwarded to the sandbox wrapper.""" mock_proc = AsyncMock() mock_proc.communicate.return_value = (b"sandboxed", b"") mock_proc.returncode = 0 + tool_bin = tmp_path / "tool-bin" + tool_cache = tmp_path / "tool-cache" with ( patch("nanobot.agent.tools.shell._IS_WINDOWS", False), @@ -490,17 +491,17 @@ class TestSandboxPlatform: tool = ExecTool( sandbox="bwrap", working_dir="/workspace", - sandbox_ro_binds=["/home/user/.local/bin"], - sandbox_rw_binds=["/home/user/.cache/uv"], + sandbox_ro_binds=[str(tool_bin)], + sandbox_rw_binds=[str(tool_cache)], ) await tool.execute(command="ls") kwargs = mock_wrap.call_args.kwargs assert kwargs["sandbox_ro_binds"] == [ - str(Path("/home/user/.local/bin").resolve(strict=False)) + str(tool_bin.resolve(strict=False)) ] assert kwargs["sandbox_rw_binds"] == [ - str(Path("/home/user/.cache/uv").resolve(strict=False)) + str(tool_cache.resolve(strict=False)) ] diff --git a/tests/tools/test_exec_security.py b/tests/tools/test_exec_security.py index 6e3a6f8a8..bc6d8a956 100644 --- a/tests/tools/test_exec_security.py +++ b/tests/tools/test_exec_security.py @@ -314,13 +314,14 @@ def test_exec_still_blocks_real_outside_path_via_redirect(tmp_path): assert "path outside working dir" in blocked -def test_exec_allows_absolute_path_inside_bwrap_ro_bind(tmp_path): +def test_exec_allows_absolute_path_inside_bwrap_ro_bind(tmp_path, monkeypatch): workspace = tmp_path / "workspace" workspace.mkdir() tool_bin = tmp_path / "home" / ".local" / "bin" tool_bin.mkdir(parents=True) uv = tool_bin / "uv" uv.write_text("#!/bin/sh\n") + monkeypatch.setattr("nanobot.agent.tools.shell._IS_WINDOWS", False) tool = ExecTool( working_dir=str(workspace), restrict_to_workspace=True, @@ -338,11 +339,12 @@ def test_exec_allows_absolute_path_inside_bwrap_ro_bind(tmp_path): assert blocked is None -def test_exec_allows_absolute_path_inside_bwrap_rw_bind(tmp_path): +def test_exec_allows_absolute_path_inside_bwrap_rw_bind(tmp_path, monkeypatch): workspace = tmp_path / "workspace" workspace.mkdir() cache_dir = tmp_path / "cache" cache_dir.mkdir() + monkeypatch.setattr("nanobot.agent.tools.shell._IS_WINDOWS", False) tool = ExecTool( working_dir=str(workspace), restrict_to_workspace=True, From cf6ca13b6d0be16df8c5e017286c0861ca6f9a0e Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Mon, 27 Jul 2026 00:23:53 +0800 Subject: [PATCH 15/48] fix(exec): preserve bwrap workspace masking --- docs/configuration.md | 4 ++-- nanobot/agent/tools/sandbox.py | 22 ++++++++++++++++++---- nanobot/agent/tools/shell.py | 18 +++++++++++++++--- tests/tools/test_exec_security.py | 24 ++++++++++++++++++++++++ tests/tools/test_sandbox.py | 21 +++++++++++++++++++++ 5 files changed, 80 insertions(+), 9 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index 3d4be10d9..0c94d0108 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -1995,8 +1995,8 @@ For API keys, tokens, and other secrets, see [Environment Variables for Secrets] | `tools.exec.timeout` | `60` | Default hard timeout in seconds for shell commands. Config values may exceed the per-call tool cap; set `0` to disable the hard timeout for trusted long-running commands. | | `tools.exec.pathPrepend` | `""` | Extra directories to prepend to `PATH` when running shell commands. Use this when configured tools should win executable lookup precedence, such as a Python virtual environment's `bin` or `Scripts` directory. | | `tools.exec.pathAppend` | `""` | Extra directories to append to `PATH` when running shell commands (e.g. `/usr/sbin` for `ufw`). | -| `tools.exec.sandboxRoBinds` | `[]` | Extra absolute paths to read-only bind into the `"bwrap"` sandbox with `--ro-bind-try`, such as `/home/user/.local/bin` or `/home/user/.cargo/bin` when those paths are also in `pathPrepend`/`pathAppend`. These roots are also accepted by the shell absolute-path guard only while bwrap is active. | -| `tools.exec.sandboxRwBinds` | `[]` | Extra absolute paths to read-write bind into the `"bwrap"` sandbox with `--bind-try`, for trusted tool caches or scratch directories. Use sparingly: paths listed here are intentionally writable by shell commands inside the sandbox. | +| `tools.exec.sandboxRoBinds` | `[]` | Extra absolute paths to read-only bind into the `"bwrap"` sandbox with `--ro-bind-try`, such as `/home/user/.local/bin` or `/home/user/.cargo/bin` when those paths are also in `pathPrepend`/`pathAppend`. These roots are also accepted by the shell absolute-path guard only while bwrap is active. Bind only directories whose contents are safe for agent commands to read; paths equal to or containing the active workspace are ignored so they cannot uncover its masked parent directory. | +| `tools.exec.sandboxRwBinds` | `[]` | Extra absolute paths to read-write bind into the `"bwrap"` sandbox with `--bind-try`, for trusted tool caches or scratch directories. Use sparingly: paths listed here are intentionally writable by shell commands inside the sandbox. Paths equal to or containing the active workspace are ignored. | | `tools.webuiAllowRemotePackageInstall` | `false` | When `false`, the WebUI can install missing optional packages only from a browser opened on the same machine as nanobot. Set to `true` only when a trusted remote admin is allowed to install Python packages into this nanobot environment. | | `tools.ssrfWhitelist` | `[]` | CIDR ranges exempted from the shared SSRF guard used by web fetches and HTTP/SSE MCP connections. Prefer exact host CIDRs such as `192.168.1.50/32`; broad ranges increase SSRF exposure. | | `channels.*.allowFrom` | omitted | Access control per channel. Omit to use pairing-only mode; set `["*"]` to allow everyone; or list specific user IDs. See [Pairing](#pairing) for details. | diff --git a/nanobot/agent/tools/sandbox.py b/nanobot/agent/tools/sandbox.py index 0d602abfb..4c5cbf853 100644 --- a/nanobot/agent/tools/sandbox.py +++ b/nanobot/agent/tools/sandbox.py @@ -13,7 +13,11 @@ from typing import Iterable from nanobot.config.paths import get_media_dir -def _normalize_bind_paths(paths: Iterable[str] | None) -> list[str]: +def _normalize_bind_paths( + paths: Iterable[str] | None, + *, + workspace: Path | None = None, +) -> list[str]: out: list[str] = [] seen: set[str] = set() for raw in paths or []: @@ -23,7 +27,17 @@ def _normalize_bind_paths(paths: Iterable[str] | None) -> list[str]: path = Path(os.path.expandvars(value)).expanduser() if not path.is_absolute(): continue - resolved = str(path.resolve(strict=False)) + resolved_path = path.resolve(strict=False) + if workspace is not None: + try: + workspace.relative_to(resolved_path) + except ValueError: + pass + else: + # A later bind of the workspace or one of its parents could + # cover the tmpfs that hides the config directory. + continue + resolved = str(resolved_path) if resolved in seen: continue seen.add(resolved) @@ -79,9 +93,9 @@ def _bwrap( "--bind", str(ws), str(ws), "--ro-bind-try", str(media), str(media), # read-only access to media ] - for p in _normalize_bind_paths(sandbox_ro_binds): + for p in _normalize_bind_paths(sandbox_ro_binds, workspace=ws): args += ["--ro-bind-try", p, p] - for p in _normalize_bind_paths(sandbox_rw_binds): + for p in _normalize_bind_paths(sandbox_rw_binds, workspace=ws): args += ["--bind-try", p, p] args += ["--chdir", sandbox_cwd, "--", "sh", "-c", command] return shlex.join(args) diff --git a/nanobot/agent/tools/shell.py b/nanobot/agent/tools/shell.py index be3e56056..71b369c83 100644 --- a/nanobot/agent/tools/shell.py +++ b/nanobot/agent/tools/shell.py @@ -809,7 +809,9 @@ class ExecTool(Tool): if workspace_root else None ) - sandbox_bind_roots = self._active_sandbox_bind_roots() + sandbox_bind_roots = self._active_sandbox_bind_roots( + resolved_workspace or cwd_path + ) for raw in self._extract_absolute_paths(cmd): try: @@ -960,7 +962,17 @@ class ExecTool(Tool): roots.append(resolved) return roots - def _active_sandbox_bind_roots(self) -> list[Path]: + def _active_sandbox_bind_roots( + self, + workspace_root: Path | None = None, + ) -> list[Path]: if self.sandbox != "bwrap" or _IS_WINDOWS: return [] - return [*self.sandbox_ro_binds, *self.sandbox_rw_binds] + roots = [*self.sandbox_ro_binds, *self.sandbox_rw_binds] + if workspace_root is None: + return roots + return [ + root + for root in roots + if not is_path_within(workspace_root, root) + ] diff --git a/tests/tools/test_exec_security.py b/tests/tools/test_exec_security.py index bc6d8a956..508cd1e64 100644 --- a/tests/tools/test_exec_security.py +++ b/tests/tools/test_exec_security.py @@ -387,6 +387,30 @@ def test_exec_bind_roots_do_not_widen_guard_without_bwrap(tmp_path): assert "path outside working dir" in blocked +def test_exec_bwrap_bind_parent_does_not_widen_workspace_guard(tmp_path, monkeypatch): + workspace = tmp_path / "workspace" + workspace.mkdir() + secret = tmp_path / "config.json" + secret.write_text("secret") + monkeypatch.setattr("nanobot.agent.tools.shell._IS_WINDOWS", False) + tool = ExecTool( + working_dir=str(workspace), + restrict_to_workspace=True, + sandbox="bwrap", + sandbox_ro_binds=[str(tmp_path)], + ) + + blocked = tool._guard_command( + f"cat {secret}", + str(workspace), + restrict_to_workspace=True, + workspace_root=str(workspace), + ) + + assert blocked is not None + assert "path outside working dir" in blocked + + # --- format command blocking ----------------------------------------------- diff --git a/tests/tools/test_sandbox.py b/tests/tools/test_sandbox.py index 752713517..c83ef4220 100644 --- a/tests/tools/test_sandbox.py +++ b/tests/tools/test_sandbox.py @@ -201,6 +201,27 @@ class TestBwrapBackend: assert "relative/bin" not in tokens assert "relative/cache" not in tokens + def test_custom_workspace_parent_binds_are_ignored(self, tmp_path): + ws = tmp_path / "private" / "project" + parent = ws.parent.resolve(strict=False) + + result = wrap_command( + "bwrap", + "cat ../config.json", + str(ws), + str(ws), + sandbox_ro_binds=[str(parent)], + sandbox_rw_binds=[str(parent)], + ) + tokens = _parse(result) + + ro_try_indices = [i for i, token in enumerate(tokens) if token == "--ro-bind-try"] + ro_try_pairs = {(tokens[i + 1], tokens[i + 2]) for i in ro_try_indices} + bind_try_indices = [i for i, token in enumerate(tokens) if token == "--bind-try"] + bind_try_pairs = {(tokens[i + 1], tokens[i + 2]) for i in bind_try_indices} + assert (str(parent), str(parent)) not in ro_try_pairs + assert (str(parent), str(parent)) not in bind_try_pairs + class TestUnknownBackend: def test_raises_value_error(self, tmp_path): From f7bf4c972ebce2c362186a1a9f3ab43a4fa06a2a Mon Sep 17 00:00:00 2001 From: santhreal <64453045+santhreal@users.noreply.github.com> Date: Sat, 25 Jul 2026 21:47:05 -0700 Subject: [PATCH 16/48] fix(pairing): treat null approved/pending maps as empty --- nanobot/pairing/store.py | 35 ++++++++++++++++++++++++++++++----- tests/pairing/test_store.py | 24 ++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 5 deletions(-) diff --git a/nanobot/pairing/store.py b/nanobot/pairing/store.py index afafb1ecd..c26f4e0ba 100644 --- a/nanobot/pairing/store.py +++ b/nanobot/pairing/store.py @@ -44,8 +44,18 @@ def _load() -> dict[str, Any]: logger.warning("Corrupted pairing store, resetting") return {"approved": {}, "pending": {}} + # JSON stores may contain null maps after partial edits; treat like {}. + approved = data.get("approved") or {} + if not isinstance(approved, dict): + approved = {} + data["approved"] = approved + pending = data.get("pending") or {} + if not isinstance(pending, dict): + pending = {} + data["pending"] = pending + # Convert approved lists to str sets for O(1) lookup. - for channel, users in data.get("approved", {}).items(): + for channel, users in approved.items(): if not isinstance(users, list): users = [] data["approved"][channel] = {str(u) for u in users} @@ -56,9 +66,15 @@ def _save(data: dict[str, Any]) -> None: path = _store_path() path.parent.mkdir(parents=True, exist_ok=True) # Convert sets back to lists for JSON serialization + approved = data.get("approved") or {} + pending = data.get("pending") or {} + if not isinstance(approved, dict): + approved = {} + if not isinstance(pending, dict): + pending = {} payload = { - "approved": {ch: sorted(list(users)) for ch, users in data.get("approved", {}).items()}, - "pending": dict(data.get("pending", {})), + "approved": {ch: sorted(list(users)) for ch, users in approved.items()}, + "pending": dict(pending), } _write_text_atomic(path, json.dumps(payload, indent=2, ensure_ascii=False)) @@ -66,10 +82,18 @@ def _save(data: dict[str, Any]) -> None: def _gc_pending(data: dict[str, Any]) -> None: """Remove expired pending entries in-place.""" now = time.time() - pending: dict[str, Any] = data.get("pending", {}) - expired = [code for code, info in pending.items() if info.get("expires_at", 0) < now] + pending: dict[str, Any] = data.get("pending") or {} + if not isinstance(pending, dict): + data["pending"] = {} + return + expired = [ + code + for code, info in pending.items() + if not isinstance(info, dict) or info.get("expires_at", 0) < now + ] for code in expired: del pending[code] + data["pending"] = pending def generate_code( @@ -152,6 +176,7 @@ def list_pending() -> list[dict[str, Any]]: return [ {"code": code, **info} for code, info in data.get("pending", {}).items() + if isinstance(info, dict) ] diff --git a/tests/pairing/test_store.py b/tests/pairing/test_store.py index 56c84f5d8..766657f65 100644 --- a/tests/pairing/test_store.py +++ b/tests/pairing/test_store.py @@ -257,3 +257,27 @@ def test_load_treats_null_approved_channel_list_as_empty(tmp_path, monkeypatch): assert store.is_approved("telegram", "123") is False assert store.is_approved("discord", "456") is True assert store.get_approved("telegram") == [] + + +def test_load_treats_null_approved_and_pending_maps_as_empty(tmp_path, monkeypatch): + """Top-level approved/pending null must not crash pairing load or list_pending.""" + path = tmp_path / "pairing.json" + path.write_text( + '{"approved": null, "pending": null}', + encoding="utf-8", + ) + monkeypatch.setattr(store, "_store_path", lambda: path) + assert store.is_approved("telegram", "123") is False + assert store.list_pending() == [] + assert store.get_approved("telegram") == [] + + +def test_list_pending_skips_null_pending_entries(tmp_path, monkeypatch): + """Null pending entry values must be dropped instead of crashing list_pending.""" + path = tmp_path / "pairing.json" + path.write_text( + '{"approved": {}, "pending": {"ABCD-EFGH": null}}', + encoding="utf-8", + ) + monkeypatch.setattr(store, "_store_path", lambda: path) + assert store.list_pending() == [] From d236883e2d304c94a887ce6de6d6b721e7213164 Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Sun, 26 Jul 2026 17:21:56 +0800 Subject: [PATCH 17/48] fix(pairing): reject malformed store entries --- nanobot/pairing/store.py | 14 +++++++++++++- tests/pairing/test_store.py | 23 +++++++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/nanobot/pairing/store.py b/nanobot/pairing/store.py index c26f4e0ba..490369cbc 100644 --- a/nanobot/pairing/store.py +++ b/nanobot/pairing/store.py @@ -43,6 +43,9 @@ def _load() -> dict[str, Any]: except (json.JSONDecodeError, OSError): logger.warning("Corrupted pairing store, resetting") return {"approved": {}, "pending": {}} + if not isinstance(data, dict): + logger.warning("Corrupted pairing store, resetting") + return {"approved": {}, "pending": {}} # JSON stores may contain null maps after partial edits; treat like {}. approved = data.get("approved") or {} @@ -89,7 +92,15 @@ def _gc_pending(data: dict[str, Any]) -> None: expired = [ code for code, info in pending.items() - if not isinstance(info, dict) or info.get("expires_at", 0) < now + if ( + not isinstance(info, dict) + or not isinstance(info.get("channel"), str) + or not info.get("channel") + or info.get("sender_id") is None + or isinstance(info.get("expires_at"), bool) + or not isinstance(info.get("expires_at"), (int, float)) + or info["expires_at"] < now + ) ] for code in expired: del pending[code] @@ -220,6 +231,7 @@ def clear_channel(channel: str) -> dict[str, int]: """Remove approved senders and pending requests for *channel*.""" with _LOCK: data = _load() + _gc_pending(data) approved: dict[str, set[str]] = data.get("approved", {}) approved_users = approved.pop(channel, set()) diff --git a/tests/pairing/test_store.py b/tests/pairing/test_store.py index 766657f65..c4b4758af 100644 --- a/tests/pairing/test_store.py +++ b/tests/pairing/test_store.py @@ -272,6 +272,15 @@ def test_load_treats_null_approved_and_pending_maps_as_empty(tmp_path, monkeypat assert store.get_approved("telegram") == [] +@pytest.mark.parametrize("payload", ["null", "[]", "true"]) +def test_load_treats_non_object_store_as_empty(tmp_path, monkeypatch, payload): + path = tmp_path / "pairing.json" + path.write_text(payload, encoding="utf-8") + monkeypatch.setattr(store, "_store_path", lambda: path) + assert store.list_pending() == [] + assert store.is_approved("telegram", "123") is False + + def test_list_pending_skips_null_pending_entries(tmp_path, monkeypatch): """Null pending entry values must be dropped instead of crashing list_pending.""" path = tmp_path / "pairing.json" @@ -281,3 +290,17 @@ def test_list_pending_skips_null_pending_entries(tmp_path, monkeypatch): ) monkeypatch.setattr(store, "_store_path", lambda: path) assert store.list_pending() == [] + assert store.clear_channel("telegram") == {"approved": 0, "pending": 0} + + +def test_pending_gc_drops_malformed_entries(tmp_path, monkeypatch): + path = tmp_path / "pairing.json" + path.write_text( + '{"approved": {}, "pending": {' + '"bad-expiry": {"channel": "telegram", "sender_id": "123", "expires_at": null},' + '"missing-sender": {"channel": "telegram", "expires_at": 9999999999}' + "}}", + encoding="utf-8", + ) + monkeypatch.setattr(store, "_store_path", lambda: path) + assert store.list_pending() == [] From 48358147468be32f9296d9b4c3af0a323197b620 Mon Sep 17 00:00:00 2001 From: KDB <937925477@qq.com> Date: Fri, 24 Jul 2026 00:22:23 +0800 Subject: [PATCH 18/48] fix(channels): ignore confirmations after connect cancellation --- nanobot/channels/feishu/connect.py | 10 ++- nanobot/channels/feishu/tests/test_connect.py | 65 +++++++++++++++++++ nanobot/channels/weixin/connect.py | 6 ++ nanobot/channels/weixin/tests/test_connect.py | 50 ++++++++++++++ 4 files changed, 130 insertions(+), 1 deletion(-) create mode 100644 nanobot/channels/feishu/tests/test_connect.py diff --git a/nanobot/channels/feishu/connect.py b/nanobot/channels/feishu/connect.py index 12930e11d..795a4a787 100644 --- a/nanobot/channels/feishu/connect.py +++ b/nanobot/channels/feishu/connect.py @@ -127,9 +127,16 @@ class FeishuConnectStore: session.last_error = str(exc) return _pending_payload(session) - session.domain = str(result.get("domain") or session.domain) status = result.get("status") if status == "succeeded": + if self._sessions.get(session_id) is not session: + return { + "session_id": session_id, + "instance_id": session.instance_id, + "status": "cancelled", + "message": "Feishu connection cancelled.", + } + session.domain = str(result.get("domain") or session.domain) session.instance_id = feishu.save_registration_result( result, instance_id=session.instance_id, @@ -145,6 +152,7 @@ class FeishuConnectStore: "app_id": result.get("app_id"), } + session.domain = str(result.get("domain") or session.domain) if status == "failed": self._sessions.pop(session_id, None) return { diff --git a/nanobot/channels/feishu/tests/test_connect.py b/nanobot/channels/feishu/tests/test_connect.py new file mode 100644 index 000000000..278580c1b --- /dev/null +++ b/nanobot/channels/feishu/tests/test_connect.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +import asyncio +import threading +from typing import Any + +import pytest + +from nanobot.channels.feishu import runtime as feishu +from nanobot.channels.feishu.connect import FeishuConnectStore + + +@pytest.mark.asyncio +async def test_feishu_cancel_wins_over_inflight_confirmation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + poll_started = threading.Event() + release_poll = threading.Event() + saved_results: list[dict[str, Any]] = [] + + monkeypatch.setattr(feishu, "_init_registration", lambda _domain: None) + monkeypatch.setattr( + feishu, + "_begin_registration", + lambda _domain: { + "device_code": "device-cancel", + "qr_url": "https://qr.example/cancel", + "expire_in": 600, + "interval": 2, + }, + ) + + def fake_poll_registration_once(**_kwargs: Any) -> dict[str, str]: + poll_started.set() + assert release_poll.wait(timeout=5) + return { + "status": "succeeded", + "domain": "feishu", + "app_id": "late-app", + "app_secret": "late-secret", + } + + def fake_save_registration_result( + result: dict[str, Any], + **_kwargs: Any, + ) -> str: + saved_results.append(result) + return "default" + + monkeypatch.setattr(feishu, "poll_registration_once", fake_poll_registration_once) + monkeypatch.setattr(feishu, "save_registration_result", fake_save_registration_result) + + store = FeishuConnectStore() + started = await store.handle("start", {}) + query = {"session_id": [started["session_id"]]} + poll_task = asyncio.create_task(store.handle("poll", query)) + assert await asyncio.to_thread(poll_started.wait, 5) + + cancelled = await store.handle("cancel", query) + release_poll.set() + completed = await poll_task + + assert cancelled["status"] == "cancelled" + assert completed["status"] == "cancelled" + assert saved_results == [] diff --git a/nanobot/channels/weixin/connect.py b/nanobot/channels/weixin/connect.py index a56a7a2f9..a254b64d3 100644 --- a/nanobot/channels/weixin/connect.py +++ b/nanobot/channels/weixin/connect.py @@ -130,6 +130,12 @@ class WeixinConnectStore: status = status_data.get("status", "") if status == "confirmed": + if self._sessions.get(session_id) is not session: + return { + "session_id": session_id, + "status": "cancelled", + "message": "WeChat login cancelled.", + } token = str(status_data.get("bot_token", "") or "") if not token: self._sessions.pop(session_id, None) diff --git a/nanobot/channels/weixin/tests/test_connect.py b/nanobot/channels/weixin/tests/test_connect.py index 47e4a8257..e201ce04e 100644 --- a/nanobot/channels/weixin/tests/test_connect.py +++ b/nanobot/channels/weixin/tests/test_connect.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import json from typing import Any @@ -97,3 +98,52 @@ async def test_weixin_reconnect_keeps_existing_account_until_scan_succeeds( cancelled = await store.cancel(started["session_id"]) assert cancelled["status"] == "cancelled" assert json.loads(state_file.read_text(encoding="utf-8")) == existing + + +@pytest.mark.asyncio +async def test_weixin_cancel_wins_over_inflight_confirmation( + tmp_path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + state_dir = tmp_path / "weixin-state" + config_path = tmp_path / "config.json" + save_config( + Config.model_validate({"channels": {"weixin": {"stateDir": str(state_dir)}}}), + config_path, + ) + monkeypatch.setattr("nanobot.config.loader._current_config_path", config_path) + + poll_started = asyncio.Event() + release_poll = asyncio.Event() + + async def fake_fetch_qr_code(self: WeixinChannel) -> tuple[str, str]: + return "qr-cancel", "https://qr.example/cancel" + + async def fake_api_get_with_base( + self: WeixinChannel, + **_kwargs: Any, + ) -> dict[str, str]: + poll_started.set() + await release_poll.wait() + return { + "status": "confirmed", + "bot_token": "late-token", + "ilink_user_id": "late-user", + } + + monkeypatch.setattr(WeixinChannel, "_fetch_qr_code", fake_fetch_qr_code) + monkeypatch.setattr(WeixinChannel, "_api_get_with_base", fake_api_get_with_base) + + store = WeixinConnectStore() + started = await store.handle("start", {}) + query = {"session_id": [started["session_id"]]} + poll_task = asyncio.create_task(store.handle("poll", query)) + await asyncio.wait_for(poll_started.wait(), timeout=5) + + cancelled = await store.handle("cancel", query) + release_poll.set() + completed = await poll_task + + assert cancelled["status"] == "cancelled" + assert completed["status"] == "cancelled" + assert not (state_dir / "account.json").exists() From 2e2f15dd0c120693881cbad94b8ebf10c3410013 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Fri, 24 Jul 2026 09:53:35 +0800 Subject: [PATCH 19/48] fix(channels): serialize Feishu connect completion --- nanobot/channels/feishu/connect.py | 44 +++++++------- nanobot/channels/feishu/tests/test_connect.py | 57 +++++++++++++++++++ 2 files changed, 81 insertions(+), 20 deletions(-) diff --git a/nanobot/channels/feishu/connect.py b/nanobot/channels/feishu/connect.py index 795a4a787..3258d1550 100644 --- a/nanobot/channels/feishu/connect.py +++ b/nanobot/channels/feishu/connect.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio import json import secrets +import threading import time from dataclasses import dataclass from typing import Any @@ -41,6 +42,7 @@ class FeishuConnectStore: def __init__(self) -> None: self._sessions: dict[str, FeishuConnectSession] = {} + self._completion_lock = threading.Lock() async def handle(self, action: str, query: QueryParams) -> dict[str, Any]: """Handle one generic settings connection action.""" @@ -58,7 +60,7 @@ class FeishuConnectStore: if action == "poll": return await asyncio.to_thread(self.poll, session_id) if action == "cancel": - return self.cancel(session_id) + return await asyncio.to_thread(self.cancel, session_id) raise ChannelConnectError(f"unsupported Feishu connect action: {action}", status=404) def start( @@ -129,28 +131,29 @@ class FeishuConnectStore: status = result.get("status") if status == "succeeded": - if self._sessions.get(session_id) is not session: + with self._completion_lock: + if self._sessions.get(session_id) is not session: + return { + "session_id": session_id, + "instance_id": session.instance_id, + "status": "cancelled", + "message": "Feishu connection cancelled.", + } + session.domain = str(result.get("domain") or session.domain) + session.instance_id = feishu.save_registration_result( + result, + instance_id=session.instance_id, + name=session.instance_name, + ) + self._sessions.pop(session_id, None) return { "session_id": session_id, "instance_id": session.instance_id, - "status": "cancelled", - "message": "Feishu connection cancelled.", + "status": "succeeded", + "message": "Feishu is connected.", + "domain": session.domain, + "app_id": result.get("app_id"), } - session.domain = str(result.get("domain") or session.domain) - session.instance_id = feishu.save_registration_result( - result, - instance_id=session.instance_id, - name=session.instance_name, - ) - self._sessions.pop(session_id, None) - return { - "session_id": session_id, - "instance_id": session.instance_id, - "status": "succeeded", - "message": "Feishu is connected.", - "domain": session.domain, - "app_id": result.get("app_id"), - } session.domain = str(result.get("domain") or session.domain) if status == "failed": @@ -166,7 +169,8 @@ class FeishuConnectStore: return _pending_payload(session) def cancel(self, session_id: str) -> dict[str, Any]: - session = self._sessions.pop(session_id, None) + with self._completion_lock: + session = self._sessions.pop(session_id, None) return { "session_id": session_id, "instance_id": session.instance_id if session else DEFAULT_INSTANCE_ID, diff --git a/nanobot/channels/feishu/tests/test_connect.py b/nanobot/channels/feishu/tests/test_connect.py index 278580c1b..433a8db0b 100644 --- a/nanobot/channels/feishu/tests/test_connect.py +++ b/nanobot/channels/feishu/tests/test_connect.py @@ -63,3 +63,60 @@ async def test_feishu_cancel_wins_over_inflight_confirmation( assert cancelled["status"] == "cancelled" assert completed["status"] == "cancelled" assert saved_results == [] + + +@pytest.mark.asyncio +async def test_feishu_cancel_does_not_interleave_with_registration_save( + monkeypatch: pytest.MonkeyPatch, +) -> None: + save_started = threading.Event() + release_save = threading.Event() + + monkeypatch.setattr(feishu, "_init_registration", lambda _domain: None) + monkeypatch.setattr( + feishu, + "_begin_registration", + lambda _domain: { + "device_code": "device-lock", + "qr_url": "https://qr.example/lock", + "expire_in": 600, + "interval": 2, + }, + ) + monkeypatch.setattr( + feishu, + "poll_registration_once", + lambda **_kwargs: { + "status": "succeeded", + "domain": "feishu", + "app_id": "saved-app", + "app_secret": "saved-secret", + }, + ) + + def fake_save_registration_result( + _result: dict[str, Any], + **_kwargs: Any, + ) -> str: + save_started.set() + assert release_save.wait(timeout=5) + return "default" + + monkeypatch.setattr(feishu, "save_registration_result", fake_save_registration_result) + + store = FeishuConnectStore() + started = await store.handle("start", {}) + query = {"session_id": [started["session_id"]]} + poll_task = asyncio.create_task(store.handle("poll", query)) + assert await asyncio.to_thread(save_started.wait, 5) + + cancel_task = asyncio.create_task(store.handle("cancel", query)) + await asyncio.sleep(0) + assert not cancel_task.done() + + release_save.set() + completed = await poll_task + cancelled = await cancel_task + + assert completed["status"] == "succeeded" + assert cancelled["status"] == "cancelled" From 9aae7485d6a7126d1e8bc4eff2a4b3e472991a76 Mon Sep 17 00:00:00 2001 From: amplifierplus <160200579+amplifierplus@users.noreply.github.com> Date: Thu, 23 Jul 2026 17:01:46 +0800 Subject: [PATCH 20/48] fix(mcp): normalize local schema refs --- nanobot/agent/tools/mcp.py | 104 ++++++++++++++++++++++++++++++----- tests/tools/test_mcp_tool.py | 73 ++++++++++++++++++++++++ 2 files changed, 162 insertions(+), 15 deletions(-) diff --git a/nanobot/agent/tools/mcp.py b/nanobot/agent/tools/mcp.py index 3cf1147c6..7902e66bf 100644 --- a/nanobot/agent/tools/mcp.py +++ b/nanobot/agent/tools/mcp.py @@ -315,13 +315,76 @@ def _extract_nullable_branch(options: Any) -> tuple[dict[str, Any], bool] | None return None -def _normalize_schema_for_openai(schema: Any) -> dict[str, Any]: - """Normalize only nullable JSON Schema patterns for tool definitions.""" - if not isinstance(schema, dict): - return {"type": "object", "properties": {}} +def _resolve_local_schema_ref(root: dict[str, Any], ref: str) -> Any: + """Resolve a local JSON Pointer without accepting remote references.""" + if ref == "#": + return root + if not ref.startswith("#/"): + raise ValueError("not a local JSON Pointer") + current: Any = root + for raw_part in ref[2:].split("/"): + part = raw_part.replace("~1", "/").replace("~0", "~") + if isinstance(current, dict): + current = current[part] + elif isinstance(current, list): + current = current[int(part)] + else: + raise KeyError(part) + return current + + +def _rewrite_local_schema_refs(schema: dict[str, Any]) -> dict[str, Any]: + """Hoist arbitrary local JSON-Pointer refs into provider-compatible ``$defs``.""" + rewritten_refs: dict[str, str] = {} + generated_defs: dict[str, Any] = {} + + def rewrite(value: Any) -> Any: + if isinstance(value, list): + return [rewrite(item) for item in value] + if not isinstance(value, dict): + return value + + rewritten = dict(value) + ref = rewritten.get("$ref") + is_rewritable_ref = isinstance(ref, str) and ( + ref == "#" or (ref.startswith("#/") and not ref.startswith("#/$defs/")) + ) + if is_rewritable_ref: + name = rewritten_refs.get(ref) + if name is None: + try: + target = _resolve_local_schema_ref(schema, ref) + except (KeyError, IndexError, TypeError, ValueError): + logger.warning("MCP tool schema contains an unresolved local $ref: {}", ref) + else: + assert isinstance(ref, str) + name = f"ref_{hashlib.sha256(ref.encode()).hexdigest()[:12]}" + existing_defs = schema.get("$defs") + while isinstance(existing_defs, dict) and name in existing_defs: + name += "_" + rewritten_refs[ref] = name + # Reserve the name before descending so recursive refs terminate. + generated_defs[name] = {} + generated_defs[name] = rewrite(target) + if name is not None: + rewritten["$ref"] = f"#/$defs/{name}" + + return {key: rewrite(item) for key, item in rewritten.items()} + + result = rewrite(schema) + if generated_defs: + existing_defs = result.get("$defs") + result["$defs"] = { + **(existing_defs if isinstance(existing_defs, dict) else {}), + **generated_defs, + } + return result + + +def _normalize_nullable_schema(schema: dict[str, Any]) -> dict[str, Any]: + """Normalize nullable forms in structural subschemas only.""" normalized = dict(schema) - raw_type = normalized.get("type") if isinstance(raw_type, list): non_null = [item for item in raw_type if item != "null"] @@ -339,23 +402,34 @@ def _normalize_schema_for_openai(schema: Any) -> dict[str, Any]: normalized["nullable"] = True break - if "properties" in normalized and isinstance(normalized["properties"], dict): + if isinstance(normalized.get("properties"), dict): normalized["properties"] = { - name: _normalize_schema_for_openai(prop) if isinstance(prop, dict) else prop + name: _normalize_nullable_schema(prop) if isinstance(prop, dict) else prop for name, prop in normalized["properties"].items() } + if isinstance(normalized.get("items"), dict): + normalized["items"] = _normalize_nullable_schema(normalized["items"]) + if isinstance(normalized.get("$defs"), dict): + normalized["$defs"] = { + name: _normalize_nullable_schema(definition) + if isinstance(definition, dict) + else definition + for name, definition in normalized["$defs"].items() + } - if "items" in normalized and isinstance(normalized["items"], dict): - normalized["items"] = _normalize_schema_for_openai(normalized["items"]) - - if normalized.get("type") != "object": - return normalized - - normalized.setdefault("properties", {}) - normalized.setdefault("required", []) + if normalized.get("type") == "object": + normalized.setdefault("properties", {}) + normalized.setdefault("required", []) return normalized +def _normalize_schema_for_openai(schema: Any) -> dict[str, Any]: + """Normalize MCP JSON Schema patterns for tool definitions.""" + if not isinstance(schema, dict): + return {"type": "object", "properties": {}} + return _normalize_nullable_schema(_rewrite_local_schema_refs(schema)) + + class _MCPWrapperBase(Tool): """Common reconnect handling for wrappers bound to one MCP server session.""" diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py index 8bf932a75..233c51fad 100644 --- a/tests/tools/test_mcp_tool.py +++ b/tests/tools/test_mcp_tool.py @@ -236,6 +236,79 @@ def test_wrapper_normalizes_nullable_property_anyof() -> None: } +def test_wrapper_hoists_recursive_local_refs_into_defs() -> None: + recursive_items_ref = "#/properties/filter/properties/items" + tool_def = SimpleNamespace( + name="search_dataset", + description="search tool", + inputSchema={ + "type": "object", + "properties": { + "filter": { + "type": "object", + "properties": { + "items": { + "type": "array", + "items": {"$ref": recursive_items_ref}, + } + }, + "required": ["items"], + } + }, + }, + ) + + wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def) + + generated_ref = wrapper.parameters["properties"]["filter"]["properties"]["items"][ + "items" + ]["$ref"] + assert generated_ref.startswith("#/$defs/ref_") + generated_name = generated_ref.removeprefix("#/$defs/") + generated_schema = wrapper.parameters["$defs"][generated_name] + assert generated_schema["type"] == "array" + assert generated_schema["items"]["$ref"] == generated_ref + + +def test_wrapper_hoists_root_self_ref_into_defs() -> None: + tool_def = SimpleNamespace( + name="tree", + description="tree tool", + inputSchema={ + "type": "object", + "properties": { + "children": {"type": "array", "items": {"$ref": "#"}}, + }, + }, + ) + + wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def) + + generated_ref = wrapper.parameters["properties"]["children"]["items"]["$ref"] + assert generated_ref.startswith("#/$defs/ref_") + generated_name = generated_ref.removeprefix("#/$defs/") + assert wrapper.parameters["$defs"][generated_name]["properties"]["children"]["items"] == { + "$ref": generated_ref + } + + +def test_wrapper_preserves_existing_defs_refs() -> None: + tool_def = SimpleNamespace( + name="demo", + description="demo tool", + inputSchema={ + "type": "object", + "$defs": {"value": {"type": "string"}}, + "properties": {"value": {"$ref": "#/$defs/value"}}, + }, + ) + + wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def) + + assert wrapper.parameters["properties"]["value"]["$ref"] == "#/$defs/value" + assert wrapper.parameters["$defs"]["value"]["type"] == "string" + + def test_normalize_windows_stdio_command_is_noop_off_windows( monkeypatch: pytest.MonkeyPatch, ) -> None: From c1899e2cb44944a989f4592737facdcd99733c2e Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Sun, 26 Jul 2026 17:26:35 +0800 Subject: [PATCH 21/48] fix(mcp): decode URI-encoded schema refs --- nanobot/agent/tools/mcp.py | 25 ++++++++++++++++++------- tests/tools/test_mcp_tool.py | 21 +++++++++++++++++++++ 2 files changed, 39 insertions(+), 7 deletions(-) diff --git a/nanobot/agent/tools/mcp.py b/nanobot/agent/tools/mcp.py index 7902e66bf..73d66b0dd 100644 --- a/nanobot/agent/tools/mcp.py +++ b/nanobot/agent/tools/mcp.py @@ -317,13 +317,17 @@ def _extract_nullable_branch(options: Any) -> tuple[dict[str, Any], bool] | None def _resolve_local_schema_ref(root: dict[str, Any], ref: str) -> Any: """Resolve a local JSON Pointer without accepting remote references.""" - if ref == "#": + if not ref.startswith("#"): + raise ValueError("not a local JSON Pointer") + + pointer = urllib.parse.unquote(ref[1:], errors="strict") + if not pointer: return root - if not ref.startswith("#/"): + if not pointer.startswith("/"): raise ValueError("not a local JSON Pointer") current: Any = root - for raw_part in ref[2:].split("/"): + for raw_part in pointer[1:].split("/"): part = raw_part.replace("~1", "/").replace("~0", "~") if isinstance(current, dict): current = current[part] @@ -347,15 +351,22 @@ def _rewrite_local_schema_refs(schema: dict[str, Any]) -> dict[str, Any]: rewritten = dict(value) ref = rewritten.get("$ref") - is_rewritable_ref = isinstance(ref, str) and ( - ref == "#" or (ref.startswith("#/") and not ref.startswith("#/$defs/")) - ) + is_rewritable_ref = False + if isinstance(ref, str) and not ref.startswith("#/$defs/"): + try: + pointer = urllib.parse.unquote(ref[1:], errors="strict") + except (UnicodeDecodeError, ValueError): + pass + else: + is_rewritable_ref = ref.startswith("#") and ( + not pointer or pointer.startswith("/") + ) if is_rewritable_ref: name = rewritten_refs.get(ref) if name is None: try: target = _resolve_local_schema_ref(schema, ref) - except (KeyError, IndexError, TypeError, ValueError): + except (KeyError, IndexError, TypeError, UnicodeDecodeError, ValueError): logger.warning("MCP tool schema contains an unresolved local $ref: {}", ref) else: assert isinstance(ref, str) diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py index 233c51fad..05f57ea4f 100644 --- a/tests/tools/test_mcp_tool.py +++ b/tests/tools/test_mcp_tool.py @@ -309,6 +309,27 @@ def test_wrapper_preserves_existing_defs_refs() -> None: assert wrapper.parameters["$defs"]["value"]["type"] == "string" +def test_wrapper_resolves_uri_encoded_json_pointer() -> None: + tool_def = SimpleNamespace( + name="demo", + description="demo tool", + inputSchema={ + "type": "object", + "properties": { + "space name/value": {"type": "string"}, + "alias": {"$ref": "#/properties/space%20name~1value"}, + }, + }, + ) + + wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def) + + generated_ref = wrapper.parameters["properties"]["alias"]["$ref"] + assert generated_ref.startswith("#/$defs/ref_") + generated_name = generated_ref.removeprefix("#/$defs/") + assert wrapper.parameters["$defs"][generated_name] == {"type": "string"} + + def test_normalize_windows_stdio_command_is_noop_off_windows( monkeypatch: pytest.MonkeyPatch, ) -> None: From b19039f9d09b7425958f0673d0f1877b0064c63f Mon Sep 17 00:00:00 2001 From: chengyongru Date: Thu, 23 Jul 2026 16:37:20 +0800 Subject: [PATCH 22/48] fix(agent): preserve length-recovered output --- nanobot/agent/runner.py | 32 +++++++++++++- tests/agent/test_runner_core.py | 64 +++++++++++++++++++++++++++ tests/agent/test_runner_hooks.py | 49 ++++++++++++++++++++ tests/agent/test_runner_injections.py | 36 +++++++++++++++ 4 files changed, 180 insertions(+), 1 deletion(-) diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index 9071a34ab..833447638 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -60,6 +60,18 @@ _MAX_LENGTH_RECOVERIES = 3 _MAX_INJECTIONS_PER_TURN = 3 _MAX_INJECTION_CYCLES = 5 + +def _restore_outer_whitespace(content: str, original: str | None) -> str: + """Restore boundary whitespace stripped while cleaning one recovered segment.""" + if not original: + return content + leading_size = len(original) - len(original.lstrip()) + trailing_size = len(original) - len(original.rstrip()) + leading = original[:leading_size] + trailing = original[-trailing_size:] if trailing_size else "" + return f"{leading}{content}{trailing}" + + @dataclass(slots=True) class AgentRunSpec: """Configuration for a single agent execution.""" @@ -381,6 +393,9 @@ class AgentRunner: workspace_violation_counts: dict[str, int] = {} empty_content_retries = 0 length_recovery_count = 0 + # Segments from one uninterrupted length-recovery chain. Tool work or + # injected user input starts a new logical answer and clears the chain. + length_recovery_parts: list[str] = [] had_injections = False injection_cycles = 0 compacted_tool_call_ids: set[str] = set() @@ -418,6 +433,7 @@ class AgentRunner: context.response = response context.tool_calls = list(response.tool_calls) + original_content = response.content reasoning_text, cleaned_content = extract_reasoning( response.reasoning_content, response.thinking_blocks, @@ -519,6 +535,7 @@ class AgentRunner: ) empty_content_retries = 0 length_recovery_count = 0 + length_recovery_parts.clear() # Checkpoint 1: drain injections after tools, before next LLM call _drained, injection_cycles = await self._try_drain_injections( spec, messages, None, injection_cycles, @@ -567,11 +584,15 @@ class AgentRunner: context.response = response context.usage = dict(raw_usage) context.tool_calls = list(response.tool_calls) + original_content = response.content clean = hook.finalize_content(context, response.content) if response.finish_reason == "length" and not is_blank_text(clean): length_recovery_count += 1 if length_recovery_count <= _MAX_LENGTH_RECOVERIES: + length_recovery_parts.append( + _restore_outer_whitespace(clean, original_content) + ) logger.info( "Output truncated on turn {} for {} ({}/{}); continuing", iteration, @@ -614,6 +635,7 @@ class AgentRunner: await hook.on_stream_end(context, resuming=should_continue) if should_continue: + length_recovery_parts.clear() await hook.after_iteration(context) continue @@ -635,6 +657,7 @@ class AgentRunner: ) if should_continue: had_injections = True + length_recovery_parts.clear() continue break if is_blank_text(clean): @@ -652,6 +675,7 @@ class AgentRunner: ) if should_continue: had_injections = True + length_recovery_parts.clear() continue break @@ -671,7 +695,13 @@ class AgentRunner: "pending_tool_calls": [], }, ) - final_content = clean + if length_recovery_parts: + final_content = ( + "".join(length_recovery_parts) + + _restore_outer_whitespace(clean, original_content) + ).strip() + else: + final_content = clean context.final_content = final_content context.stop_reason = stop_reason await hook.after_iteration(context) diff --git a/tests/agent/test_runner_core.py b/tests/agent/test_runner_core.py index 39b74725e..584df875c 100644 --- a/tests/agent/test_runner_core.py +++ b/tests/agent/test_runner_core.py @@ -450,6 +450,70 @@ async def test_runner_uses_specific_message_after_empty_finalization_retry(): assert result.stop_reason == "empty_final_response" +@pytest.mark.asyncio +async def test_runner_length_recovery_returns_all_segments(): + """Recovered output segments are returned together instead of only the tail.""" + from nanobot.agent.runner import AgentRunner + + provider = MagicMock(spec=LLMProvider) + provider.chat_with_retry = AsyncMock(side_effect=[ + LLMResponse(content="first ", finish_reason="length"), + LLMResponse(content="second ", finish_reason="length"), + LLMResponse(content="third", finish_reason="stop"), + ]) + tools = MagicMock() + tools.get_definitions.return_value = [] + + runner = AgentRunner() + result = await runner.run(make_run_spec(provider, + initial_messages=[{"role": "user", "content": "give a long answer"}], + tools=tools, + model="test-model", + max_iterations=5, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + )) + + assert result.final_content == "first second third" + assert [ + message["content"] + for message in result.messages + if message.get("role") == "assistant" + ] == ["first", "second", "third"] + assert provider.chat_with_retry.await_count == 3 + + +@pytest.mark.asyncio +async def test_runner_length_recovery_does_not_leak_across_tool_calls(): + """A recovered prefix belongs only to its contiguous response chain.""" + from nanobot.agent.runner import AgentRunner + + provider = MagicMock(spec=LLMProvider) + provider.chat_with_retry = AsyncMock(side_effect=[ + LLMResponse(content="working", finish_reason="length"), + LLMResponse( + content=None, + tool_calls=[ToolCallRequest(id="call_1", name="read_file", arguments={"path": "x"})], + finish_reason="tool_calls", + ), + LLMResponse(content="final answer", finish_reason="stop"), + ]) + tools = MagicMock() + tools.get_definitions.return_value = [] + tools.execute = AsyncMock(return_value="file content") + + runner = AgentRunner() + result = await runner.run(make_run_spec(provider, + initial_messages=[{"role": "user", "content": "inspect a file"}], + tools=tools, + model="test-model", + max_iterations=5, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + )) + + assert result.final_content == "final answer" + assert result.tools_used == ["read_file"] + + @pytest.mark.asyncio async def test_runner_empty_response_does_not_break_tool_chain(): """An empty intermediate response must not kill an ongoing tool chain. diff --git a/tests/agent/test_runner_hooks.py b/tests/agent/test_runner_hooks.py index f441d285e..0bc8a6fe6 100644 --- a/tests/agent/test_runner_hooks.py +++ b/tests/agent/test_runner_hooks.py @@ -143,6 +143,55 @@ async def test_runner_streaming_hook_receives_deltas_and_end_signal(): provider.chat_with_retry.assert_not_awaited() +@pytest.mark.asyncio +async def test_runner_length_recovery_streams_segments_once_and_returns_all_content(): + from nanobot.agent.hook import AgentHook, AgentHookContext + from nanobot.agent.runner import AgentRunner + + provider = MagicMock(spec=LLMProvider) + streamed: list[str] = [] + endings: list[bool] = [] + responses = iter([ + LLMResponse(content="first ", finish_reason="length"), + LLMResponse(content="second", finish_reason="stop"), + ]) + + async def chat_stream_with_retry(*, on_content_delta, **kwargs): + response = next(responses) + await on_content_delta(response.content or "") + return response + + provider.chat_stream_with_retry = chat_stream_with_retry + provider.chat_with_retry = AsyncMock() + tools = MagicMock() + tools.get_definitions.return_value = [] + + class StreamingHook(AgentHook): + def wants_streaming(self) -> bool: + return True + + async def on_stream(self, context: AgentHookContext, delta: str) -> None: + streamed.append(delta) + + async def on_stream_end(self, context: AgentHookContext, *, resuming: bool) -> None: + endings.append(resuming) + + runner = AgentRunner() + result = await runner.run(make_run_spec(provider, + initial_messages=[{"role": "user", "content": "give a long answer"}], + tools=tools, + model="test-model", + max_iterations=3, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + hook=StreamingHook(), + )) + + assert result.final_content == "first second" + assert streamed == ["first ", "second"] + assert endings == [True, False] + provider.chat_with_retry.assert_not_awaited() + + @pytest.mark.asyncio async def test_runner_passes_cached_tokens_to_hook_context(): """Hook context.usage should contain cached_tokens.""" diff --git a/tests/agent/test_runner_injections.py b/tests/agent/test_runner_injections.py index 9b96b4a8c..cd7c76a45 100644 --- a/tests/agent/test_runner_injections.py +++ b/tests/agent/test_runner_injections.py @@ -352,6 +352,42 @@ async def test_checkpoint2_injects_after_final_response_with_resuming_stream(): assert stream_end_calls[-1] is False +@pytest.mark.asyncio +async def test_injected_followup_starts_new_length_recovery_chain(): + """Recovered content from the prior answer must not prefix a follow-up reply.""" + from nanobot.agent.runner import AgentRunner + from nanobot.bus.events import InboundMessage + + provider = MagicMock() + provider.chat_with_retry = AsyncMock(side_effect=[ + LLMResponse(content="first", finish_reason="length"), + LLMResponse(content="second", finish_reason="stop"), + LLMResponse(content="follow-up answer", finish_reason="stop"), + ]) + tools = MagicMock() + tools.get_definitions.return_value = [] + + injection_queue = asyncio.Queue() + inject_cb = _make_injection_callback(injection_queue) + await injection_queue.put( + InboundMessage(channel="cli", sender_id="u", chat_id="c", content="follow-up question") + ) + + runner = AgentRunner() + result = await runner.run(make_run_spec(provider, + initial_messages=[{"role": "user", "content": "give a long answer"}], + tools=tools, + model="test-model", + max_iterations=5, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + injection_callback=inject_cb, + )) + + assert result.had_injections is True + assert result.final_content == "follow-up answer" + assert provider.chat_with_retry.await_count == 3 + + @pytest.mark.asyncio async def test_checkpoint2_preserves_final_response_in_history_before_followup(): """A follow-up injected after a final answer must still see that answer in history.""" From df2e5b7225685aecfa4f69b461f55026ef660f71 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Thu, 23 Jul 2026 17:11:09 +0800 Subject: [PATCH 23/48] fix(agent): anchor truncated response continuations --- nanobot/agent/runner.py | 2 +- nanobot/utils/runtime.py | 21 +++++++++++++++++---- tests/utils/test_runtime.py | 14 ++++++++++++++ 3 files changed, 32 insertions(+), 5 deletions(-) create mode 100644 tests/utils/test_runtime.py diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index 833447638..9428be951 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -607,7 +607,7 @@ class AgentRunner: reasoning_content=response.reasoning_content, thinking_blocks=response.thinking_blocks, )) - messages.append(build_length_recovery_message()) + messages.append(build_length_recovery_message(clean)) await hook.after_iteration(context) continue diff --git a/nanobot/utils/runtime.py b/nanobot/utils/runtime.py index 4f6599f68..8ffe62770 100644 --- a/nanobot/utils/runtime.py +++ b/nanobot/utils/runtime.py @@ -14,6 +14,7 @@ _MAX_REPEAT_EXTERNAL_LOOKUPS = 2 # Third same-target workspace violation in a turn escalates to "stop retrying". _MAX_REPEAT_WORKSPACE_VIOLATIONS = 2 +_LENGTH_RECOVERY_TAIL_CHARS = 500 EMPTY_FINAL_RESPONSE_MESSAGE = ( "I completed the tool steps but couldn't produce a final answer. " @@ -33,8 +34,10 @@ BUDGET_EXHAUSTED_FINALIZATION_PROMPT = ( ) LENGTH_RECOVERY_PROMPT = ( - "Output limit reached. Continue exactly where you left off " - "— no recap, no apology. Break remaining work into smaller steps if needed." + "The previous assistant response was cut off. Continue the same response from its " + "exact endpoint. Output only new continuation text in the same language and style. " + "Do not acknowledge this instruction, restart the response, repeat its title or any " + "existing text, recap, or apologize." ) SUSTAINED_GOAL_CONTINUE_PROMPT = ( @@ -79,9 +82,19 @@ def build_budget_exhausted_finalization_message() -> dict[str, str]: return {"role": "user", "content": BUDGET_EXHAUSTED_FINALIZATION_PROMPT} -def build_length_recovery_message() -> dict[str, str]: +def build_length_recovery_message(content: str) -> dict[str, str]: """Prompt the model to continue after hitting output token limit.""" - return {"role": "user", "content": LENGTH_RECOVERY_PROMPT} + tail = content[-_LENGTH_RECOVERY_TAIL_CHARS:] + prompt = ( + f"{LENGTH_RECOVERY_PROMPT}\n\n" + "The following tail was already delivered to the user. Treat it as immutable " + "context and do not output it again:\n" + "\n" + f"{tail}\n" + "\n" + "Begin with the text that belongs immediately after this tail." + ) + return {"role": "user", "content": prompt} def build_goal_continue_message(custom: str | None = None) -> dict[str, str]: diff --git a/tests/utils/test_runtime.py b/tests/utils/test_runtime.py new file mode 100644 index 000000000..56e2f2bc6 --- /dev/null +++ b/tests/utils/test_runtime.py @@ -0,0 +1,14 @@ +from nanobot.utils.runtime import build_length_recovery_message + + +def test_length_recovery_message_anchors_the_existing_tail() -> None: + omitted_prefix = "OMITTED_PREFIX" + tail = "x" * 500 + + message = build_length_recovery_message(omitted_prefix + tail) + + assert message["role"] == "user" + assert omitted_prefix not in message["content"] + assert f"\n{tail}\n" in message["content"] + assert "Output only new continuation text" in message["content"] + assert "Break remaining work into smaller steps" not in message["content"] From 154cbc1974bb8aa315ad686b7e56c1563d17142c Mon Sep 17 00:00:00 2001 From: chengyongru Date: Thu, 23 Jul 2026 17:14:58 +0800 Subject: [PATCH 24/48] refactor(agent): trim recovery tail anchor --- nanobot/utils/runtime.py | 2 +- tests/utils/test_runtime.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/nanobot/utils/runtime.py b/nanobot/utils/runtime.py index 8ffe62770..e755fa5be 100644 --- a/nanobot/utils/runtime.py +++ b/nanobot/utils/runtime.py @@ -14,7 +14,7 @@ _MAX_REPEAT_EXTERNAL_LOOKUPS = 2 # Third same-target workspace violation in a turn escalates to "stop retrying". _MAX_REPEAT_WORKSPACE_VIOLATIONS = 2 -_LENGTH_RECOVERY_TAIL_CHARS = 500 +_LENGTH_RECOVERY_TAIL_CHARS = 64 EMPTY_FINAL_RESPONSE_MESSAGE = ( "I completed the tool steps but couldn't produce a final answer. " diff --git a/tests/utils/test_runtime.py b/tests/utils/test_runtime.py index 56e2f2bc6..40aa5b054 100644 --- a/tests/utils/test_runtime.py +++ b/tests/utils/test_runtime.py @@ -3,7 +3,7 @@ from nanobot.utils.runtime import build_length_recovery_message def test_length_recovery_message_anchors_the_existing_tail() -> None: omitted_prefix = "OMITTED_PREFIX" - tail = "x" * 500 + tail = "x" * 64 message = build_length_recovery_message(omitted_prefix + tail) From 1d2ed6e4d2c57cf3d1a3218da00fd53ed80214f9 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Thu, 23 Jul 2026 17:40:08 +0800 Subject: [PATCH 25/48] fix(agent): reset recovery chains across injections Reset both the recovered segments and retry budget whenever injected input starts a new logical answer. Cover fatal tool-error boundaries and rename the prompt test module so pytest can collect the full suite. --- nanobot/agent/runner.py | 5 ++++ tests/agent/test_runner_injections.py | 26 +++++++++++++------ ...ime.py => test_length_recovery_runtime.py} | 2 ++ 3 files changed, 25 insertions(+), 8 deletions(-) rename tests/utils/{test_runtime.py => test_length_recovery_runtime.py} (91%) diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index 9428be951..a898936e4 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -520,6 +520,8 @@ class AgentRunner: ) if should_continue: had_injections = True + length_recovery_count = 0 + length_recovery_parts.clear() continue break await self._emit_checkpoint( @@ -635,6 +637,7 @@ class AgentRunner: await hook.on_stream_end(context, resuming=should_continue) if should_continue: + length_recovery_count = 0 length_recovery_parts.clear() await hook.after_iteration(context) continue @@ -657,6 +660,7 @@ class AgentRunner: ) if should_continue: had_injections = True + length_recovery_count = 0 length_recovery_parts.clear() continue break @@ -675,6 +679,7 @@ class AgentRunner: ) if should_continue: had_injections = True + length_recovery_count = 0 length_recovery_parts.clear() continue break diff --git a/tests/agent/test_runner_injections.py b/tests/agent/test_runner_injections.py index cd7c76a45..f76ad10bd 100644 --- a/tests/agent/test_runner_injections.py +++ b/tests/agent/test_runner_injections.py @@ -354,15 +354,18 @@ async def test_checkpoint2_injects_after_final_response_with_resuming_stream(): @pytest.mark.asyncio async def test_injected_followup_starts_new_length_recovery_chain(): - """Recovered content from the prior answer must not prefix a follow-up reply.""" + """A follow-up gets a fresh recovery budget and no content from the prior answer.""" from nanobot.agent.runner import AgentRunner from nanobot.bus.events import InboundMessage provider = MagicMock() provider.chat_with_retry = AsyncMock(side_effect=[ - LLMResponse(content="first", finish_reason="length"), - LLMResponse(content="second", finish_reason="stop"), - LLMResponse(content="follow-up answer", finish_reason="stop"), + LLMResponse(content="first-1 ", finish_reason="length"), + LLMResponse(content="first-2 ", finish_reason="length"), + LLMResponse(content="first-3 ", finish_reason="length"), + LLMResponse(content="first-final", finish_reason="stop"), + LLMResponse(content="follow-up ", finish_reason="length"), + LLMResponse(content="answer", finish_reason="stop"), ]) tools = MagicMock() tools.get_definitions.return_value = [] @@ -378,14 +381,14 @@ async def test_injected_followup_starts_new_length_recovery_chain(): initial_messages=[{"role": "user", "content": "give a long answer"}], tools=tools, model="test-model", - max_iterations=5, + max_iterations=8, max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, injection_callback=inject_cb, )) assert result.had_injections is True assert result.final_content == "follow-up answer" - assert provider.chat_with_retry.await_count == 3 + assert provider.chat_with_retry.await_count == 6 @pytest.mark.asyncio @@ -1348,7 +1351,7 @@ async def test_dispatch_republishes_leftover_queue_messages(tmp_path): @pytest.mark.asyncio async def test_drain_injections_on_fatal_tool_error(): - """Pending injections should be drained even when a fatal tool error occurs.""" + """A fatal tool error must not leak recovered content into an injected follow-up.""" from nanobot.agent.runner import AgentRunner from nanobot.bus.events import InboundMessage @@ -1358,12 +1361,18 @@ async def test_drain_injections_on_fatal_tool_error(): async def chat_with_retry(*, messages, **kwargs): call_count["n"] += 1 if call_count["n"] == 1: + return LLMResponse( + content="stale prefix ", + finish_reason="length", + usage={}, + ) + if call_count["n"] == 2: return LLMResponse( content="", tool_calls=[ToolCallRequest(id="c1", name="exec", arguments={"cmd": "bad"})], usage={}, ) - # Second call: respond normally to the injected follow-up + # Third call: respond normally to the injected follow-up. return LLMResponse(content="reply to follow-up", tool_calls=[], usage={}) provider.chat_with_retry = chat_with_retry @@ -1391,6 +1400,7 @@ async def test_drain_injections_on_fatal_tool_error(): assert result.had_injections is True assert result.final_content == "reply to follow-up" + assert call_count["n"] == 3 # The injection should be in the messages history injected = [ m for m in result.messages diff --git a/tests/utils/test_runtime.py b/tests/utils/test_length_recovery_runtime.py similarity index 91% rename from tests/utils/test_runtime.py rename to tests/utils/test_length_recovery_runtime.py index 40aa5b054..4ce2c9812 100644 --- a/tests/utils/test_runtime.py +++ b/tests/utils/test_length_recovery_runtime.py @@ -1,3 +1,5 @@ +"""Tests for length-recovery prompt construction.""" + from nanobot.utils.runtime import build_length_recovery_message From 3cc5a98d9f9e00b450f36b269743220ee55ba1f6 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Thu, 23 Jul 2026 17:54:03 +0800 Subject: [PATCH 26/48] refactor(agent): derive recovery count from segments --- nanobot/agent/runner.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index a898936e4..18e93aad8 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -392,7 +392,6 @@ class AgentRunner: # Per-turn throttle for repeated attempts against the same outside target. workspace_violation_counts: dict[str, int] = {} empty_content_retries = 0 - length_recovery_count = 0 # Segments from one uninterrupted length-recovery chain. Tool work or # injected user input starts a new logical answer and clears the chain. length_recovery_parts: list[str] = [] @@ -520,7 +519,6 @@ class AgentRunner: ) if should_continue: had_injections = True - length_recovery_count = 0 length_recovery_parts.clear() continue break @@ -536,7 +534,6 @@ class AgentRunner: }, ) empty_content_retries = 0 - length_recovery_count = 0 length_recovery_parts.clear() # Checkpoint 1: drain injections after tools, before next LLM call _drained, injection_cycles = await self._try_drain_injections( @@ -590,8 +587,7 @@ class AgentRunner: clean = hook.finalize_content(context, response.content) if response.finish_reason == "length" and not is_blank_text(clean): - length_recovery_count += 1 - if length_recovery_count <= _MAX_LENGTH_RECOVERIES: + if len(length_recovery_parts) < _MAX_LENGTH_RECOVERIES: length_recovery_parts.append( _restore_outer_whitespace(clean, original_content) ) @@ -599,7 +595,7 @@ class AgentRunner: "Output truncated on turn {} for {} ({}/{}); continuing", iteration, spec.session_key or "default", - length_recovery_count, + len(length_recovery_parts), _MAX_LENGTH_RECOVERIES, ) if hook.wants_streaming(): @@ -637,7 +633,6 @@ class AgentRunner: await hook.on_stream_end(context, resuming=should_continue) if should_continue: - length_recovery_count = 0 length_recovery_parts.clear() await hook.after_iteration(context) continue @@ -660,7 +655,6 @@ class AgentRunner: ) if should_continue: had_injections = True - length_recovery_count = 0 length_recovery_parts.clear() continue break @@ -679,7 +673,6 @@ class AgentRunner: ) if should_continue: had_injections = True - length_recovery_count = 0 length_recovery_parts.clear() continue break From 27a00c7a4f1bf952b2b224b49546f8d23dde7904 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Thu, 23 Jul 2026 18:35:41 +0800 Subject: [PATCH 27/48] fix(webui): merge length recovery stream segments --- nanobot/agent/hook.py | 1 + nanobot/agent/loop.py | 31 ++++++-- nanobot/agent/progress_hook.py | 8 +- nanobot/agent/runner.py | 1 + nanobot/agent/turn_delivery.py | 11 ++- nanobot/bus/outbound_events.py | 2 + nanobot/channels/base.py | 4 + nanobot/channels/discord/runtime.py | 1 + nanobot/channels/feishu/runtime.py | 1 + nanobot/channels/manager.py | 24 +++++- nanobot/channels/matrix/runtime.py | 1 + nanobot/channels/mattermost/runtime.py | 1 + nanobot/channels/telegram/runtime.py | 1 + nanobot/channels/websocket/runtime.py | 9 ++- .../websocket/tests/test_websocket_channel.py | 33 +++++++++ nanobot/channels/weixin/runtime.py | 1 + nanobot/webui/transcript.py | 8 +- tests/agent/test_loop_progress.py | 46 ++++++++++++ tests/agent/test_runner_hooks.py | 3 + tests/bus/test_outbound_events.py | 11 ++- .../test_channel_manager_delta_coalescing.py | 19 ++++- tests/channels/test_channel_plugins.py | 16 ++-- tests/utils/test_webui_transcript.py | 20 +++++ webui/src/hooks/useNanobotStream.ts | 26 +++++-- webui/src/lib/types.ts | 2 + webui/src/tests/useNanobotStream.test.tsx | 73 +++++++++++++++++++ 26 files changed, 325 insertions(+), 29 deletions(-) diff --git a/nanobot/agent/hook.py b/nanobot/agent/hook.py index 12a4e7170..8e4dd5ffe 100644 --- a/nanobot/agent/hook.py +++ b/nanobot/agent/hook.py @@ -25,6 +25,7 @@ class AgentHookContext: tool_events: list[dict[str, str]] = field(default_factory=list) streamed_content: bool = False streamed_reasoning: bool = False + stream_continues_current_message: 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 f079867bb..e972b6be8 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio import dataclasses +import inspect import os import time from collections.abc import Mapping @@ -860,9 +861,9 @@ class AgentLoop: """Run the agent iteration loop. *on_stream*: called with each content delta during streaming. - *on_stream_end(resuming)*: called when a streaming session finishes. - ``resuming=True`` means tool calls follow (spinner should restart); - ``resuming=False`` means this is the final response. + *on_stream_end(resuming, merge_next)*: called when a streaming session finishes. + ``resuming=True`` means the active turn continues. ``merge_next=True`` means + the next text segment belongs to the same user-visible assistant message. Returns (final_content, tools_used, messages, stop_reason, had_injections). """ @@ -1385,6 +1386,19 @@ class AgentLoop: if ctx.on_stream is not None: stream_callback = ctx.on_stream stream_end_callback = ctx.on_stream_end + stream_end_accepts_merge_next = False + if stream_end_callback is not None: + try: + stream_end_signature = inspect.signature(stream_end_callback) + stream_end_accepts_merge_next = ( + "merge_next" in stream_end_signature.parameters + or any( + parameter.kind is inspect.Parameter.VAR_KEYWORD + for parameter in stream_end_signature.parameters.values() + ) + ) + except (TypeError, ValueError): + pass segment_streamed_content = False async def _tracked_stream(delta: str) -> None: @@ -1393,12 +1407,19 @@ class AgentLoop: segment_streamed_content = True await stream_callback(delta) - async def _tracked_stream_end(*, resuming: bool = False) -> None: + async def _tracked_stream_end( + *, + resuming: bool = False, + merge_next: bool = False, + ) -> None: nonlocal segment_streamed_content ctx.streamed_content = segment_streamed_content segment_streamed_content = False if stream_end_callback is not None: - await stream_end_callback(resuming=resuming) + if merge_next and stream_end_accepts_merge_next: + await stream_end_callback(resuming=resuming, merge_next=True) + else: + await stream_end_callback(resuming=resuming) ctx.on_stream = _tracked_stream ctx.on_stream_end = _tracked_stream_end diff --git a/nanobot/agent/progress_hook.py b/nanobot/agent/progress_hook.py index 205c20192..826093d9a 100644 --- a/nanobot/agent/progress_hook.py +++ b/nanobot/agent/progress_hook.py @@ -85,7 +85,13 @@ class AgentProgressHook(AgentHook): async def on_stream_end(self, context: AgentHookContext, *, resuming: bool) -> None: await self.emit_reasoning_end() if self._on_stream_end: - await self._on_stream_end(resuming=resuming) + kwargs: dict[str, bool] = {"resuming": resuming} + if ( + context.stream_continues_current_message + and self._on_progress_accepts(self._on_stream_end, "merge_next") + ): + kwargs["merge_next"] = True + await self._on_stream_end(**kwargs) self._stream_buf = "" self._think_extractor.reset() diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index 18e93aad8..3d217822b 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -599,6 +599,7 @@ class AgentRunner: _MAX_LENGTH_RECOVERIES, ) if hook.wants_streaming(): + context.stream_continues_current_message = True await hook.on_stream_end(context, resuming=True) messages.append(build_assistant_message( clean, diff --git a/nanobot/agent/turn_delivery.py b/nanobot/agent/turn_delivery.py index 2eeb5aab5..a415551d8 100644 --- a/nanobot/agent/turn_delivery.py +++ b/nanobot/agent/turn_delivery.py @@ -285,7 +285,12 @@ class TurnDelivery: ) ) - async def _publish_stream_end(self, *, resuming: bool = False) -> None: + async def _publish_stream_end( + self, + *, + resuming: bool = False, + merge_next: bool = False, + ) -> None: await self.bus.publish_outbound( outbound_message_for_event( channel=self.delivery_message.channel, @@ -293,8 +298,10 @@ class TurnDelivery: event=StreamEndEvent( stream_id=self._stream_id(), resuming=resuming, + merge_next=merge_next, ), metadata=self.delivery_message.metadata, ) ) - self._stream_segment += 1 + if not merge_next: + self._stream_segment += 1 diff --git a/nanobot/bus/outbound_events.py b/nanobot/bus/outbound_events.py index 3b03068e2..1a5b8d551 100644 --- a/nanobot/bus/outbound_events.py +++ b/nanobot/bus/outbound_events.py @@ -46,6 +46,7 @@ class StreamEndEvent(OutboundEvent): content: str = "" stream_id: str | None = None resuming: bool = False + merge_next: bool = False @dataclass(frozen=True) @@ -176,6 +177,7 @@ def _legacy_event_from_metadata(msg: OutboundMessage) -> OutboundEvent | None: content=msg.content, stream_id=_metadata_str(meta, "_stream_id"), resuming=bool(meta.get("_resuming")), + merge_next=bool(meta.get("_merge_next")), ) if meta.get("_stream_delta"): return StreamDeltaEvent( diff --git a/nanobot/channels/base.py b/nanobot/channels/base.py index c4f51b7ac..01a794a44 100644 --- a/nanobot/channels/base.py +++ b/nanobot/channels/base.py @@ -110,6 +110,7 @@ class BaseChannel(ABC): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: """Deliver a streaming text chunk. @@ -118,6 +119,9 @@ class BaseChannel(ABC): Stateful implementations should key buffers by ``stream_id`` rather than only by ``chat_id`` when it is provided. + + ``merge_next`` marks a resumable provider boundary whose next text + segment belongs to the same user-visible message. """ pass diff --git a/nanobot/channels/discord/runtime.py b/nanobot/channels/discord/runtime.py index c6d858d68..b20baca1f 100644 --- a/nanobot/channels/discord/runtime.py +++ b/nanobot/channels/discord/runtime.py @@ -489,6 +489,7 @@ class DiscordChannel(BaseChannel): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: """Progressive Discord delivery: send once, then edit until the stream ends.""" client = self._client diff --git a/nanobot/channels/feishu/runtime.py b/nanobot/channels/feishu/runtime.py index 51f6d8b56..1541fd979 100644 --- a/nanobot/channels/feishu/runtime.py +++ b/nanobot/channels/feishu/runtime.py @@ -2216,6 +2216,7 @@ class FeishuChannel(BaseChannel): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: """Progressive streaming via CardKit: create card on first delta, stream-update on subsequent. diff --git a/nanobot/channels/manager.py b/nanobot/channels/manager.py index 03fe89beb..3fbe32ede 100644 --- a/nanobot/channels/manager.py +++ b/nanobot/channels/manager.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio import hashlib +import inspect from collections.abc import Callable, Iterable from contextlib import suppress from pathlib import Path @@ -763,13 +764,29 @@ class ChannelManager: msg: OutboundMessage, event: StreamDeltaEvent | StreamEndEvent, ) -> None: + kwargs: dict[str, Any] = { + "stream_id": event.stream_id, + "stream_end": isinstance(event, StreamEndEvent), + "resuming": event.resuming if isinstance(event, StreamEndEvent) else False, + } + if isinstance(event, StreamEndEvent) and event.merge_next: + try: + signature = inspect.signature(channel.send_delta) + if ( + "merge_next" in signature.parameters + or any( + parameter.kind is inspect.Parameter.VAR_KEYWORD + for parameter in signature.parameters.values() + ) + ): + kwargs["merge_next"] = True + except (TypeError, ValueError): + pass await channel.send_delta( msg.chat_id, msg.content, msg.metadata, - stream_id=event.stream_id, - stream_end=isinstance(event, StreamEndEvent), - resuming=event.resuming if isinstance(event, StreamEndEvent) else False, + **kwargs, ) @staticmethod @@ -850,6 +867,7 @@ class ChannelManager: final_event = StreamEndEvent( stream_id=next_stream_id, resuming=next_event.resuming, + merge_next=next_event.merge_next, ) # Stream ended - stop coalescing this stream break diff --git a/nanobot/channels/matrix/runtime.py b/nanobot/channels/matrix/runtime.py index 2544d64e9..f0cccfb49 100644 --- a/nanobot/channels/matrix/runtime.py +++ b/nanobot/channels/matrix/runtime.py @@ -598,6 +598,7 @@ class MatrixChannel(BaseChannel): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: relates_to = self._build_thread_relates_to(metadata) diff --git a/nanobot/channels/mattermost/runtime.py b/nanobot/channels/mattermost/runtime.py index 756daed73..d15bcd3fe 100644 --- a/nanobot/channels/mattermost/runtime.py +++ b/nanobot/channels/mattermost/runtime.py @@ -515,6 +515,7 @@ class MattermostChannel(BaseChannel): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: if not self._http_client: return diff --git a/nanobot/channels/telegram/runtime.py b/nanobot/channels/telegram/runtime.py index 3a8020740..4183a82c4 100644 --- a/nanobot/channels/telegram/runtime.py +++ b/nanobot/channels/telegram/runtime.py @@ -923,6 +923,7 @@ class TelegramChannel(BaseChannel): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: """Progressive message editing: send on first delta, edit on subsequent ones.""" if not self._app: diff --git a/nanobot/channels/websocket/runtime.py b/nanobot/channels/websocket/runtime.py index 6741ad43f..cdd8e4f89 100644 --- a/nanobot/channels/websocket/runtime.py +++ b/nanobot/channels/websocket/runtime.py @@ -995,13 +995,18 @@ class WebSocketChannel(BaseChannel): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: conns = list(self._subs.get(chat_id, ())) meta = metadata or {} stream_key = (chat_id, str(stream_id or "")) if stream_end: body: dict[str, Any] = {"event": "stream_end", "chat_id": chat_id} - buffered = self._stream_text_buffers.pop(stream_key, []) + buffered = ( + self._stream_text_buffers.setdefault(stream_key, []) + if merge_next + else self._stream_text_buffers.pop(stream_key, []) + ) if delta: buffered.append(delta) full_text = "".join(buffered) @@ -1019,6 +1024,8 @@ class WebSocketChannel(BaseChannel): body["stream_id"] = stream_id if stream_end and resuming: body["resuming"] = True + if stream_end and merge_next: + body["merge_next"] = True self._transcripts.prepare_and_append( chat_id, body, diff --git a/nanobot/channels/websocket/tests/test_websocket_channel.py b/nanobot/channels/websocket/tests/test_websocket_channel.py index fcefdb234..92eb84bb7 100644 --- a/nanobot/channels/websocket/tests/test_websocket_channel.py +++ b/nanobot/channels/websocket/tests/test_websocket_channel.py @@ -1350,6 +1350,39 @@ async def test_send_delta_marks_resuming_stream_end() -> None: assert payload["resuming"] is True +@pytest.mark.asyncio +async def test_send_delta_keeps_buffer_across_merged_stream_boundary() -> None: + bus = MagicMock() + channel = WebSocketChannel( + {"enabled": True, "allowFrom": ["*"], "streaming": True}, + bus, + gateway=_basic_handler(bus), + ) + mock_ws = AsyncMock() + channel._attach(mock_ws, "chat-1") + + await channel.send_delta("chat-1", "first ", stream_id="sid") + await channel.send_delta( + "chat-1", + "", + stream_id="sid", + stream_end=True, + resuming=True, + merge_next=True, + ) + await channel.send_delta("chat-1", "second", stream_id="sid") + await channel.send_delta("chat-1", "", stream_id="sid", stream_end=True) + + payloads = [json.loads(call.args[0]) for call in mock_ws.send.await_args_list] + assert payloads[1]["merge_next"] is True + assert payloads[1]["resuming"] is True + assert [payload["text"] for payload in payloads if payload["event"] == "delta"] == [ + "first ", + "second", + ] + assert ("chat-1", "sid") not in channel._stream_text_buffers + + @pytest.mark.asyncio async def test_send_delta_stream_end_includes_inline_final_text() -> None: bus = MagicMock() diff --git a/nanobot/channels/weixin/runtime.py b/nanobot/channels/weixin/runtime.py index eab09c84d..976013c08 100644 --- a/nanobot/channels/weixin/runtime.py +++ b/nanobot/channels/weixin/runtime.py @@ -1243,6 +1243,7 @@ class WeixinChannel(BaseChannel): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: """Deliver a streamed reply to WeChat. diff --git a/nanobot/webui/transcript.py b/nanobot/webui/transcript.py index 63a9aac95..2c8a73290 100644 --- a/nanobot/webui/transcript.py +++ b/nanobot/webui/transcript.py @@ -1770,6 +1770,7 @@ def replay_transcript_to_ui_messages( buffer_message_id = None buffer_parts = [] continue + merge_next = rec.get("resuming") is True and rec.get("merge_next") is True final_text = rec.get("text") if isinstance(final_text, str): if buffer_message_id is None: @@ -1794,8 +1795,11 @@ def replay_transcript_to_ui_messages( **_turn_fields(rec, "answer"), } break - buffer_message_id = None - buffer_parts = [] + if merge_next: + buffer_parts = [final_text] + if not merge_next: + buffer_message_id = None + buffer_parts = [] continue if ev == "reasoning_delta": diff --git a/tests/agent/test_loop_progress.py b/tests/agent/test_loop_progress.py index 3b4906456..aca35927e 100644 --- a/tests/agent/test_loop_progress.py +++ b/tests/agent/test_loop_progress.py @@ -534,6 +534,52 @@ class TestToolEventProgress: assert turn_end_msgs[0].content == "" provider.chat_with_retry.assert_not_awaited() + @pytest.mark.asyncio + async def test_length_recovery_keeps_one_user_visible_stream( + self, + tmp_path: Path, + ) -> None: + bus = MessageBus() + provider = MagicMock() + provider.supports_progress_deltas = True + provider.get_default_model.return_value = "test-model" + responses = iter([ + LLMResponse(content="first-", finish_reason="length"), + LLMResponse(content="second", finish_reason="stop"), + ]) + + async def chat_stream_with_retry(*, on_content_delta, **kwargs): + response = next(responses) + await on_content_delta(response.content or "") + return response + + provider.chat_stream_with_retry = chat_stream_with_retry + provider.chat_with_retry = AsyncMock() + loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model") + _attach_webui_runtime_events(loop, bus) + loop.tools.get_definitions = MagicMock(return_value=[]) + loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] + + await loop._dispatch(InboundMessage( + channel="websocket", + sender_id="u1", + chat_id="chat1", + content="give a long answer", + metadata={"_wants_stream": True}, + )) + + outbound = [] + while bus.outbound_size > 0: + outbound.append(await bus.consume_outbound()) + + deltas = [m.event for m in outbound if isinstance(m.event, StreamDeltaEvent)] + endings = [m.event for m in outbound if isinstance(m.event, StreamEndEvent)] + + assert [event.content for event in deltas] == ["first-", "second"] + assert [event.resuming for event in endings] == [True, False] + assert [event.merge_next for event in endings] == [True, False] + assert {event.stream_id for event in [*deltas, *endings]} == {deltas[0].stream_id} + @pytest.mark.asyncio async def test_non_streamed_finalization_is_delivered_as_regular_message( self, diff --git a/tests/agent/test_runner_hooks.py b/tests/agent/test_runner_hooks.py index 0bc8a6fe6..508f663e0 100644 --- a/tests/agent/test_runner_hooks.py +++ b/tests/agent/test_runner_hooks.py @@ -151,6 +151,7 @@ async def test_runner_length_recovery_streams_segments_once_and_returns_all_cont provider = MagicMock(spec=LLMProvider) streamed: list[str] = [] endings: list[bool] = [] + merge_next: list[bool] = [] responses = iter([ LLMResponse(content="first ", finish_reason="length"), LLMResponse(content="second", finish_reason="stop"), @@ -175,6 +176,7 @@ async def test_runner_length_recovery_streams_segments_once_and_returns_all_cont async def on_stream_end(self, context: AgentHookContext, *, resuming: bool) -> None: endings.append(resuming) + merge_next.append(context.stream_continues_current_message) runner = AgentRunner() result = await runner.run(make_run_spec(provider, @@ -189,6 +191,7 @@ async def test_runner_length_recovery_streams_segments_once_and_returns_all_cont assert result.final_content == "first second" assert streamed == ["first ", "second"] assert endings == [True, False] + assert merge_next == [True, False] provider.chat_with_retry.assert_not_awaited() diff --git a/tests/bus/test_outbound_events.py b/tests/bus/test_outbound_events.py index 0beb7e60c..99eb259c3 100644 --- a/tests/bus/test_outbound_events.py +++ b/tests/bus/test_outbound_events.py @@ -94,7 +94,12 @@ def test_legacy_stream_metadata_flags_create_runtime_events() -> None: channel="websocket", chat_id="chat-1", content="", - metadata={"_stream_end": True, "_stream_id": "s1", "_resuming": True}, + metadata={ + "_stream_end": True, + "_stream_id": "s1", + "_resuming": True, + "_merge_next": True, + }, ) delta_event = outbound_event_from_message(delta) @@ -106,6 +111,7 @@ def test_legacy_stream_metadata_flags_create_runtime_events() -> None: assert isinstance(end_event, StreamEndEvent) assert end_event.stream_id == "s1" assert end_event.resuming is True + assert end_event.merge_next is True def test_legacy_webui_runtime_metadata_flags_create_runtime_events() -> None: @@ -221,7 +227,7 @@ def test_replace_outbound_event_keeps_routing_metadata() -> None: updated = replace_outbound_event( msg, - StreamEndEvent(stream_id="s1", resuming=True), + StreamEndEvent(stream_id="s1", resuming=True, merge_next=True), content="hello world", ) @@ -230,6 +236,7 @@ def test_replace_outbound_event_keeps_routing_metadata() -> None: assert isinstance(updated.event, StreamEndEvent) assert updated.event.stream_id == "s1" assert updated.event.resuming is True + assert updated.event.merge_next is True def test_streamed_response_event_keeps_final_content_outside_event_payload() -> None: diff --git a/tests/channels/test_channel_manager_delta_coalescing.py b/tests/channels/test_channel_manager_delta_coalescing.py index de60124e2..8a009d848 100644 --- a/tests/channels/test_channel_manager_delta_coalescing.py +++ b/tests/channels/test_channel_manager_delta_coalescing.py @@ -49,6 +49,7 @@ class MockChannel(BaseChannel): stream_id=None, stream_end=False, resuming=False, + merge_next=False, ): return await self._send_delta_mock( chat_id, @@ -57,6 +58,7 @@ class MockChannel(BaseChannel): stream_id=stream_id, stream_end=stream_end, resuming=resuming, + merge_next=merge_next, ) @@ -92,11 +94,17 @@ def _end( chat_id: str = "chat1", stream_id: str | None = None, resuming: bool = False, + merge_next: bool = False, ): return outbound_message_for_event( channel="mock", chat_id=chat_id, - event=StreamEndEvent(content=content, stream_id=stream_id, resuming=resuming), + event=StreamEndEvent( + content=content, + stream_id=stream_id, + resuming=resuming, + merge_next=merge_next, + ), ) @@ -137,6 +145,7 @@ class TestDeltaCoalescing: stream_id=None, stream_end=False, resuming=False, + merge_next=False, ) @pytest.mark.asyncio @@ -184,13 +193,19 @@ class TestDeltaCoalescing: @pytest.mark.asyncio async def test_stream_end_terminates_coalescing(self, manager, bus): await bus.publish_outbound(_delta("Hello")) - await bus.publish_outbound(_end(" world")) + await bus.publish_outbound(_end( + " world", + resuming=True, + merge_next=True, + )) first_msg = await bus.consume_outbound() merged, pending = manager._coalesce_stream_deltas(first_msg) assert merged.content == "Hello world" assert isinstance(merged.event, StreamEndEvent) + assert merged.event.resuming is True + assert merged.event.merge_next is True assert len(pending) == 0 @pytest.mark.asyncio diff --git a/tests/channels/test_channel_plugins.py b/tests/channels/test_channel_plugins.py index 50ff6905f..93a8b31d2 100644 --- a/tests/channels/test_channel_plugins.py +++ b/tests/channels/test_channel_plugins.py @@ -2818,7 +2818,7 @@ async def test_send_with_retry_no_retry_when_max_is_zero(): @pytest.mark.asyncio async def test_send_with_retry_calls_send_delta(): """_send_with_retry should call send_delta for stream delta events.""" - calls: list[tuple[str, str, str | None, bool, bool]] = [] + calls: list[tuple[str, str, str | None, bool, bool, bool]] = [] class _StreamingChannel(BaseChannel): name = "streaming" @@ -2842,8 +2842,9 @@ async def test_send_with_retry_calls_send_delta(): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: - calls.append((chat_id, delta, stream_id, stream_end, resuming)) + calls.append((chat_id, delta, stream_id, stream_end, resuming, merge_next)) fake_config = SimpleNamespace( channels=ChannelsConfig(send_max_retries=3), @@ -2865,13 +2866,18 @@ async def test_send_with_retry_calls_send_delta(): end = outbound_message_for_event( channel="streaming", chat_id="123", - event=StreamEndEvent(content="", stream_id="s1", resuming=True), + event=StreamEndEvent( + content="", + stream_id="s1", + resuming=True, + merge_next=True, + ), ) await mgr._send_with_retry(mgr.channels["streaming"], end) assert calls == [ - ("123", "test delta", "s1", False, False), - ("123", "", "s1", True, True), + ("123", "test delta", "s1", False, False, False), + ("123", "", "s1", True, True, True), ] diff --git a/tests/utils/test_webui_transcript.py b/tests/utils/test_webui_transcript.py index 067a45fad..921982f6b 100644 --- a/tests/utils/test_webui_transcript.py +++ b/tests/utils/test_webui_transcript.py @@ -1121,6 +1121,26 @@ def test_replay_keeps_interrupted_pre_tool_text_in_activity() -> None: assert msgs[2]["content"] == "Done. Open index.html to play." +def test_replay_merges_length_recovery_segments_into_one_assistant_message() -> None: + msgs = replay_transcript_to_ui_messages([ + {"event": "delta", "chat_id": "t-stream", "text": "first "}, + { + "event": "stream_end", + "chat_id": "t-stream", + "text": "first ", + "resuming": True, + "merge_next": True, + }, + {"event": "delta", "chat_id": "t-stream", "text": "second"}, + {"event": "stream_end", "chat_id": "t-stream"}, + {"event": "turn_end", "chat_id": "t-stream"}, + ]) + + assert len(msgs) == 1 + assert msgs[0]["role"] == "assistant" + assert msgs[0]["content"] == "first second" + + def test_replay_tool_events_dedupes_finish_after_start() -> None: msgs = replay_transcript_to_ui_messages([ { diff --git a/webui/src/hooks/useNanobotStream.ts b/webui/src/hooks/useNanobotStream.ts index 4d2ee120e..1f18e2007 100644 --- a/webui/src/hooks/useNanobotStream.ts +++ b/webui/src/hooks/useNanobotStream.ts @@ -26,7 +26,7 @@ import type { } from "@/lib/types"; interface StreamBuffer { - /** ID of the assistant message currently receiving deltas (cleared on ``stream_end``). */ + /** ID of the assistant message currently receiving deltas (cleared when its segment closes). */ messageId: string; } @@ -780,15 +780,20 @@ export function useNanobotStream( ?? findStreamingAssistantIndex(next, closedAssistantStreamIdsRef.current, turn); if (targetIndex !== null) { const target = next[targetIndex]; - next = replaceMessageAt(next, targetIndex, { + const merged = { ...target, content: finalAnswerText, isStreaming: true, ...turn, - }); + }; + next = replaceMessageAt(next, targetIndex, merged); + if (!options?.closeAnswerSegment) { + closedAssistantStreamIdsRef.current.delete(merged.id); + activeAssistantRef.current = { id: merged.id, index: targetIndex }; + buffer.current = { messageId: merged.id }; + } } else { const id = crypto.randomUUID(); - closedAssistantStreamIdsRef.current.add(id); next = [ ...next, { @@ -800,6 +805,12 @@ export function useNanobotStream( createdAt: Date.now(), }, ]; + if (options?.closeAnswerSegment) { + closedAssistantStreamIdsRef.current.add(id); + } else { + activeAssistantRef.current = { id, index: next.length - 1 }; + buffer.current = { messageId: id }; + } } } if (options?.closeAnswerSegment) closeActiveAssistantStream(); @@ -911,8 +922,9 @@ export function useNanobotStream( if (ev.event === "stream_end") { const turn = turnFieldsFromEvent(ev, "answer"); + const mergeNext = ev.resuming === true && ev.merge_next === true; flushPendingStreamEvents({ - closeAnswerSegment: true, + closeAnswerSegment: !mergeNext, ...(typeof ev.text === "string" ? { finalAnswerText: ev.text } : {}), turn, }); @@ -920,7 +932,9 @@ export function useNanobotStream( if (ev.resuming) { cancelStreamEndTimer(); setIsStreaming(true); - setMessages((prev) => finalizeStreamedTurn(prev, turn)); + if (!mergeNext) { + setMessages((prev) => finalizeStreamedTurn(prev, turn)); + } return; } scheduleStreamEndTimer(turn); diff --git a/webui/src/lib/types.ts b/webui/src/lib/types.ts index 5a9ce09ef..1db55490d 100644 --- a/webui/src/lib/types.ts +++ b/webui/src/lib/types.ts @@ -1118,6 +1118,8 @@ export type InboundEvent = text?: string; /** This answer segment ended, but the active agent turn will continue. */ resuming?: boolean; + /** The next answer segment continues this same assistant message. */ + merge_next?: boolean; } & InboundTurnMetadata) | ({ event: "reasoning_delta"; diff --git a/webui/src/tests/useNanobotStream.test.tsx b/webui/src/tests/useNanobotStream.test.tsx index c5cd49a31..ebd2db0d9 100644 --- a/webui/src/tests/useNanobotStream.test.tsx +++ b/webui/src/tests/useNanobotStream.test.tsx @@ -2009,6 +2009,79 @@ describe("useNanobotStream", () => { expect(result.current.messages.every((message) => !message.isStreaming)).toBe(true); }); + it("keeps length-recovery segments in one assistant message", async () => { + const fake = fakeClient(); + const { result } = renderHook(() => useNanobotStream("chat-length", EMPTY_MESSAGES), { + wrapper: wrap(fake.client), + }); + + act(() => { + result.current.send("give a long answer"); + }); + const activeTurnId = fake.client.sendMessage.mock.calls.at(-1)![3]?.turnId; + + act(() => { + fake.emit("chat-length", { + event: "delta", + chat_id: "chat-length", + text: "first ", + turn_id: activeTurnId, + }); + }); + await flushStreamFrame(); + const assistantId = result.current.messages[1].id; + + act(() => { + fake.emit("chat-length", { + event: "stream_end", + chat_id: "chat-length", + text: "first ", + resuming: true, + merge_next: true, + turn_id: activeTurnId, + }); + }); + + expect(result.current.messages).toHaveLength(2); + expect(result.current.messages[1]).toMatchObject({ + id: assistantId, + content: "first ", + isStreaming: true, + }); + + act(() => { + fake.emit("chat-length", { + event: "delta", + chat_id: "chat-length", + text: "second", + turn_id: activeTurnId, + }); + }); + await flushStreamFrame(); + + expect(result.current.messages).toHaveLength(2); + expect(result.current.messages[1]).toMatchObject({ + id: assistantId, + content: "first second", + isStreaming: true, + }); + + act(() => { + fake.emit("chat-length", { + event: "turn_end", + chat_id: "chat-length", + turn_id: activeTurnId, + }); + }); + + expect(result.current.messages).toHaveLength(2); + expect(result.current.messages[1]).toMatchObject({ + id: assistantId, + content: "first second", + isStreaming: false, + }); + }); + it("keeps streaming alive across stream_end when tool activity follows", async () => { const fake = fakeClient(); const onTurnEnd = vi.fn(); From e6baecafcd78a9eacaa4fe4c62528332c3b26c97 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Fri, 24 Jul 2026 10:20:43 +0800 Subject: [PATCH 28/48] fix(agent): close length recovery lifecycle gaps --- nanobot/agent/loop.py | 15 +++- nanobot/agent/runner.py | 22 +++-- nanobot/agent/turn_delivery.py | 8 ++ nanobot/channels/mattermost/runtime.py | 6 +- .../tests/test_mattermost_channel.py | 27 ++++++ tests/agent/test_loop_progress.py | 90 +++++++++++++++++++ tests/agent/test_runner_core.py | 34 +++++++ 7 files changed, 195 insertions(+), 7 deletions(-) diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index e972b6be8..eab6f92e3 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -1070,7 +1070,12 @@ class AgentLoop: # Push final content through stream so streaming channels (e.g. Feishu) # update the card instead of leaving it empty. if on_stream and on_stream_end and should_stream: - await on_stream(result.final_content or "") + stream_content = ( + result.pending_stream_content + if result.pending_stream_content is not None + else result.final_content or "" + ) + await on_stream(stream_content) await on_stream_end(resuming=False) elif result.stop_reason == "error": logger.error("LLM returned error: {}", (result.final_content or "")[:200]) @@ -1217,6 +1222,14 @@ class AgentLoop: for _, coordinator in self._automation_turn_coordinators: coordinator.complete(msg, error=asyncio.CancelledError()) logger.info("Task cancelled for session {}", session_key) + try: + await delivery.abort_stream() + except Exception: + logger.debug( + "Could not close stream for cancelled session {}", + session_key, + exc_info=True, + ) # Preserve partial context from the interrupted turn so # the user does not lose tool results and assistant # messages accumulated before /stop. The checkpoint was diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index 3d217822b..acc54abbf 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -113,6 +113,8 @@ class AgentRunResult: error: str | None = None tool_events: list[dict[str, str]] = field(default_factory=list) had_injections: bool = False + # Terminal tail to emit when the preceding final-content prefix was already streamed. + pending_stream_content: str | None = None class AgentRunner: @@ -398,6 +400,7 @@ class AgentRunner: had_injections = False injection_cycles = 0 compacted_tool_call_ids: set[str] = set() + pending_stream_content: str | None = None governance_config = ContextGovernanceConfig( provider=spec.runtime.provider, model=spec.runtime.model, @@ -718,17 +721,25 @@ class AgentRunner: ) if drained_after_max_iterations: had_injections = True - final_content = None + terminal_content = None if spec.finalize_on_max_iterations: - final_content = await self._try_finalize_after_max_iterations( + terminal_content = await self._try_finalize_after_max_iterations( spec, hook, messages, usage, ) - if final_content is None: - final_content = self._max_iterations_fallback(spec) - self._append_final_message(messages, final_content) + if terminal_content is None: + terminal_content = self._max_iterations_fallback(spec) + if length_recovery_parts: + terminal_tail = f"\n\n{terminal_content.lstrip()}" + final_content = ( + "".join(length_recovery_parts).rstrip() + terminal_tail + ).strip() + pending_stream_content = terminal_tail + else: + final_content = terminal_content + self._append_final_message(messages, terminal_content) return AgentRunResult( final_content=final_content, @@ -739,6 +750,7 @@ class AgentRunner: error=error, tool_events=tool_events, had_injections=had_injections, + pending_stream_content=pending_stream_content, ) def _build_request_kwargs( diff --git a/nanobot/agent/turn_delivery.py b/nanobot/agent/turn_delivery.py index a415551d8..9f2aba0a4 100644 --- a/nanobot/agent/turn_delivery.py +++ b/nanobot/agent/turn_delivery.py @@ -126,6 +126,7 @@ class TurnDelivery: lifecycle_message: InboundMessage = field(init=False) _stream_base_id: str | None = field(init=False, default=None) _stream_segment: int = field(init=False, default=0) + _stream_open: bool = field(init=False, default=False) def __post_init__(self) -> None: self.delivery_message = dataclasses.replace( @@ -284,6 +285,7 @@ class TurnDelivery: metadata=self.delivery_message.metadata, ) ) + self._stream_open = True async def _publish_stream_end( self, @@ -303,5 +305,11 @@ class TurnDelivery: metadata=self.delivery_message.metadata, ) ) + self._stream_open = merge_next if not merge_next: self._stream_segment += 1 + + async def abort_stream(self) -> None: + """Close an interrupted stream so stateful channels can release its buffer.""" + if self._stream_open: + await self._publish_stream_end() diff --git a/nanobot/channels/mattermost/runtime.py b/nanobot/channels/mattermost/runtime.py index d15bcd3fe..473dac316 100644 --- a/nanobot/channels/mattermost/runtime.py +++ b/nanobot/channels/mattermost/runtime.py @@ -533,7 +533,11 @@ class MattermostChannel(BaseChannel): final += delta if resuming: - self._clear_stream_state(stream_id) + if merge_next: + self._stream_buffers[stream_id] = final + self._stream_committed[stream_id] = final + else: + self._clear_stream_state(stream_id) return if final and not meta.get("_progress"): diff --git a/nanobot/channels/mattermost/tests/test_mattermost_channel.py b/nanobot/channels/mattermost/tests/test_mattermost_channel.py index 3b27ccf1c..f9afcdd15 100644 --- a/nanobot/channels/mattermost/tests/test_mattermost_channel.py +++ b/nanobot/channels/mattermost/tests/test_mattermost_channel.py @@ -582,6 +582,33 @@ async def test_stream_end_keyword_resuming_does_not_post_or_mark_done(): assert "s1" not in channel._stream_buffers +@pytest.mark.asyncio +async def test_stream_end_merge_next_preserves_buffer_until_final_end(): + channel, fake = _make_channel() + channel._self_id = "bot_id" + fake.set_post_response("/api/v4/posts", {"id": "stream_post_1"}) + await channel.send_delta("chan_1", "first ", stream_id="s1") + + await channel.send_delta( + "chan_1", + "boundary ", + stream_id="s1", + stream_end=True, + resuming=True, + merge_next=True, + ) + + assert channel._stream_buffers["s1"] == "first boundary " + + await channel.send_delta("chan_1", "second", stream_id="s1") + await channel.send_delta("chan_1", "", stream_id="s1", stream_end=True) + + posts = [call for call in fake.post_calls if call["path"] == "/api/v4/posts"] + assert len(posts) == 1 + assert posts[0]["json"]["message"] == "first boundary second" + assert "s1" not in channel._stream_buffers + + @pytest.mark.asyncio async def test_stream_end_failure_keeps_buffer_for_retry(): channel, fake = _make_channel() diff --git a/tests/agent/test_loop_progress.py b/tests/agent/test_loop_progress.py index aca35927e..39ecb727f 100644 --- a/tests/agent/test_loop_progress.py +++ b/tests/agent/test_loop_progress.py @@ -580,6 +580,96 @@ class TestToolEventProgress: assert [event.merge_next for event in endings] == [True, False] assert {event.stream_id for event in [*deltas, *endings]} == {deltas[0].stream_id} + @pytest.mark.asyncio + async def test_length_recovery_at_max_iterations_streams_only_missing_tail( + self, + tmp_path: Path, + ) -> None: + bus = MessageBus() + provider = MagicMock() + provider.supports_progress_deltas = True + provider.get_default_model.return_value = "test-model" + + async def chat_stream_with_retry(*, on_content_delta, **kwargs): + await on_content_delta("partial") + return LLMResponse(content="partial", finish_reason="length") + + provider.chat_stream_with_retry = chat_stream_with_retry + provider.chat_with_retry = AsyncMock( + return_value=LLMResponse(content="summary", finish_reason="stop") + ) + loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model") + _attach_webui_runtime_events(loop, bus) + loop.max_iterations = 1 + loop.tools.get_definitions = MagicMock(return_value=[]) + loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] + + await loop._dispatch(InboundMessage( + channel="websocket", + sender_id="u1", + chat_id="chat1", + content="give a long answer", + metadata={"_wants_stream": True}, + )) + + outbound = [] + while bus.outbound_size > 0: + outbound.append(await bus.consume_outbound()) + + deltas = [m.event for m in outbound if isinstance(m.event, StreamDeltaEvent)] + endings = [m.event for m in outbound if isinstance(m.event, StreamEndEvent)] + final = [m for m in outbound if isinstance(m.event, StreamedResponseEvent)] + + assert [event.content for event in deltas] == ["partial", "\n\nsummary"] + assert [event.merge_next for event in endings] == [True, False] + assert {event.stream_id for event in [*deltas, *endings]} == {deltas[0].stream_id} + assert [message.content for message in final] == ["partial\n\nsummary"] + + @pytest.mark.asyncio + async def test_cancelled_length_recovery_closes_merged_stream( + self, + tmp_path: Path, + ) -> None: + bus = MessageBus() + provider = MagicMock() + provider.get_default_model.return_value = "test-model" + loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model") + + async def cancel_after_merge( + _msg: InboundMessage, + *, + on_stream, + on_stream_end, + **_kwargs, + ): + assert on_stream is not None + assert on_stream_end is not None + await on_stream("partial") + await on_stream_end(resuming=True, merge_next=True) + raise asyncio.CancelledError + + loop._process_message = cancel_after_merge # type: ignore[method-assign] + + with pytest.raises(asyncio.CancelledError): + await loop._dispatch(InboundMessage( + channel="websocket", + sender_id="u1", + chat_id="chat1", + content="give a long answer", + metadata={"_wants_stream": True}, + )) + + outbound = [] + while bus.outbound_size > 0: + outbound.append(await bus.consume_outbound()) + + endings = [m.event for m in outbound if isinstance(m.event, StreamEndEvent)] + assert [(event.resuming, event.merge_next) for event in endings] == [ + (True, True), + (False, False), + ] + assert endings[0].stream_id == endings[1].stream_id + @pytest.mark.asyncio async def test_non_streamed_finalization_is_delivered_as_regular_message( self, diff --git a/tests/agent/test_runner_core.py b/tests/agent/test_runner_core.py index 584df875c..883fe24d8 100644 --- a/tests/agent/test_runner_core.py +++ b/tests/agent/test_runner_core.py @@ -482,6 +482,40 @@ async def test_runner_length_recovery_returns_all_segments(): assert provider.chat_with_retry.await_count == 3 +@pytest.mark.asyncio +async def test_runner_length_recovery_preserves_prefix_at_max_iterations(): + """Budget exhaustion must not replace output already produced by recovery.""" + from nanobot.agent.runner import AgentRunner + + provider = MagicMock(spec=LLMProvider) + provider.chat_with_retry = AsyncMock( + return_value=LLMResponse(content="partial answer", finish_reason="length") + ) + tools = MagicMock() + tools.get_definitions.return_value = [] + + runner = AgentRunner() + result = await runner.run(make_run_spec( + provider, + initial_messages=[{"role": "user", "content": "give a long answer"}], + tools=tools, + model="test-model", + max_iterations=1, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + finalize_on_max_iterations=False, + max_iterations_message="limit reached", + )) + + assert result.stop_reason == "max_iterations" + assert result.final_content == "partial answer\n\nlimit reached" + assert result.pending_stream_content == "\n\nlimit reached" + assert [ + message["content"] + for message in result.messages + if message.get("role") == "assistant" + ] == ["partial answer", "limit reached"] + + @pytest.mark.asyncio async def test_runner_length_recovery_does_not_leak_across_tool_calls(): """A recovered prefix belongs only to its contiguous response chain.""" From b55b76d75574d74c1ff5356bace3b76ad12ef399 Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Mon, 27 Jul 2026 01:33:04 +0800 Subject: [PATCH 29/48] fix(streaming): preserve recovered segments across channels --- nanobot/agent/runner.py | 17 +++++++ nanobot/channels/discord/runtime.py | 4 ++ .../discord/tests/test_discord_channel.py | 30 ++++++++++++ nanobot/channels/feishu/runtime.py | 4 ++ .../feishu/tests/test_feishu_streaming.py | 21 +++++++++ nanobot/channels/matrix/runtime.py | 4 ++ .../matrix/tests/test_matrix_channel.py | 23 +++++++++ nanobot/channels/telegram/runtime.py | 4 ++ .../telegram/tests/test_telegram_channel.py | 27 +++++++++++ nanobot/channels/weixin/runtime.py | 4 ++ .../weixin/tests/test_weixin_channel.py | 23 +++++++++ tests/agent/test_loop_progress.py | 47 +++++++++++++++++++ 12 files changed, 208 insertions(+) diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index acc54abbf..671208ce8 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -613,6 +613,23 @@ class AgentRunner: await hook.after_iteration(context) continue + # Some streaming providers recover with a complete response but no + # content deltas. When an earlier length segment is already visible, + # emit this terminal segment into the same stream; otherwise the + # regular full response would duplicate the visible prefix. + if ( + length_recovery_parts + and hook.wants_streaming() + and not context.streamed_content + and response.finish_reason != "error" + and not is_blank_text(clean) + ): + await hook.on_stream( + context, + _restore_outer_whitespace(clean, original_content), + ) + context.streamed_content = True + assistant_message: dict[str, Any] | None = None if response.finish_reason != "error" and not is_blank_text(clean): assistant_message = build_assistant_message( diff --git a/nanobot/channels/discord/runtime.py b/nanobot/channels/discord/runtime.py index b20baca1f..bab06fe24 100644 --- a/nanobot/channels/discord/runtime.py +++ b/nanobot/channels/discord/runtime.py @@ -497,6 +497,10 @@ class DiscordChannel(BaseChannel): self.logger.warning("client not ready; dropping stream delta") return + if stream_end and merge_next: + if not delta: + return + stream_end = False if stream_end: buf = self._stream_bufs.get(chat_id) if not buf or buf.message is None or not buf.text: diff --git a/nanobot/channels/discord/tests/test_discord_channel.py b/nanobot/channels/discord/tests/test_discord_channel.py index d86b56b7c..a749363b0 100644 --- a/nanobot/channels/discord/tests/test_discord_channel.py +++ b/nanobot/channels/discord/tests/test_discord_channel.py @@ -754,6 +754,36 @@ async def test_send_delta_streams_by_editing_message(monkeypatch) -> None: assert owner._stream_bufs == {} +@pytest.mark.asyncio +async def test_send_delta_merge_next_keeps_one_message(monkeypatch) -> None: + owner = DiscordChannel(DiscordConfig(enabled=True, allow_from=["*"]), MessageBus()) + client = _FakeDiscordClient(owner, intents=None) + owner._client = client + owner._running = True + target = _FakeChannel(channel_id=123) + client.channels[123] = target + + times = iter([1.0, 3.0, 5.0]) + monkeypatch.setattr("nanobot.channels.discord.runtime.time.monotonic", lambda: next(times, 5.0)) + + await owner.send_delta( + "123", + "first-", + stream_id="s1", + stream_end=True, + merge_next=True, + ) + await owner.send_delta("123", "second", stream_id="s1") + await owner.send_delta("123", "", stream_id="s1", stream_end=True) + + assert target.sent_payloads == [{"content": "first-"}] + assert target.sent_messages[0].edits == [ + {"content": "first-second"}, + {"content": "first-second"}, + ] + assert owner._stream_bufs == {} + + @pytest.mark.asyncio async def test_send_delta_stream_end_splits_oversized_reply(monkeypatch) -> None: owner = DiscordChannel(DiscordConfig(enabled=True, allow_from=["*"]), MessageBus()) diff --git a/nanobot/channels/feishu/runtime.py b/nanobot/channels/feishu/runtime.py index 1541fd979..a9753e6a1 100644 --- a/nanobot/channels/feishu/runtime.py +++ b/nanobot/channels/feishu/runtime.py @@ -2232,6 +2232,10 @@ class FeishuChannel(BaseChannel): rid_type = "chat_id" if chat_id.startswith("oc_") else "open_id" # --- stream end: final update or fallback --- + if stream_end and merge_next: + if not delta: + return + stream_end = False if stream_end: message_id = meta.get("message_id") # Only finalize the OnIt -> DONE reaction transition on the truly diff --git a/nanobot/channels/feishu/tests/test_feishu_streaming.py b/nanobot/channels/feishu/tests/test_feishu_streaming.py index 7540d5a30..2b67f5a57 100644 --- a/nanobot/channels/feishu/tests/test_feishu_streaming.py +++ b/nanobot/channels/feishu/tests/test_feishu_streaming.py @@ -285,6 +285,27 @@ class TestSendDelta: settings_call = ch._client.cardkit.v1.card.settings.call_args[0][0] assert settings_call.body.sequence == 5 # after final content seq 4 + @pytest.mark.asyncio + async def test_stream_end_merge_next_preserves_buffer(self): + ch = _make_channel() + ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf( + text="first-", + card_id="card_1", + sequence=3, + last_edit=time.monotonic(), + ) + + await ch.send_delta( + "oc_chat1", + "boundary", + stream_end=True, + merge_next=True, + ) + + assert ch._stream_bufs["oc_chat1"].text == "first-boundary" + ch._client.cardkit.v1.card_element.content.assert_not_called() + ch._client.cardkit.v1.card.settings.assert_not_called() + @pytest.mark.asyncio async def test_stream_end_fallback_when_no_card_id(self): """If card creation failed, stream_end falls back to a plain card message.""" diff --git a/nanobot/channels/matrix/runtime.py b/nanobot/channels/matrix/runtime.py index f0cccfb49..992aba7a9 100644 --- a/nanobot/channels/matrix/runtime.py +++ b/nanobot/channels/matrix/runtime.py @@ -602,6 +602,10 @@ class MatrixChannel(BaseChannel): ) -> None: relates_to = self._build_thread_relates_to(metadata) + if stream_end and merge_next: + if not delta: + return + stream_end = False if stream_end: stream_key = _matrix_stream_key(chat_id, stream_id) buf = self._stream_bufs.pop(stream_key, None) diff --git a/nanobot/channels/matrix/tests/test_matrix_channel.py b/nanobot/channels/matrix/tests/test_matrix_channel.py index 58ed088ba..bcc6f3af4 100644 --- a/nanobot/channels/matrix/tests/test_matrix_channel.py +++ b/nanobot/channels/matrix/tests/test_matrix_channel.py @@ -1937,6 +1937,29 @@ async def test_send_delta_stream_end_replaces_existing_message() -> None: } +@pytest.mark.asyncio +async def test_send_delta_merge_next_preserves_buffer() -> None: + channel = MatrixChannel(_make_config(), MessageBus()) + client = _FakeAsyncClient("", "", "", None) + channel.client = client + channel._stream_bufs["!room:matrix.org"] = matrix_module._StreamBuf( + text="first-", + event_id="event-1", + last_edit=100.0, + ) + channel.monotonic_time = lambda: 100.1 + + await channel.send_delta( + "!room:matrix.org", + "boundary", + stream_end=True, + merge_next=True, + ) + + assert channel._stream_bufs["!room:matrix.org"].text == "first-boundary" + assert client.room_send_calls == [] + + @pytest.mark.asyncio async def test_send_delta_keeps_same_room_stream_ids_independent(monkeypatch) -> None: channel = MatrixChannel(_make_config(), MessageBus()) diff --git a/nanobot/channels/telegram/runtime.py b/nanobot/channels/telegram/runtime.py index 4183a82c4..9e42b2df1 100644 --- a/nanobot/channels/telegram/runtime.py +++ b/nanobot/channels/telegram/runtime.py @@ -931,6 +931,10 @@ class TelegramChannel(BaseChannel): meta = metadata or {} int_chat_id = int(chat_id) + if stream_end and merge_next: + if not delta: + return + stream_end = False if stream_end: buf = self._stream_bufs.get(chat_id) if not buf or not buf.message_id or not buf.text: diff --git a/nanobot/channels/telegram/tests/test_telegram_channel.py b/nanobot/channels/telegram/tests/test_telegram_channel.py index ff62713a4..498aa892e 100644 --- a/nanobot/channels/telegram/tests/test_telegram_channel.py +++ b/nanobot/channels/telegram/tests/test_telegram_channel.py @@ -675,6 +675,33 @@ async def test_send_delta_stream_end_raises_and_keeps_buffer_on_failure() -> Non assert "123" in channel._stream_bufs +@pytest.mark.asyncio +async def test_send_delta_merge_next_preserves_buffer() -> None: + channel = TelegramChannel( + TelegramConfig(enabled=True, token="123:abc", allow_from=["*"]), + MessageBus(), + ) + channel._app = _FakeApp(lambda: None) + channel._app.bot.edit_message_text = AsyncMock() + channel._stream_bufs["123"] = _StreamBuf( + text="first-", + message_id=7, + last_edit=float("inf"), + stream_id="s:0", + ) + + await channel.send_delta( + "123", + "boundary", + stream_id="s:0", + stream_end=True, + merge_next=True, + ) + + assert channel._stream_bufs["123"].text == "first-boundary" + channel._app.bot.edit_message_text.assert_not_awaited() + + @pytest.mark.asyncio async def test_send_delta_stream_end_treats_not_modified_as_success() -> None: from telegram.error import BadRequest diff --git a/nanobot/channels/weixin/runtime.py b/nanobot/channels/weixin/runtime.py index 976013c08..ea5d88b59 100644 --- a/nanobot/channels/weixin/runtime.py +++ b/nanobot/channels/weixin/runtime.py @@ -1257,6 +1257,10 @@ class WeixinChannel(BaseChannel): return is_end = stream_end or bool(meta.get("_stream_end")) buffer_key = stream_id or chat_id + if is_end and merge_next: + if delta: + self._stream_buffers.setdefault(buffer_key, []).append(delta) + return # Accumulate intermediate deltas. The stream_end message's own content # (present when the manager coalesces deltas into the end message) is # folded into `full` below instead of appended here, so a send retry diff --git a/nanobot/channels/weixin/tests/test_weixin_channel.py b/nanobot/channels/weixin/tests/test_weixin_channel.py index d0c1c6e41..7058767e5 100644 --- a/nanobot/channels/weixin/tests/test_weixin_channel.py +++ b/nanobot/channels/weixin/tests/test_weixin_channel.py @@ -1824,6 +1824,29 @@ async def test_stream_end_flushes_buffered_answer() -> None: assert "wx-user" not in channel._stream_buffers +@pytest.mark.asyncio +async def test_stream_end_merge_next_preserves_buffer_until_final_end() -> None: + channel, _bus = _make_channel() + channel._client = object() + channel._token = "token" + channel._context_tokens["wx-user"] = "ctx-1" + channel._context_token_at["wx-user"] = time.time() + channel._send_text = AsyncMock() + + await channel.send_delta( + "wx-user", + "first-", + stream_id="s1", + stream_end=True, + merge_next=True, + ) + await channel.send_delta("wx-user", "second", stream_id="s1") + await channel.send_delta("wx-user", "", stream_id="s1", stream_end=True) + + channel._send_text.assert_awaited_once_with("wx-user", "first-second", "ctx-1") + assert "s1" not in channel._stream_buffers + + @pytest.mark.asyncio async def test_stream_end_send_failure_keeps_buffer_for_retry() -> None: channel, _bus = _make_channel() diff --git a/tests/agent/test_loop_progress.py b/tests/agent/test_loop_progress.py index 39ecb727f..220ae646d 100644 --- a/tests/agent/test_loop_progress.py +++ b/tests/agent/test_loop_progress.py @@ -580,6 +580,53 @@ class TestToolEventProgress: assert [event.merge_next for event in endings] == [True, False] assert {event.stream_id for event in [*deltas, *endings]} == {deltas[0].stream_id} + @pytest.mark.asyncio + async def test_length_recovery_streams_non_delta_terminal_segment( + self, + tmp_path: Path, + ) -> None: + bus = MessageBus() + provider = MagicMock() + provider.supports_progress_deltas = True + provider.get_default_model.return_value = "test-model" + call_count = 0 + + async def chat_stream_with_retry(*, on_content_delta, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + await on_content_delta("first-") + return LLMResponse(content="first-", finish_reason="length") + return LLMResponse(content="second", finish_reason="stop") + + provider.chat_stream_with_retry = chat_stream_with_retry + provider.chat_with_retry = AsyncMock() + loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model") + _attach_webui_runtime_events(loop, bus) + loop.tools.get_definitions = MagicMock(return_value=[]) + loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] + + await loop._dispatch(InboundMessage( + channel="websocket", + sender_id="u1", + chat_id="chat1", + content="give a long answer", + metadata={"_wants_stream": True}, + )) + + outbound = [] + while bus.outbound_size > 0: + outbound.append(await bus.consume_outbound()) + + deltas = [m.event for m in outbound if isinstance(m.event, StreamDeltaEvent)] + endings = [m.event for m in outbound if isinstance(m.event, StreamEndEvent)] + final = [m for m in outbound if m.content == "first-second"] + + assert [event.content for event in deltas] == ["first-", "second"] + assert [event.merge_next for event in endings] == [True, False] + assert len(final) == 1 + assert isinstance(final[0].event, StreamedResponseEvent) + @pytest.mark.asyncio async def test_length_recovery_at_max_iterations_streams_only_missing_tail( self, From 15e42059bd7dd9a13a9171a9360cb4e924d91cfd Mon Sep 17 00:00:00 2001 From: shixi-li Date: Thu, 23 Jul 2026 15:33:30 +0800 Subject: [PATCH 30/48] fix(memory): progress past completed no-op batches --- nanobot/agent/memory.py | 9 ++-- nanobot/cli/commands.py | 27 +++++----- nanobot/command/builtin.py | 17 +++---- tests/command/test_builtin_dream.py | 78 ++++++++++++++++++++++++++--- 4 files changed, 98 insertions(+), 33 deletions(-) diff --git a/nanobot/agent/memory.py b/nanobot/agent/memory.py index eb122fd8d..3b126abbd 100644 --- a/nanobot/agent/memory.py +++ b/nanobot/agent/memory.py @@ -47,9 +47,9 @@ class MemoryStore: """Pure file I/O for memory files: MEMORY.md, history.jsonl, SOUL.md, USER.md.""" _DEFAULT_MAX_HISTORY = 1000 - # Durable files whose real working-tree delta grounds Dream commit messages - # and the cursor-advance gate. Deliberately excludes memory/.dream_cursor so - # that advancing the cursor itself is never mistaken for a productive edit. + # Durable files whose real working-tree delta grounds Dream commit messages. + # Deliberately excludes memory/.dream_cursor so progress bookkeeping never + # appears as a durable-memory edit in the audit record. _DREAM_CONTENT_PATHS = ("SOUL.md", "USER.md", "memory/MEMORY.md") # Per-file cap when embedding current contents into the Dream prompt. The # durable files are tiny in practice (~5 KB total), but a runaway file must @@ -586,8 +586,7 @@ class MemoryStore: """Structured summary of uncommitted changes to the durable memory files. Returns "" when git is unavailable or no content file changed. This is - the ground-truth input for diff-grounded Dream commit messages and for - gating cursor advance on real edits (never on LLM self-report). + the ground-truth input for diff-grounded Dream commit messages. """ if not self._git.is_initialized(): return "" diff --git a/nanobot/cli/commands.py b/nanobot/cli/commands.py index a91f408da..4961fcdf7 100644 --- a/nanobot/cli/commands.py +++ b/nanobot/cli/commands.py @@ -1846,20 +1846,23 @@ def _run_gateway( tools=store.build_dream_tools(), on_progress=_silent, ) - # Ground truth: the real file delta, not the LLM's self-report. + # The real file delta grounds the audit record; clean completion + # decides whether this history batch has finished processing. diff_body = store.dream_content_diff() - productive = bool(diff_body) or ( - not store.git.is_initialized() - and MemoryStore.dream_run_completed(resp) - ) - if productive: + completed = MemoryStore.dream_run_completed(resp) + if completed: store.set_last_dream_cursor(last_cursor) - logger.info("Dream cron job completed, cursor advanced to {}", last_cursor) - elif MemoryStore.dream_run_completed(resp): - logger.info( - "Dream cron job completed with no memory changes; " - "cursor not advanced", - ) + if diff_body: + logger.info( + "Dream cron job completed, cursor advanced to {}", + last_cursor, + ) + else: + logger.info( + "Dream cron job completed with no memory changes; " + "cursor advanced to {}", + last_cursor, + ) else: logger.warning( "Dream cron job did not complete; cursor remains at {}", diff --git a/nanobot/command/builtin.py b/nanobot/command/builtin.py index a92890015..672f22020 100644 --- a/nanobot/command/builtin.py +++ b/nanobot/command/builtin.py @@ -437,17 +437,16 @@ async def cmd_dream(ctx: CommandContext) -> OutboundMessage: on_progress=_silent, ) elapsed = time.monotonic() - t0 - # Ground truth: the real file delta, not the LLM's self-report. + # The real file delta grounds the audit record; clean completion + # decides whether this history batch has finished processing. diff_body = store.dream_content_diff() - productive = bool(diff_body) or ( - not store.git.is_initialized() - and MemoryStore.dream_run_completed(resp) - ) - if productive: + completed = MemoryStore.dream_run_completed(resp) + if completed: store.set_last_dream_cursor(last_cursor) - content = f"Dream completed in {elapsed:.1f}s." - elif MemoryStore.dream_run_completed(resp): - content = f"Dream completed in {elapsed:.1f}s; no memory changes." + if diff_body: + content = f"Dream completed in {elapsed:.1f}s." + else: + content = f"Dream completed in {elapsed:.1f}s; no memory changes." else: content = ( f"Dream did not complete after {elapsed:.1f}s; " diff --git a/tests/command/test_builtin_dream.py b/tests/command/test_builtin_dream.py index c3fdcf7a1..df7bbf0e3 100644 --- a/tests/command/test_builtin_dream.py +++ b/tests/command/test_builtin_dream.py @@ -225,7 +225,7 @@ def _build_runnable_dream( @pytest.mark.asyncio async def test_dream_advances_cursor_when_diff_nonempty(tmp_path) -> None: - """A real file delta => productive run => cursor advances (Tier 3).""" + """A completed run with a real file delta advances the cursor.""" ctx, store = _build_runnable_dream(tmp_path, initialized=True, content_diff="SOUL.md: +1 -0") await cmd_dream(ctx) await asyncio.sleep(0) @@ -233,19 +233,83 @@ async def test_dream_advances_cursor_when_diff_nonempty(tmp_path) -> None: @pytest.mark.asyncio -async def test_dream_keeps_cursor_on_completed_noop(tmp_path) -> None: - """Completed run with no file changes must NOT advance the cursor, so the - history batch is reconsidered next run instead of silently swallowed.""" +async def test_dream_advances_cursor_on_completed_noop(tmp_path) -> None: + """A completed no-op has processed the batch and must not repeat it.""" ctx, store = _build_runnable_dream(tmp_path, initialized=True, content_diff="") await cmd_dream(ctx) await asyncio.sleep(0) - assert store._last_dream_cursor == 5 # unchanged + assert store._last_dream_cursor == 42 + assert "no memory changes" in ctx.loop.bus.outbound[0].content + + +@pytest.mark.asyncio +async def test_dream_keeps_cursor_when_incomplete_with_diff(tmp_path) -> None: + """An incomplete run remains retryable even if it left a partial edit.""" + ctx, store = _build_runnable_dream( + tmp_path, + initialized=True, + content_diff="SOUL.md: +1 -0", + stop_reason="length", + ) + await cmd_dream(ctx) + await asyncio.sleep(0) + assert store._last_dream_cursor == 5 + assert "did not complete" in ctx.loop.bus.outbound[0].content + + +@pytest.mark.asyncio +async def test_dream_noop_batch_unlocks_following_history(tmp_path) -> None: + """A no-op first batch must not starve later history entries.""" + workspace = tmp_path / "workspace" + workspace.mkdir() + store = MemoryStore(workspace) + store.write_soul("# Soul") + store.write_memory("# Memory") + for index in range(1, 22): + store.append_history(f"entry-{index:02d}") + store.git.init() + + processed_prompts: list[str] = [] + + async def process_direct(prompt, *args, **kwargs): + processed_prompts.append(prompt) + return OutboundMessage( + channel="cli", + chat_id="direct", + content="done", + metadata={"_stop_reason": "completed"}, + ) + + msg = InboundMessage(channel="cli", sender_id="u1", chat_id="direct", content="/dream") + bus = _FakeBus() + sessions_dir = tmp_path / "sessions" + sessions_dir.mkdir() + loop = SimpleNamespace( + bus=bus, + context=SimpleNamespace(memory=store, timezone="UTC"), + sessions=SimpleNamespace(sessions_dir=sessions_dir), + process_direct=process_direct, + ) + ctx = CommandContext(msg=msg, session=None, key=msg.session_key, raw="/dream", args="", loop=loop) + + await cmd_dream(ctx) + await asyncio.sleep(0) + + assert len(processed_prompts) == 1 + assert "entry-20" in processed_prompts[0] + assert "entry-21" not in processed_prompts[0] + assert store.get_last_dream_cursor() == 20 + next_result = store.build_dream_prompt() + assert next_result is not None + next_prompt, next_cursor = next_result + assert next_cursor == 21 + assert "entry-21" in next_prompt + assert "entry-01" not in next_prompt @pytest.mark.asyncio async def test_dream_non_git_falls_back_to_completion_gate(tmp_path) -> None: - """Without git there is no diff signal; productivity falls back to the - completion check so non-git workspaces keep working.""" + """Non-git workspaces use the same clean-completion gate.""" ctx, store = _build_runnable_dream( tmp_path, initialized=False, content_diff="", stop_reason="completed", ) From 4e2640f2d2f66cdbed4d917ae75f30a49bd24dda Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Mon, 27 Jul 2026 01:52:06 +0800 Subject: [PATCH 31/48] fix(memory): keep failed Dream batches retryable --- nanobot/agent/memory.py | 34 ++++++++++++++++++++++++++--- nanobot/cli/commands.py | 10 ++++++--- nanobot/command/builtin.py | 13 ++++++----- tests/command/test_builtin_dream.py | 25 +++++++++++++++++++++ 4 files changed, 70 insertions(+), 12 deletions(-) diff --git a/nanobot/agent/memory.py b/nanobot/agent/memory.py index 3b126abbd..5bc4bf3b3 100644 --- a/nanobot/agent/memory.py +++ b/nanobot/agent/memory.py @@ -43,6 +43,26 @@ if TYPE_CHECKING: # MemoryStore — pure file I/O layer # --------------------------------------------------------------------------- + +class DreamRunProgress: + """Track tool failures that make a nominally completed Dream run unsafe to advance.""" + + def __init__(self) -> None: + self.had_tool_errors = False + + async def __call__( + self, + *_args: Any, + tool_events: list[dict[str, Any]] | None = None, + **_kwargs: Any, + ) -> None: + if any( + isinstance(event, dict) and event.get("phase") == "error" + for event in tool_events or () + ): + self.had_tool_errors = True + + class MemoryStore: """Pure file I/O for memory files: MEMORY.md, history.jsonl, SOUL.md, USER.md.""" @@ -635,10 +655,18 @@ class MemoryStore: return tools @staticmethod - def dream_run_completed(resp: object | None) -> bool: - """Return True only when an ephemeral Dream agent turn completed cleanly.""" + def dream_run_completed( + resp: object | None, + *, + had_tool_errors: bool = False, + ) -> bool: + """Return True only when a Dream turn completed without tool failures.""" metadata = getattr(resp, "metadata", None) - return isinstance(metadata, dict) and metadata.get("_stop_reason") == "completed" + return ( + not had_tool_errors + and isinstance(metadata, dict) + and metadata.get("_stop_reason") == "completed" + ) # -- message formatting utility ------------------------------------------ diff --git a/nanobot/cli/commands.py b/nanobot/cli/commands.py index 4961fcdf7..e3db4ee7a 100644 --- a/nanobot/cli/commands.py +++ b/nanobot/cli/commands.py @@ -1824,12 +1824,13 @@ def _run_gateway( # Dream is an internal job — run directly, not through the agent loop. if job.name == "dream": - from nanobot.agent.memory import MemoryStore + from nanobot.agent.memory import DreamRunProgress, MemoryStore dream_session_key = MemoryStore.dream_session_key prune_dream_sessions = MemoryStore.prune_dream_sessions store = agent.context.memory + progress = DreamRunProgress() resp = None diff_body = "" try: @@ -1844,12 +1845,15 @@ def _run_gateway( session_key=key, ephemeral=True, tools=store.build_dream_tools(), - on_progress=_silent, + on_progress=progress, ) # The real file delta grounds the audit record; clean completion # decides whether this history batch has finished processing. diff_body = store.dream_content_diff() - completed = MemoryStore.dream_run_completed(resp) + completed = MemoryStore.dream_run_completed( + resp, + had_tool_errors=progress.had_tool_errors, + ) if completed: store.set_last_dream_cursor(last_cursor) if diff_body: diff --git a/nanobot/command/builtin.py b/nanobot/command/builtin.py index 672f22020..3b3befa82 100644 --- a/nanobot/command/builtin.py +++ b/nanobot/command/builtin.py @@ -404,16 +404,14 @@ async def cmd_dream(ctx: CommandContext) -> OutboundMessage: msg = ctx.msg async def _run_dream(): - async def _silent(*_args, **_kwargs): - pass - - from nanobot.agent.memory import MemoryStore + from nanobot.agent.memory import DreamRunProgress, MemoryStore dream_session_key = MemoryStore.dream_session_key build_dream_commit_message = MemoryStore.build_dream_commit_message prune_dream_sessions = MemoryStore.prune_dream_sessions store = loop.context.memory + progress = DreamRunProgress() content = "" resp = None diff_body = "" @@ -434,13 +432,16 @@ async def cmd_dream(ctx: CommandContext) -> OutboundMessage: session_key=key, ephemeral=True, tools=store.build_dream_tools(), - on_progress=_silent, + on_progress=progress, ) elapsed = time.monotonic() - t0 # The real file delta grounds the audit record; clean completion # decides whether this history batch has finished processing. diff_body = store.dream_content_diff() - completed = MemoryStore.dream_run_completed(resp) + completed = MemoryStore.dream_run_completed( + resp, + had_tool_errors=progress.had_tool_errors, + ) if completed: store.set_last_dream_cursor(last_cursor) if diff_body: diff --git a/tests/command/test_builtin_dream.py b/tests/command/test_builtin_dream.py index df7bbf0e3..b90df244a 100644 --- a/tests/command/test_builtin_dream.py +++ b/tests/command/test_builtin_dream.py @@ -192,6 +192,7 @@ def _build_runnable_dream( initialized: bool, content_diff: str, stop_reason: str = "completed", + tool_error: bool = False, ) -> tuple[CommandContext, _FakeStore]: """Build a /dream ctx whose run is driven by a canned stop reason + diff.""" msg = InboundMessage(channel="cli", sender_id="u1", chat_id="direct", content="/dream") @@ -203,6 +204,15 @@ def _build_runnable_dream( ) async def process_direct(*args, **kwargs): + if tool_error: + await kwargs["on_progress"]( + "", + tool_events=[{ + "phase": "error", + "name": "edit_file", + "error": "edit failed", + }], + ) return OutboundMessage( channel="cli", chat_id="direct", @@ -257,6 +267,21 @@ async def test_dream_keeps_cursor_when_incomplete_with_diff(tmp_path) -> None: assert "did not complete" in ctx.loop.bus.outbound[0].content +@pytest.mark.asyncio +async def test_dream_keeps_cursor_when_completed_after_tool_error(tmp_path) -> None: + """A soft tool failure must not masquerade as a verified no-op.""" + ctx, store = _build_runnable_dream( + tmp_path, + initialized=True, + content_diff="", + tool_error=True, + ) + await cmd_dream(ctx) + await asyncio.sleep(0) + assert store._last_dream_cursor == 5 + assert "did not complete" in ctx.loop.bus.outbound[0].content + + @pytest.mark.asyncio async def test_dream_noop_batch_unlocks_following_history(tmp_path) -> None: """A no-op first batch must not starve later history entries.""" From 7aab7e8830af2686e7ba3a114e8b7d9151d4b341 Mon Sep 17 00:00:00 2001 From: Andrew Khmylov Date: Wed, 22 Jul 2026 19:14:17 +0100 Subject: [PATCH 32/48] feat(agent): make idle compaction scan interval configurable Before this change, idle compaction is triggered every 1 second if the incoming message stream is idle. When triggered, it enumerates all session files, loads and parses them, and then checks their expiration. This becomes too CPU-intensive, especially on low-power devices like Raspberry Pi. It's unlikely that you actually need to compact every second over the long time. This change adds a configurable throttling for idle-compaction. Default behavior is unchanged. --- docs/configuration.md | 6 ++-- nanobot/agent/loop.py | 22 ++++++++++--- nanobot/config/schema.py | 4 +++ tests/agent/test_auto_compact.py | 54 +++++++++++++++++++++++++++++++- 4 files changed, 78 insertions(+), 8 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index 0c94d0108..04699b2aa 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -2158,7 +2158,8 @@ When a user is idle for longer than a configured threshold, nanobot **proactivel { "agents": { "defaults": { - "idleCompactAfterMinutes": 15 + "idleCompactAfterMinutes": 15, + "idleCompactCheckIntervalSeconds": 0 } } } @@ -2167,11 +2168,12 @@ When a user is idle for longer than a configured threshold, nanobot **proactivel | Option | Default | Description | |--------|---------|-------------| | `agents.defaults.idleCompactAfterMinutes` | `15` | Minutes of idle time before auto-compaction starts. Set to `0` to disable. The default is close to a typical LLM KV cache expiry window, so stale sessions get compacted before the user returns. | +| `agents.defaults.idleCompactCheckIntervalSeconds` | `0` | Minimum number of seconds between scans for idle sessions. | `sessionTtlMinutes` remains accepted as a legacy alias for backward compatibility, but `idleCompactAfterMinutes` is the preferred config key going forward. How it works: -1. **Idle detection**: On each idle tick (~1 s), checks all sessions for expiration. +1. **Idle detection**: On each idle tick (~1 s), checks all sessions for expiration, subject to the minimum interval set by `idleCompactCheckIntervalSeconds`. 2. **Background compaction**: Idle sessions summarize the older live prefix via LLM and keep the most recent legal suffix (currently 8 messages). 3. **Summary injection**: When the user returns, the summary is injected as runtime context (one-shot, not persisted) alongside the retained recent suffix. 4. **Restart-safe resume**: The summary is also mirrored into session metadata so it can still be recovered after a process restart. diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index eab6f92e3..84b1f8d77 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -297,6 +297,7 @@ class AgentLoop: runtime_model_publisher: Callable[[str, str | None], None] | None = None, restart_mode: str = "auto", local_trigger_store: Any | None = None, + idle_compact_check_interval_seconds: int = 0, ): from nanobot.config.schema import ToolsConfig @@ -445,6 +446,8 @@ class AgentLoop: consolidator=self.consolidator, session_ttl_minutes=session_ttl_minutes, ) + self._idle_compact_check_interval_s = idle_compact_check_interval_seconds + self._next_idle_compact_check_at = time.monotonic() if model_preset: self.set_model_preset(model_preset, publish_update=False) self._register_default_tools(provider_snapshot_loader=provider_snapshot_loader) @@ -500,6 +503,7 @@ class AgentLoop: unified_session=defaults.unified_session, disabled_skills=defaults.disabled_skills, session_ttl_minutes=defaults.session_ttl_minutes, + idle_compact_check_interval_seconds=defaults.idle_compact_check_interval_seconds, consolidation_ratio=defaults.consolidation_ratio, tools_config=config.tools, model_presets=preset_helpers.configured_model_presets(config), @@ -1081,6 +1085,18 @@ class AgentLoop: logger.error("LLM returned error: {}", (result.final_content or "")[:200]) return result.final_content, result.tools_used, result.messages, result.stop_reason, result.had_injections + def _check_expired_sessions_if_due(self) -> None: + """Scan idle sessions no more often than the configured interval.""" + now = time.monotonic() + if now < self._next_idle_compact_check_at: + return + self._next_idle_compact_check_at = now + self._idle_compact_check_interval_s + self.auto_compact.check_expired( + self._schedule_background, + self.runtime_for_session, + active_session_keys=self._pending_queues.keys(), + ) + async def run(self) -> None: """Run the agent loop, dispatching messages as tasks to stay responsive to /stop.""" self._running = True @@ -1092,11 +1108,7 @@ class AgentLoop: try: msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0) except asyncio.TimeoutError: - self.auto_compact.check_expired( - self._schedule_background, - self.runtime_for_session, - active_session_keys=self._pending_queues.keys(), - ) + self._check_expired_sessions_if_due() continue except asyncio.CancelledError: # Preserve real task cancellation so shutdown can complete cleanly. diff --git a/nanobot/config/schema.py b/nanobot/config/schema.py index 8d264ce6b..452c54067 100644 --- a/nanobot/config/schema.py +++ b/nanobot/config/schema.py @@ -154,6 +154,10 @@ class AgentDefaults(Base): validation_alias=AliasChoices("idleCompactAfterMinutes", "sessionTtlMinutes"), serialization_alias="idleCompactAfterMinutes", ) # Auto-compact idle threshold in minutes (0 = disabled) + idle_compact_check_interval_seconds: int = Field( + default=0, + ge=0, + ) # Minimum interval in seconds between scans for idle sessions consolidation_ratio: float = Field( default=0.5, ge=0.1, diff --git a/tests/agent/test_auto_compact.py b/tests/agent/test_auto_compact.py index ce0f22059..af27ab6ee 100644 --- a/tests/agent/test_auto_compact.py +++ b/tests/agent/test_auto_compact.py @@ -11,7 +11,7 @@ from nanobot.agent.loop import AgentLoop from nanobot.bus.events import InboundMessage from nanobot.bus.queue import MessageBus from nanobot.command import CommandContext -from nanobot.config.schema import AgentDefaults +from nanobot.config.schema import AgentDefaults, Config from nanobot.providers.base import LLMResponse @@ -180,12 +180,64 @@ class TestSessionTTLConfig: assert data["idleCompactAfterMinutes"] == 30 assert "sessionTtlMinutes" not in data + def test_idle_scan_interval_defaults_to_zero(self): + """The default should preserve a scan on every idle tick.""" + defaults = AgentDefaults() + assert defaults.idle_compact_check_interval_seconds == 0 + + def test_idle_scan_interval_uses_camel_case_config_key(self): + """The JSON config should use the standard camelCase alias.""" + defaults = AgentDefaults.model_validate({"idleCompactCheckIntervalSeconds": 10}) + assert defaults.idle_compact_check_interval_seconds == 10 + data = defaults.model_dump(mode="json", by_alias=True) + assert data["idleCompactCheckIntervalSeconds"] == 10 + def test_session_file_cap_is_internal_constant(self): """Session file cap should remain an internal constant, not a config field.""" from nanobot.session.manager import FILE_MAX_MESSAGES assert FILE_MAX_MESSAGES == 2000 +class TestIdleScanThrottling: + """Test scheduling of full idle-session scans.""" + + def test_configured_idle_scan_interval_throttles_checks(self, tmp_path, monkeypatch): + """The configured interval should reach the loop and gate session scans.""" + ticks = iter((1_000.0, 1_000.0, 1_009.999, 1_010.0)) + monkeypatch.setattr("nanobot.agent.loop.time.monotonic", lambda: next(ticks)) + config = Config.model_validate({ + "agents": { + "defaults": { + "workspace": str(tmp_path), + "idleCompactCheckIntervalSeconds": 10, + } + } + }) + provider = MagicMock() + provider.get_default_model.return_value = "test-model" + loop = AgentLoop.from_config(config, provider=provider) + loop.auto_compact.check_expired = MagicMock() + + loop._check_expired_sessions_if_due() + loop.auto_compact.check_expired.assert_called_once() + loop._check_expired_sessions_if_due() + loop.auto_compact.check_expired.assert_called_once() + loop._check_expired_sessions_if_due() + + assert loop.auto_compact.check_expired.call_count == 2 + + def test_default_idle_scan_interval_checks_every_tick(self, tmp_path, monkeypatch): + """The zero default should leave each idle tick eligible to scan.""" + monkeypatch.setattr("nanobot.agent.loop.time.monotonic", lambda: 1_000.0) + loop = _make_loop(tmp_path) + loop.auto_compact.check_expired = MagicMock() + + loop._check_expired_sessions_if_due() + loop._check_expired_sessions_if_due() + + assert loop.auto_compact.check_expired.call_count == 2 + + class TestAgentLoopTTLParam: """Test that AutoCompact receives and stores session_ttl_minutes.""" From 68717937e89dd45787462ea9136a05fbefe0a24b Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Mon, 27 Jul 2026 02:08:03 +0800 Subject: [PATCH 33/48] fix(agent): throttle idle scans by default --- docs/configuration.md | 6 +++--- nanobot/config/schema.py | 2 +- tests/agent/test_auto_compact.py | 10 +++++----- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index 04699b2aa..533e159f3 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -2159,7 +2159,7 @@ When a user is idle for longer than a configured threshold, nanobot **proactivel "agents": { "defaults": { "idleCompactAfterMinutes": 15, - "idleCompactCheckIntervalSeconds": 0 + "idleCompactCheckIntervalSeconds": 60 } } } @@ -2168,12 +2168,12 @@ When a user is idle for longer than a configured threshold, nanobot **proactivel | Option | Default | Description | |--------|---------|-------------| | `agents.defaults.idleCompactAfterMinutes` | `15` | Minutes of idle time before auto-compaction starts. Set to `0` to disable. The default is close to a typical LLM KV cache expiry window, so stale sessions get compacted before the user returns. | -| `agents.defaults.idleCompactCheckIntervalSeconds` | `0` | Minimum number of seconds between scans for idle sessions. | +| `agents.defaults.idleCompactCheckIntervalSeconds` | `60` | Minimum number of seconds between scans for idle sessions. Set to `0` to scan on every idle tick (~1 s). | `sessionTtlMinutes` remains accepted as a legacy alias for backward compatibility, but `idleCompactAfterMinutes` is the preferred config key going forward. How it works: -1. **Idle detection**: On each idle tick (~1 s), checks all sessions for expiration, subject to the minimum interval set by `idleCompactCheckIntervalSeconds`. +1. **Idle detection**: On each idle tick (~1 s), checks whether an idle-session scan is due. By default, the full scan runs at most once per minute. 2. **Background compaction**: Idle sessions summarize the older live prefix via LLM and keep the most recent legal suffix (currently 8 messages). 3. **Summary injection**: When the user returns, the summary is injected as runtime context (one-shot, not persisted) alongside the retained recent suffix. 4. **Restart-safe resume**: The summary is also mirrored into session metadata so it can still be recovered after a process restart. diff --git a/nanobot/config/schema.py b/nanobot/config/schema.py index 452c54067..a36ded81a 100644 --- a/nanobot/config/schema.py +++ b/nanobot/config/schema.py @@ -155,7 +155,7 @@ class AgentDefaults(Base): serialization_alias="idleCompactAfterMinutes", ) # Auto-compact idle threshold in minutes (0 = disabled) idle_compact_check_interval_seconds: int = Field( - default=0, + default=60, ge=0, ) # Minimum interval in seconds between scans for idle sessions consolidation_ratio: float = Field( diff --git a/tests/agent/test_auto_compact.py b/tests/agent/test_auto_compact.py index af27ab6ee..957d04989 100644 --- a/tests/agent/test_auto_compact.py +++ b/tests/agent/test_auto_compact.py @@ -180,10 +180,10 @@ class TestSessionTTLConfig: assert data["idleCompactAfterMinutes"] == 30 assert "sessionTtlMinutes" not in data - def test_idle_scan_interval_defaults_to_zero(self): - """The default should preserve a scan on every idle tick.""" + def test_idle_scan_interval_defaults_to_sixty_seconds(self): + """The config default should avoid scanning all sessions every idle tick.""" defaults = AgentDefaults() - assert defaults.idle_compact_check_interval_seconds == 0 + assert defaults.idle_compact_check_interval_seconds == 60 def test_idle_scan_interval_uses_camel_case_config_key(self): """The JSON config should use the standard camelCase alias.""" @@ -226,8 +226,8 @@ class TestIdleScanThrottling: assert loop.auto_compact.check_expired.call_count == 2 - def test_default_idle_scan_interval_checks_every_tick(self, tmp_path, monkeypatch): - """The zero default should leave each idle tick eligible to scan.""" + def test_zero_idle_scan_interval_checks_every_tick(self, tmp_path, monkeypatch): + """An explicit zero should leave each idle tick eligible to scan.""" monkeypatch.setattr("nanobot.agent.loop.time.monotonic", lambda: 1_000.0) loop = _make_loop(tmp_path) loop.auto_compact.check_expired = MagicMock() From 14e692e40d69e677e3cae68c95efa0a01685a053 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E6=98=8E=E6=8C=AF?= Date: Mon, 22 Jun 2026 11:21:33 +0800 Subject: [PATCH 34/48] feat(dingtalk): add disable_private_chat to reject 1:1 DMs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a `disable_private_chat` config flag (JSON alias `disablePrivateChat`, default False) to the DingTalk channel. When enabled, any non-group (1:1) message is rejected with a Chinese notice directing the user to group chat ("该机器人未开启私聊，请在群聊中与我对话。") before permission/pairing logic runs, so even allowlisted senders are redirected. Group messages are unaffected. Co-Authored-By: Claude --- nanobot/channels/dingtalk/runtime.py | 16 ++++ .../dingtalk/tests/test_dingtalk_channel.py | 79 +++++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/nanobot/channels/dingtalk/runtime.py b/nanobot/channels/dingtalk/runtime.py index 7e990fe02..b1ae38c2b 100644 --- a/nanobot/channels/dingtalk/runtime.py +++ b/nanobot/channels/dingtalk/runtime.py @@ -175,6 +175,7 @@ class DingTalkConfig(Base): allow_remote_media_redirects: bool = False remote_media_redirect_allowed_hosts: list[str] = Field(default_factory=list) group_user_isolation: bool = False # If True, each user in group chat gets their own session + disable_private_chat: bool = False # If True, reject 1:1 DMs with a notice; group chats only class DingTalkChannel(BaseChannel): @@ -750,6 +751,21 @@ class DingTalkChannel(BaseChannel): session_key = None if is_group and self.config.group_user_isolation: session_key = f"{self.name}:group:{conversation_id}:{sender_id}" + + if not is_group and self.config.disable_private_chat: + # Private chat is disabled: reply with a notice and drop the + # message before any permission/pairing logic runs, so even + # allowlisted users are redirected to group chat. + self.logger.info("private chat disabled; rejecting DM from {}", sender_name) + await self.send( + OutboundMessage( + channel=self.name, + chat_id=str(chat_id), + content="该机器人未开启私聊，请在群聊中与我对话。", + ) + ) + return + await self._handle_message( sender_id=sender_id, chat_id=chat_id, diff --git a/nanobot/channels/dingtalk/tests/test_dingtalk_channel.py b/nanobot/channels/dingtalk/tests/test_dingtalk_channel.py index 884cf1682..9004760a8 100644 --- a/nanobot/channels/dingtalk/tests/test_dingtalk_channel.py +++ b/nanobot/channels/dingtalk/tests/test_dingtalk_channel.py @@ -153,6 +153,85 @@ async def test_group_user_isolation_true_separates_sessions() -> None: assert msg1.chat_id == msg2.chat_id == "group:conv123" +@pytest.mark.asyncio +async def test_dm_rejected_when_private_chat_disabled(monkeypatch) -> None: + """With disable_private_chat=True, a 1:1 DM is rejected: nothing reaches the + bus (no session is created) and the bot replies with a notice directing the + user to group chat. Even allowlisted senders are blocked in DMs.""" + config = DingTalkConfig( + client_id="app", + client_secret="secret", + allow_from=["*"], # even allowlisted senders are blocked in DMs + disable_private_chat=True, + ) + bus = MessageBus() + channel = DingTalkChannel(config, bus) + + async def fake_get_token(): + return "test-token" + + monkeypatch.setattr(channel, "_get_access_token", fake_get_token) + channel._http = _FakeHttp() + + await channel._on_message( + "hello", + sender_id="user1", + sender_name="Alice", + conversation_type="1", + ) + + # No inbound message was published -> no session created + assert bus.inbound.empty() + + # A notice was sent back to the DM user via the private-chat API + assert len(channel._http.calls) == 1 + call = channel._http.calls[0] + assert call["url"] == "https://api.dingtalk.com/v1.0/robot/oToMessages/batchSend" + assert call["json"]["msgKey"] == "sampleMarkdown" + assert call["json"]["userIds"] == ["user1"] + assert "该机器人未开启私聊，请在群聊中与我对话。" in call["json"]["msgParam"] + + +@pytest.mark.asyncio +async def test_dm_allowed_when_private_chat_not_disabled() -> None: + """By default (disable_private_chat=False), a 1:1 DM still reaches the bus.""" + config = DingTalkConfig(client_id="app", client_secret="secret", allow_from=["*"]) + bus = MessageBus() + channel = DingTalkChannel(config, bus) + + await channel._on_message( + "hello", + sender_id="user1", + sender_name="Alice", + conversation_type="1", + ) + + msg = await bus.consume_inbound() + assert msg.chat_id == "user1" + assert msg.metadata["conversation_type"] == "1" + + +@pytest.mark.asyncio +async def test_group_message_allowed_when_private_chat_disabled() -> None: + """Disabling private chat must not affect group messages.""" + config = DingTalkConfig( + client_id="app", client_secret="secret", allow_from=["*"], disable_private_chat=True + ) + bus = MessageBus() + channel = DingTalkChannel(config, bus) + + await channel._on_message( + "hello", + sender_id="user1", + sender_name="Alice", + conversation_type="2", + conversation_id="conv123", + ) + + msg = await bus.consume_inbound() + assert msg.chat_id == "group:conv123" + + @pytest.mark.asyncio async def test_group_send_uses_group_messages_api() -> None: config = DingTalkConfig(client_id="app", client_secret="secret", allow_from=["*"]) From 205889f9e02179902b20b1c35dda1b59b1bd03b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E6=98=8E=E6=8C=AF?= Date: Mon, 22 Jun 2026 11:51:13 +0800 Subject: [PATCH 35/48] feat(dingtalk): prefix group replies with sender mention In group chats, prefix the outbound markdown reply with an H1 naming the sender (# @) so the addressed user can spot it in a busy group. Private replies are sent verbatim. Visual only: DingTalk markdown robot messages do not push real @ notifications (that would require staffId plumbing and a different message type). sender_name is read from OutboundMessage.metadata, which the agent loop already propagates from inbound metadata. Co-Authored-By: Claude --- nanobot/channels/dingtalk/runtime.py | 11 +++- .../dingtalk/tests/test_dingtalk_channel.py | 51 +++++++++++++++++++ 2 files changed, 60 insertions(+), 2 deletions(-) diff --git a/nanobot/channels/dingtalk/runtime.py b/nanobot/channels/dingtalk/runtime.py index b1ae38c2b..1a6cb25fb 100644 --- a/nanobot/channels/dingtalk/runtime.py +++ b/nanobot/channels/dingtalk/runtime.py @@ -713,8 +713,15 @@ class DingTalkChannel(BaseChannel): if not token: raise RuntimeError("DingTalk access token unavailable") - if msg.content and msg.content.strip(): - if not await self._send_markdown_text(token, msg.chat_id, msg.content.strip()): + content = msg.content.strip() if msg.content else "" + if content: + # In group chats, prefix the reply with a markdown header naming the + # sender so the addressed user can spot the reply. Visual only — + # DingTalk's markdown robot messages do not push real @ notifications. + sender_name = msg.metadata.get("sender_name") if msg.metadata else None + if msg.chat_id.startswith("group:") and sender_name: + content = f"# @{sender_name}\n\n{content}" + if not await self._send_markdown_text(token, msg.chat_id, content): raise RuntimeError("DingTalk text message was not delivered") for media_ref in msg.media or []: diff --git a/nanobot/channels/dingtalk/tests/test_dingtalk_channel.py b/nanobot/channels/dingtalk/tests/test_dingtalk_channel.py index 9004760a8..73406f825 100644 --- a/nanobot/channels/dingtalk/tests/test_dingtalk_channel.py +++ b/nanobot/channels/dingtalk/tests/test_dingtalk_channel.py @@ -1,4 +1,5 @@ import asyncio +import json import zipfile from io import BytesIO from types import SimpleNamespace @@ -252,6 +253,56 @@ async def test_group_send_uses_group_messages_api() -> None: assert call["json"]["msgKey"] == "sampleMarkdown" +@pytest.mark.asyncio +async def test_group_send_prepends_sender_mention(monkeypatch) -> None: + """Group replies are prefixed with a markdown header naming the sender.""" + config = DingTalkConfig(client_id="app", client_secret="secret", allow_from=["*"]) + channel = DingTalkChannel(config, MessageBus()) + channel._http = _FakeHttp() + + async def _fake_token() -> str: + return "token" + + monkeypatch.setattr(channel, "_get_access_token", _fake_token) + + await channel.send( + OutboundMessage( + channel="dingtalk", + chat_id="group:conv123", + content="hello", + metadata={"sender_name": "Alice"}, + ) + ) + + sent_text = json.loads(channel._http.calls[0]["json"]["msgParam"])["text"] + assert sent_text == "# @Alice\n\nhello" + + +@pytest.mark.asyncio +async def test_private_send_does_not_prepend_mention(monkeypatch) -> None: + """Private replies are sent verbatim, without the sender header.""" + config = DingTalkConfig(client_id="app", client_secret="secret", allow_from=["*"]) + channel = DingTalkChannel(config, MessageBus()) + channel._http = _FakeHttp() + + async def _fake_token() -> str: + return "token" + + monkeypatch.setattr(channel, "_get_access_token", _fake_token) + + await channel.send( + OutboundMessage( + channel="dingtalk", + chat_id="user1", # private chat: no "group:" prefix + content="hello", + metadata={"sender_name": "Alice"}, + ) + ) + + sent_text = json.loads(channel._http.calls[0]["json"]["msgParam"])["text"] + assert sent_text == "hello" + + @pytest.mark.asyncio async def test_handler_uses_voice_recognition_text_when_text_is_empty(monkeypatch) -> None: bus = MessageBus() From 9f3dee0192896471f819b35c3adaa9d41d0bc170 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E6=98=8E=E6=8C=AF?= Date: Mon, 22 Jun 2026 19:12:17 +0800 Subject: [PATCH 36/48] docs(dingtalk): clarify disable_private_chat intent in comments Addresses automated review: document that the guard is an intentional hard group-only switch (allowlisted DMs blocked by design) and that str() guards a None sender_id. Comment-only. Co-Authored-By: Claude --- nanobot/channels/dingtalk/runtime.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/nanobot/channels/dingtalk/runtime.py b/nanobot/channels/dingtalk/runtime.py index 1a6cb25fb..cbba682da 100644 --- a/nanobot/channels/dingtalk/runtime.py +++ b/nanobot/channels/dingtalk/runtime.py @@ -760,14 +760,15 @@ class DingTalkChannel(BaseChannel): session_key = f"{self.name}:group:{conversation_id}:{sender_id}" if not is_group and self.config.disable_private_chat: - # Private chat is disabled: reply with a notice and drop the - # message before any permission/pairing logic runs, so even - # allowlisted users are redirected to group chat. + # Group-only kill switch: drop DMs with a notice *before* any + # allow_from / pairing check, so even allowlisted senders are + # redirected — intentional, this is a hard private-chat guard + # rather than an authorization decision. No session is created. self.logger.info("private chat disabled; rejecting DM from {}", sender_name) await self.send( OutboundMessage( channel=self.name, - chat_id=str(chat_id), + chat_id=str(chat_id), # str() guards a None sender_id content="该机器人未开启私聊，请在群聊中与我对话。", ) ) From addaf2d3fcc321ecbb203b266ded4614557d7b28 Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Mon, 27 Jul 2026 02:26:24 +0800 Subject: [PATCH 37/48] fix(dingtalk): harden group reply sender labels --- nanobot/channels/dingtalk/runtime.py | 27 ++++++-- .../dingtalk/tests/test_dingtalk_channel.py | 62 ++++++++++++++++++- 2 files changed, 82 insertions(+), 7 deletions(-) diff --git a/nanobot/channels/dingtalk/runtime.py b/nanobot/channels/dingtalk/runtime.py index cbba682da..dd3989153 100644 --- a/nanobot/channels/dingtalk/runtime.py +++ b/nanobot/channels/dingtalk/runtime.py @@ -24,6 +24,17 @@ from nanobot.security.network import validate_resolved_url, validate_url_target DINGTALK_MAX_REMOTE_MEDIA_BYTES = 20 * 1024 * 1024 DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS = 3 +_DINGTALK_MARKDOWN_INLINE_SPECIALS = frozenset(r"\`*_{}[]()<>#+-.!|~") +_DINGTALK_SENDER_NAME_MAX_CHARS = 80 + + +def _escape_markdown_sender_name(value: str) -> str: + """Render an untrusted display name as one bounded Markdown-safe line.""" + normalized = " ".join(value.split())[:_DINGTALK_SENDER_NAME_MAX_CHARS] + return "".join( + f"\\{char}" if char in _DINGTALK_MARKDOWN_INLINE_SPECIALS else char + for char in normalized + ) try: from dingtalk_stream import ( @@ -719,8 +730,13 @@ class DingTalkChannel(BaseChannel): # sender so the addressed user can spot the reply. Visual only — # DingTalk's markdown robot messages do not push real @ notifications. sender_name = msg.metadata.get("sender_name") if msg.metadata else None - if msg.chat_id.startswith("group:") and sender_name: - content = f"# @{sender_name}\n\n{content}" + safe_sender_name = ( + _escape_markdown_sender_name(sender_name) + if isinstance(sender_name, str) + else "" + ) + if msg.chat_id.startswith("group:") and safe_sender_name: + content = f"# @{safe_sender_name}\n\n{content}" if not await self._send_markdown_text(token, msg.chat_id, content): raise RuntimeError("DingTalk text message was not delivered") @@ -741,7 +757,7 @@ class DingTalkChannel(BaseChannel): async def _on_message( self, content: str, - sender_id: str, + sender_id: str | None, sender_name: str, conversation_type: str | None = None, conversation_id: str | None = None, @@ -753,6 +769,9 @@ class DingTalkChannel(BaseChannel): """ try: self.logger.info("inbound: {} from {}", content, sender_name) + if not sender_id: + self.logger.warning("dropping DingTalk message without a sender ID") + return is_group = conversation_type == "2" and conversation_id chat_id = f"group:{conversation_id}" if is_group else sender_id session_key = None @@ -768,7 +787,7 @@ class DingTalkChannel(BaseChannel): await self.send( OutboundMessage( channel=self.name, - chat_id=str(chat_id), # str() guards a None sender_id + chat_id=chat_id, content="该机器人未开启私聊，请在群聊中与我对话。", ) ) diff --git a/nanobot/channels/dingtalk/tests/test_dingtalk_channel.py b/nanobot/channels/dingtalk/tests/test_dingtalk_channel.py index 73406f825..2721a646d 100644 --- a/nanobot/channels/dingtalk/tests/test_dingtalk_channel.py +++ b/nanobot/channels/dingtalk/tests/test_dingtalk_channel.py @@ -10,15 +10,15 @@ import pytest # Check optional dingtalk dependencies before running tests try: - from nanobot.channels import dingtalk - DINGTALK_AVAILABLE = getattr(dingtalk, "DINGTALK_AVAILABLE", False) + import nanobot.channels.dingtalk.runtime as dingtalk_module + + DINGTALK_AVAILABLE = dingtalk_module.DINGTALK_AVAILABLE except ImportError: DINGTALK_AVAILABLE = False if not DINGTALK_AVAILABLE: pytest.skip("DingTalk dependencies not installed (dingtalk-stream)", allow_module_level=True) -import nanobot.channels.dingtalk.runtime as dingtalk_module from nanobot.bus.events import OutboundMessage from nanobot.bus.queue import MessageBus from nanobot.channels.dingtalk.runtime import ( @@ -154,6 +154,13 @@ async def test_group_user_isolation_true_separates_sessions() -> None: assert msg1.chat_id == msg2.chat_id == "group:conv123" +def test_disable_private_chat_uses_camel_case_config_key() -> None: + config = DingTalkConfig.model_validate({"disablePrivateChat": True}) + + assert config.disable_private_chat is True + assert config.model_dump(mode="json", by_alias=True)["disablePrivateChat"] is True + + @pytest.mark.asyncio async def test_dm_rejected_when_private_chat_disabled(monkeypatch) -> None: """With disable_private_chat=True, a 1:1 DM is rejected: nothing reaches the @@ -278,6 +285,31 @@ async def test_group_send_prepends_sender_mention(monkeypatch) -> None: assert sent_text == "# @Alice\n\nhello" +@pytest.mark.asyncio +async def test_group_send_escapes_untrusted_sender_name(monkeypatch) -> None: + """A sender nickname cannot inject extra Markdown blocks into the reply.""" + config = DingTalkConfig(client_id="app", client_secret="secret", allow_from=["*"]) + channel = DingTalkChannel(config, MessageBus()) + channel._http = _FakeHttp() + + async def _fake_token() -> str: + return "token" + + monkeypatch.setattr(channel, "_get_access_token", _fake_token) + + await channel.send( + OutboundMessage( + channel="dingtalk", + chat_id="group:conv123", + content="hello", + metadata={"sender_name": "Alice\n# [click](https://evil) *admin*"}, + ) + ) + + sent_text = json.loads(channel._http.calls[0]["json"]["msgParam"])["text"] + assert sent_text == r"# @Alice \# \[click\]\(https://evil\) \*admin\*" + "\n\nhello" + + @pytest.mark.asyncio async def test_private_send_does_not_prepend_mention(monkeypatch) -> None: """Private replies are sent verbatim, without the sender header.""" @@ -303,6 +335,30 @@ async def test_private_send_does_not_prepend_mention(monkeypatch) -> None: assert sent_text == "hello" +@pytest.mark.asyncio +async def test_message_without_sender_id_is_dropped() -> None: + """Malformed inbound events must not publish or attempt an invalid reply.""" + config = DingTalkConfig( + client_id="app", + client_secret="secret", + allow_from=["*"], + disable_private_chat=True, + ) + bus = MessageBus() + channel = DingTalkChannel(config, bus) + channel._http = _FakeHttp() + + await channel._on_message( + "hello", + sender_id=None, + sender_name="Unknown", + conversation_type="1", + ) + + assert bus.inbound.empty() + assert channel._http.calls == [] + + @pytest.mark.asyncio async def test_handler_uses_voice_recognition_text_when_text_is_empty(monkeypatch) -> None: bus = MessageBus() From 2a1f840ce2512bfc756f54744ada7f92c5f3f0fb Mon Sep 17 00:00:00 2001 From: chengyongru Date: Wed, 15 Jul 2026 11:05:55 +0800 Subject: [PATCH 38/48] fix(cli): support Codex OAuth in quick start --- nanobot/cli/onboard.py | 107 ++++++++++++++++++++++++--- tests/agent/test_onboard_logic.py | 118 +++++++++++++++++++++++++++++- 2 files changed, 213 insertions(+), 12 deletions(-) diff --git a/nanobot/cli/onboard.py b/nanobot/cli/onboard.py index abc6c53e2..1185f851e 100644 --- a/nanobot/cli/onboard.py +++ b/nanobot/cli/onboard.py @@ -3,6 +3,7 @@ import asyncio import json import types +from contextlib import suppress from dataclasses import dataclass from functools import lru_cache from typing import Any, Literal, NamedTuple, get_args, get_origin @@ -22,7 +23,7 @@ from nanobot.cli.models import ( get_model_context_limit, get_model_suggestions, ) -from nanobot.config.loader import get_config_path, load_config +from nanobot.config.loader import get_config_path, load_config, resolve_config_env_vars from nanobot.config.schema import Config, ModelPresetConfig console = Console() @@ -44,6 +45,8 @@ class _QuickStartProviderInfo(NamedTuple): default_api_base: str backend: str is_direct: bool + is_oauth: bool + default_model: str class _QuickStartEndpointChoice(NamedTuple): @@ -73,6 +76,7 @@ _BACK_PRESSED = object() # Sentinel value for back navigation _MODEL_PRESET_CACHE: set[str] = set() _QUICK_START_CUSTOM_PROVIDER_CHOICE = "Other OpenAI-compatible" +_QUICK_START_OAUTH_PROVIDERS = {"openai_codex"} _CLEAR_CHOICE = "Clear value" _QUICK_START_MENU_CHOICE = "[Q] Quick Start" @@ -1576,7 +1580,11 @@ def _get_quick_start_provider_info() -> dict[str, _QuickStartProviderInfo]: result: dict[str, _QuickStartProviderInfo] = {} for spec in PROVIDERS: - if spec.name == "custom" or spec.is_oauth or spec.is_transcription_only: + if ( + spec.name == "custom" + or spec.is_transcription_only + or (spec.is_oauth and spec.name not in _QUICK_START_OAUTH_PROVIDERS) + ): continue result[spec.name] = _QuickStartProviderInfo( display_name=spec.display_name or spec.name, @@ -1584,6 +1592,8 @@ def _get_quick_start_provider_info() -> dict[str, _QuickStartProviderInfo]: default_api_base=spec.default_api_base, backend=spec.backend, is_direct=spec.is_direct, + is_oauth=spec.is_oauth, + default_model=spec.builtin_models[0].id if spec.builtin_models else "", ) return result @@ -1599,7 +1609,64 @@ def _get_quick_start_provider_choices() -> dict[str, str]: def _quick_start_requires_api_key(provider_name: str, info: _QuickStartProviderInfo | None) -> bool: """Return whether Quick Start should ask for an API key.""" - return provider_name == "custom" or not (info and info.is_local) + return provider_name == "custom" or not (info and (info.is_local or info.is_oauth)) + + +def _quick_start_oauth_login(config: Config, provider_name: str) -> bool: + """Authenticate an OAuth provider supported by Quick Start.""" + if provider_name != "openai_codex": + console.print(f"[red]OAuth login is not supported for {provider_name}[/red]") + return False + + try: + from oauth_cli_kit import get_token, login_oauth_interactive + except ImportError: + console.print("[red]oauth_cli_kit not installed. Run: pip install oauth-cli-kit[/red]") + return False + + try: + proxy = resolve_config_env_vars(config).providers.openai_codex.proxy or None + except ValueError as exc: + console.print(f"[red]{exc}[/red]") + return False + + token = None + with suppress(Exception): + token = get_token(proxy=proxy) + if not (token and token.access): + console.print("[cyan]Starting interactive OAuth login...[/cyan]\n") + try: + token = login_oauth_interactive( + print_fn=lambda message: console.print(message), + prompt_fn=lambda prompt: _get_questionary().text(prompt).ask() or "", + proxy=proxy, + ) + except Exception as exc: + console.print(f"[red]OAuth login failed: {exc}[/red]") + return False + + if not (token and token.access): + console.print("[red]OAuth login failed[/red]") + return False + + account = getattr(token, "account_id", None) + suffix = f" [dim]{account}[/dim]" if account else "" + console.print(f"[green]Authenticated with OpenAI Codex[/green]{suffix}") + return True + + +def _quick_start_oauth_is_authenticated(config: Config, provider_name: str) -> bool: + """Return whether Quick Start can load a usable OAuth token.""" + if provider_name != "openai_codex": + return False + try: + from oauth_cli_kit import get_token + + proxy = resolve_config_env_vars(config).providers.openai_codex.proxy or None + token = get_token(proxy=proxy) + except Exception: + return False + return bool(token and token.access) def _quick_start_requires_base_url(provider_name: str, info: _QuickStartProviderInfo | None) -> bool: @@ -1710,7 +1777,11 @@ def _configure_quick_start_provider(config: Config) -> bool | object: console.print(f"[red]Unknown provider: {provider_name}[/red]") return False - model = _input_model_with_autocomplete("Model ID", "", provider_name) + model = _input_model_with_autocomplete( + "Model ID", + provider_info.default_model if provider_info else "", + provider_name, + ) if model is _BACK_PRESSED: continue model = (model or "").strip() @@ -1718,6 +1789,10 @@ def _configure_quick_start_provider(config: Config) -> bool | object: console.print("[yellow]! Model ID is required for Quick Start[/yellow]") return False + if provider_info and provider_info.is_oauth: + if not _quick_start_oauth_login(config, provider_name): + return False + if api_key is not None: provider_config.api_key = api_key if api_base: @@ -1784,17 +1859,27 @@ def _show_quick_start_summary(config: Config) -> None: _show_quick_start_progress(3) preset = config.model_presets.get("primary") provider_label = "AI provider" - has_api_key = True + credentials_ready = True + credential_name = "API key" if preset: provider_config = getattr(config.providers, preset.provider, None) - provider_label, _is_gateway, is_local, _api_base = _get_provider_info().get( - preset.provider, (preset.provider, False, False, "") - ) - has_api_key = is_local or bool(provider_config and provider_config.api_key) + provider_info = _get_quick_start_provider_info().get(preset.provider) + if provider_info: + provider_label = provider_info.display_name + if provider_info.is_oauth: + credential_name = "OAuth login" + credentials_ready = _quick_start_oauth_is_authenticated(config, preset.provider) + else: + credentials_ready = provider_info.is_local or bool( + provider_config and provider_config.api_key + ) + else: + provider_label = _get_provider_names().get(preset.provider, preset.provider) + credentials_ready = bool(provider_config and provider_config.api_key) status = "Ready" - if not has_api_key: - status = f"{provider_label} API key missing" + if not credentials_ready: + status = f"{provider_label} {credential_name} missing" rows = [ ("Status", status), diff --git a/tests/agent/test_onboard_logic.py b/tests/agent/test_onboard_logic.py index 633b34e31..e11741880 100644 --- a/tests/agent/test_onboard_logic.py +++ b/tests/agent/test_onboard_logic.py @@ -978,7 +978,14 @@ class TestMainMenuUpdate: expected_provider_names = set() seen_display_names: set[str] = set() for spec in PROVIDERS: - if spec.name == "custom" or spec.is_oauth or spec.is_transcription_only: + if ( + spec.name == "custom" + or spec.is_transcription_only + or ( + spec.is_oauth + and spec.name not in onboard_wizard._QUICK_START_OAUTH_PROVIDERS + ) + ): continue if spec.display_name in seen_display_names: continue @@ -988,9 +995,118 @@ class TestMainMenuUpdate: assert selected_provider_names == expected_provider_names assert "assemblyai" not in selected_provider_names + assert choices["OpenAI Codex"] == "openai_codex" + assert "github_copilot" not in selected_provider_names assert choices["OpenCode Zen"] == "opencode" assert choices[onboard_wizard._QUICK_START_CUSTOM_PROVIDER_CHOICE] == "custom" + def test_quick_start_openai_codex_uses_oauth_and_default_model(self, monkeypatch): + """Codex should authenticate without asking for an API key.""" + config = Config() + oauth_calls: list[tuple[Config, str]] = [] + model_prompts: list[tuple[str, str, str]] = [] + + monkeypatch.setattr(onboard_wizard, "_show_quick_start_progress", lambda *_args: None) + monkeypatch.setattr( + onboard_wizard, + "_select_with_back", + lambda *args, **kwargs: "OpenAI Codex", + ) + + def fail_api_key_prompt(*_args, **_kwargs): + raise AssertionError("OpenAI Codex Quick Start should not ask for an API key") + + def fake_model_input(prompt, current, provider): + model_prompts.append((prompt, current, provider)) + return current + + monkeypatch.setattr(onboard_wizard, "_input_text", fail_api_key_prompt) + monkeypatch.setattr(onboard_wizard, "_input_model_with_autocomplete", fake_model_input) + monkeypatch.setattr( + onboard_wizard, + "_quick_start_oauth_login", + lambda selected_config, provider: oauth_calls.append( + (selected_config, provider) + ) + or True, + ) + + assert onboard_wizard._configure_quick_start_provider(config) is True + + assert oauth_calls == [(config, "openai_codex")] + assert model_prompts == [ + ("Model ID", "openai-codex/gpt-5.6-sol", "openai_codex") + ] + assert config.providers.openai_codex.api_key is None + assert config.model_presets["primary"].provider == "openai_codex" + assert config.model_presets["primary"].model == "openai-codex/gpt-5.6-sol" + + def test_quick_start_openai_codex_login_failure_does_not_create_preset(self, monkeypatch): + """A failed Codex login must not leave a ready-looking model preset.""" + config = Config() + + monkeypatch.setattr(onboard_wizard, "_show_quick_start_progress", lambda *_args: None) + monkeypatch.setattr( + onboard_wizard, + "_select_with_back", + lambda *args, **kwargs: "OpenAI Codex", + ) + monkeypatch.setattr( + onboard_wizard, + "_input_model_with_autocomplete", + lambda *args, **kwargs: "openai-codex/gpt-5.6-sol", + ) + monkeypatch.setattr(onboard_wizard, "_quick_start_oauth_login", lambda *args: False) + + assert onboard_wizard._configure_quick_start_provider(config) is False + assert "primary" not in config.model_presets + + def test_quick_start_openai_codex_login_reuses_existing_token(self, monkeypatch): + """Quick Start should not open a new login flow when Codex is already authenticated.""" + import oauth_cli_kit + + config = Config() + token = SimpleNamespace(access="existing-token", account_id="account-123") + login_calls: list[object] = [] + + monkeypatch.setattr(oauth_cli_kit, "get_token", lambda **kwargs: token) + monkeypatch.setattr( + oauth_cli_kit, + "login_oauth_interactive", + lambda **kwargs: login_calls.append(kwargs), + ) + monkeypatch.setattr(onboard_wizard.console, "print", lambda *args, **kwargs: None) + + assert onboard_wizard._quick_start_oauth_login(config, "openai_codex") is True + assert login_calls == [] + + def test_quick_start_summary_reports_missing_codex_oauth(self, monkeypatch): + """The review step should distinguish OAuth from an API-key setup.""" + config = Config() + config.model_presets["primary"] = ModelPresetConfig( + model="openai-codex/gpt-5.6-sol", + provider="openai_codex", + ) + captured: dict[str, list[tuple[str, str]]] = {} + + monkeypatch.setattr(onboard_wizard, "_show_quick_start_progress", lambda *_args: None) + monkeypatch.setattr( + onboard_wizard, + "_quick_start_oauth_is_authenticated", + lambda *args: False, + ) + monkeypatch.setattr( + onboard_wizard, + "_print_summary_panel", + lambda rows, _title: captured.setdefault("rows", rows), + ) + + onboard_wizard._show_quick_start_summary(config) + + rows = dict(captured["rows"]) + assert rows["Status"] == "OpenAI Codex OAuth login missing" + assert rows["WebSocket channel"] == "enabled" + def test_quick_start_provider_choice_skips_advanced_prompts(self, monkeypatch): """The beginner path should ask for provider credentials and model.""" config = Config() From a4ec83fb0da0d9b7c8038b93a7e8ca916d0704af Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Sun, 26 Jul 2026 17:37:23 +0800 Subject: [PATCH 39/48] fix(cli): scope Codex proxy env resolution --- nanobot/cli/onboard.py | 11 +++++++++-- tests/agent/test_onboard_logic.py | 32 ++++++++++++++++++++++++++++++- 2 files changed, 40 insertions(+), 3 deletions(-) diff --git a/nanobot/cli/onboard.py b/nanobot/cli/onboard.py index 1185f851e..7104ae693 100644 --- a/nanobot/cli/onboard.py +++ b/nanobot/cli/onboard.py @@ -1612,6 +1612,13 @@ def _quick_start_requires_api_key(provider_name: str, info: _QuickStartProviderI return provider_name == "custom" or not (info and (info.is_local or info.is_oauth)) +def _quick_start_codex_proxy(config: Config) -> str | None: + """Resolve only the Codex proxy without validating unrelated provider secrets.""" + proxy_config = Config() + proxy_config.providers.openai_codex.proxy = config.providers.openai_codex.proxy + return resolve_config_env_vars(proxy_config).providers.openai_codex.proxy or None + + def _quick_start_oauth_login(config: Config, provider_name: str) -> bool: """Authenticate an OAuth provider supported by Quick Start.""" if provider_name != "openai_codex": @@ -1625,7 +1632,7 @@ def _quick_start_oauth_login(config: Config, provider_name: str) -> bool: return False try: - proxy = resolve_config_env_vars(config).providers.openai_codex.proxy or None + proxy = _quick_start_codex_proxy(config) except ValueError as exc: console.print(f"[red]{exc}[/red]") return False @@ -1662,7 +1669,7 @@ def _quick_start_oauth_is_authenticated(config: Config, provider_name: str) -> b try: from oauth_cli_kit import get_token - proxy = resolve_config_env_vars(config).providers.openai_codex.proxy or None + proxy = _quick_start_codex_proxy(config) token = get_token(proxy=proxy) except Exception: return False diff --git a/tests/agent/test_onboard_logic.py b/tests/agent/test_onboard_logic.py index e11741880..88947862c 100644 --- a/tests/agent/test_onboard_logic.py +++ b/tests/agent/test_onboard_logic.py @@ -1066,10 +1066,19 @@ class TestMainMenuUpdate: import oauth_cli_kit config = Config() + config.providers.openai.api_key = "${UNRELATED_MISSING_KEY}" + config.providers.openai_codex.proxy = "${CODEX_PROXY}" token = SimpleNamespace(access="existing-token", account_id="account-123") + token_proxies: list[str | None] = [] login_calls: list[object] = [] - monkeypatch.setattr(oauth_cli_kit, "get_token", lambda **kwargs: token) + monkeypatch.setenv("CODEX_PROXY", "http://127.0.0.1:8080") + monkeypatch.delenv("UNRELATED_MISSING_KEY", raising=False) + monkeypatch.setattr( + oauth_cli_kit, + "get_token", + lambda **kwargs: token_proxies.append(kwargs.get("proxy")) or token, + ) monkeypatch.setattr( oauth_cli_kit, "login_oauth_interactive", @@ -1078,7 +1087,28 @@ class TestMainMenuUpdate: monkeypatch.setattr(onboard_wizard.console, "print", lambda *args, **kwargs: None) assert onboard_wizard._quick_start_oauth_login(config, "openai_codex") is True + assert token_proxies == ["http://127.0.0.1:8080"] assert login_calls == [] + assert config.providers.openai.api_key == "${UNRELATED_MISSING_KEY}" + assert config.providers.openai_codex.proxy == "${CODEX_PROXY}" + + def test_quick_start_codex_auth_check_ignores_unrelated_missing_env(self, monkeypatch): + """OAuth readiness should depend only on the Codex proxy and token.""" + import oauth_cli_kit + + config = Config() + config.providers.anthropic.api_key = "${UNRELATED_MISSING_KEY}" + monkeypatch.delenv("UNRELATED_MISSING_KEY", raising=False) + monkeypatch.setattr( + oauth_cli_kit, + "get_token", + lambda **kwargs: SimpleNamespace(access="existing-token"), + ) + + assert ( + onboard_wizard._quick_start_oauth_is_authenticated(config, "openai_codex") + is True + ) def test_quick_start_summary_reports_missing_codex_oauth(self, monkeypatch): """The review step should distinguish OAuth from an API-key setup.""" From b695a7e87558785d319ed5eb9797a63852a8a2d0 Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Mon, 27 Jul 2026 02:43:33 +0800 Subject: [PATCH 40/48] fix(cli): harden quick start OAuth handling --- nanobot/cli/onboard.py | 15 ++++---- tests/agent/test_onboard_logic.py | 64 +++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 7 deletions(-) diff --git a/nanobot/cli/onboard.py b/nanobot/cli/onboard.py index 7104ae693..226cc3fcf 100644 --- a/nanobot/cli/onboard.py +++ b/nanobot/cli/onboard.py @@ -15,6 +15,7 @@ except ModuleNotFoundError: # pragma: no cover - exercised in environments with from loguru import logger from pydantic import BaseModel from rich.console import Console +from rich.markup import escape from rich.panel import Panel from rich.table import Table @@ -1634,30 +1635,30 @@ def _quick_start_oauth_login(config: Config, provider_name: str) -> bool: try: proxy = _quick_start_codex_proxy(config) except ValueError as exc: - console.print(f"[red]{exc}[/red]") + console.print(f"[red]{escape(str(exc))}[/red]") return False token = None with suppress(Exception): token = get_token(proxy=proxy) - if not (token and token.access): + if not getattr(token, "access", None): console.print("[cyan]Starting interactive OAuth login...[/cyan]\n") try: token = login_oauth_interactive( - print_fn=lambda message: console.print(message), + print_fn=lambda message: console.print(message, markup=False), prompt_fn=lambda prompt: _get_questionary().text(prompt).ask() or "", proxy=proxy, ) except Exception as exc: - console.print(f"[red]OAuth login failed: {exc}[/red]") + console.print(f"[red]OAuth login failed: {escape(str(exc))}[/red]") return False - if not (token and token.access): + if not getattr(token, "access", None): console.print("[red]OAuth login failed[/red]") return False account = getattr(token, "account_id", None) - suffix = f" [dim]{account}[/dim]" if account else "" + suffix = f" [dim]{escape(str(account))}[/dim]" if account else "" console.print(f"[green]Authenticated with OpenAI Codex[/green]{suffix}") return True @@ -1673,7 +1674,7 @@ def _quick_start_oauth_is_authenticated(config: Config, provider_name: str) -> b token = get_token(proxy=proxy) except Exception: return False - return bool(token and token.access) + return bool(getattr(token, "access", None)) def _quick_start_requires_base_url(provider_name: str, info: _QuickStartProviderInfo | None) -> bool: diff --git a/tests/agent/test_onboard_logic.py b/tests/agent/test_onboard_logic.py index 88947862c..7a28133ec 100644 --- a/tests/agent/test_onboard_logic.py +++ b/tests/agent/test_onboard_logic.py @@ -1092,6 +1092,55 @@ class TestMainMenuUpdate: assert config.providers.openai.api_key == "${UNRELATED_MISSING_KEY}" assert config.providers.openai_codex.proxy == "${CODEX_PROXY}" + def test_quick_start_openai_codex_runs_interactive_login_for_bad_cached_token( + self, monkeypatch + ): + """A malformed cached token should fall back to the interactive OAuth flow.""" + import oauth_cli_kit + + config = Config() + config.providers.openai_codex.proxy = "http://127.0.0.1:8080" + prompts: list[str] = [] + printed: list[tuple[tuple[object, ...], dict[str, object]]] = [] + + class FakePrompt: + def ask(self): + return "authorization-code" + + def fake_login(**kwargs): + kwargs["print_fn"]("[bold]Open the browser[/bold]") + prompts.append(kwargs["prompt_fn"]("Paste the authorization code")) + assert kwargs["proxy"] == "http://127.0.0.1:8080" + return SimpleNamespace( + access="fresh-token", + account_id="[red]account-123[/red]", + ) + + monkeypatch.setattr( + oauth_cli_kit, + "get_token", + lambda **_kwargs: SimpleNamespace(account_id="missing-access"), + ) + monkeypatch.setattr(oauth_cli_kit, "login_oauth_interactive", fake_login) + monkeypatch.setattr( + onboard_wizard, + "_get_questionary", + lambda: SimpleNamespace(text=lambda *_args, **_kwargs: FakePrompt()), + ) + monkeypatch.setattr( + onboard_wizard.console, + "print", + lambda *args, **kwargs: printed.append((args, kwargs)), + ) + + assert onboard_wizard._quick_start_oauth_login(config, "openai_codex") is True + assert prompts == ["authorization-code"] + assert any( + args == ("[bold]Open the browser[/bold]",) and kwargs == {"markup": False} + for args, kwargs in printed + ) + assert any(r"\[red]account-123\[/red]" in str(args[0]) for args, _kwargs in printed) + def test_quick_start_codex_auth_check_ignores_unrelated_missing_env(self, monkeypatch): """OAuth readiness should depend only on the Codex proxy and token.""" import oauth_cli_kit @@ -1110,6 +1159,21 @@ class TestMainMenuUpdate: is True ) + def test_quick_start_codex_auth_check_rejects_malformed_token(self, monkeypatch): + """A malformed cached token should report not-ready instead of crashing.""" + import oauth_cli_kit + + monkeypatch.setattr( + oauth_cli_kit, + "get_token", + lambda **_kwargs: SimpleNamespace(account_id="missing-access"), + ) + + assert ( + onboard_wizard._quick_start_oauth_is_authenticated(Config(), "openai_codex") + is False + ) + def test_quick_start_summary_reports_missing_codex_oauth(self, monkeypatch): """The review step should distinguish OAuth from an API-key setup.""" config = Config() From 4986590bd7e8afa1250f8875632cfc641c0a1a58 Mon Sep 17 00:00:00 2001 From: stupidloud Date: Thu, 2 Jul 2026 22:33:45 +0800 Subject: [PATCH 41/48] fix(image): pass aspect ratio and size to Gemini Flash image models The Gemini Flash image path (`generateContent`) dropped both `aspect_ratio` and `image_size`: `generate()` never forwarded them and `_generate_gemini_flash` did not accept them, so every request fell back to 1:1 / input-matched output. The Imagen path was unaffected. Forward the hints and emit them under `generationConfig.responseFormat.image` per the current Gemini API. Aspect ratio is validated against the accepted set; `imageSize` is validated against {512,1K,2K,4K} and only sent to Gemini 3+ image models, since `gemini-2.5-flash-image` supports only `aspectRatio`. Co-Authored-By: Claude Opus 4.8 (1M context) --- nanobot/providers/image_generation.py | 47 ++++++++++++++++++- tests/providers/test_image_generation.py | 57 ++++++++++++++++++++++++ 2 files changed, 102 insertions(+), 2 deletions(-) diff --git a/nanobot/providers/image_generation.py b/nanobot/providers/image_generation.py index 94b9a75e9..0d27f04e0 100644 --- a/nanobot/providers/image_generation.py +++ b/nanobot/providers/image_generation.py @@ -33,6 +33,13 @@ _AIHUBMIX_ASPECT_RATIO_SIZES = { } _GEMINI_DEFAULT_TIMEOUT_S = 120.0 _GEMINI_IMAGEN_ASPECT_RATIOS = {"1:1", "9:16", "16:9", "3:4", "4:3"} +# Aspect ratios accepted by the Gemini Flash image (generateContent) models. +_GEMINI_FLASH_ASPECT_RATIOS = { + "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", + "9:16", "16:9", "21:9", "1:4", "4:1", "1:8", "8:1", +} +# Image-size tokens accepted by Gemini 3+ image models (2.5 Flash Image ignores it). +_GEMINI_FLASH_IMAGE_SIZES = {"512", "1K", "2K", "4K"} _OLLAMA_DEFAULT_SIDE = 1024 _OLLAMA_SIZE_PRESETS = { "1K": 1024, @@ -635,7 +642,11 @@ class GeminiImageGenerationClient(ImageGenerationProvider): prompt=prompt, model=model, aspect_ratio=aspect_ratio ) return await self._generate_gemini_flash( - prompt=prompt, model=model, reference_images=reference_images or [] + prompt=prompt, + model=model, + reference_images=reference_images or [], + aspect_ratio=aspect_ratio, + image_size=image_size, ) async def _generate_imagen( @@ -691,15 +702,22 @@ class GeminiImageGenerationClient(ImageGenerationProvider): prompt: str, model: str, reference_images: list[str], + aspect_ratio: str | None = None, + image_size: str | None = None, ) -> GeneratedImageResponse: parts: list[dict[str, Any]] = [ {"inlineData": image_path_to_inline_data(path)} for path in reference_images ] parts.append({"text": prompt}) + generation_config: dict[str, Any] = {"responseModalities": ["TEXT", "IMAGE"]} + image_config = _gemini_flash_image_config(model, aspect_ratio, image_size) + if image_config: + generation_config["responseFormat"] = {"image": image_config} + body: dict[str, Any] = { "contents": [{"role": "user", "parts": parts}], - "generationConfig": {"responseModalities": ["TEXT", "IMAGE"]}, + "generationConfig": generation_config, } body.update(self.extra_body) @@ -748,6 +766,31 @@ class GeminiImageGenerationClient(ImageGenerationProvider): ) +def _gemini_flash_image_config( + model: str, + aspect_ratio: str | None, + image_size: str | None, +) -> dict[str, str]: + """Build the ``responseFormat.image`` config for Gemini Flash image models. + + Aspect ratio applies to all Flash image models; image size is only honored + by Gemini 3+ image models (``gemini-2.5-flash-image`` ignores it). + """ + config: dict[str, str] = {} + if aspect_ratio and aspect_ratio in _GEMINI_FLASH_ASPECT_RATIOS: + config["aspectRatio"] = aspect_ratio + if image_size and _gemini_flash_supports_image_size(model): + normalized = image_size.strip().upper() + if normalized in _GEMINI_FLASH_IMAGE_SIZES: + config["imageSize"] = normalized + return config + + +def _gemini_flash_supports_image_size(model: str) -> bool: + """Return whether the model honors ``imageSize`` (Gemini 3+ image models).""" + return "2.5" not in model.lower() + + async def _aihubmix_images_from_payload( client: httpx.AsyncClient, payload: dict[str, Any], diff --git a/tests/providers/test_image_generation.py b/tests/providers/test_image_generation.py index cbb0f99ad..8e3110498 100644 --- a/tests/providers/test_image_generation.py +++ b/tests/providers/test_image_generation.py @@ -422,6 +422,63 @@ async def test_gemini_flash_reference_images(tmp_path: Path) -> None: assert parts[1] == {"text": "edit this"} +def _gemini_flash_image_response() -> FakeResponse: + return FakeResponse( + { + "candidates": [ + {"content": {"parts": [{"inlineData": {"mimeType": "image/png", "data": RAW_B64}}]}} + ] + } + ) + + +@pytest.mark.asyncio +async def test_gemini_flash_forwards_aspect_ratio_and_image_size() -> None: + fake = FakeClient(_gemini_flash_image_response()) + client = GeminiImageGenerationClient(api_key="AIza-test", client=fake) # type: ignore[arg-type] + + await client.generate( + prompt="draw a cat", + model="gemini-3-pro-image", + aspect_ratio="16:9", + image_size="2K", + ) + + image_config = fake.calls[0]["json"]["generationConfig"]["responseFormat"]["image"] + assert image_config == {"aspectRatio": "16:9", "imageSize": "2K"} + + +@pytest.mark.asyncio +async def test_gemini_flash_2_5_drops_image_size() -> None: + fake = FakeClient(_gemini_flash_image_response()) + client = GeminiImageGenerationClient(api_key="AIza-test", client=fake) # type: ignore[arg-type] + + await client.generate( + prompt="draw a cat", + model="gemini-2.5-flash-image", + aspect_ratio="4:3", + image_size="1K", + ) + + image_config = fake.calls[0]["json"]["generationConfig"]["responseFormat"]["image"] + assert image_config == {"aspectRatio": "4:3"} + + +@pytest.mark.asyncio +async def test_gemini_flash_ignores_unsupported_hints() -> None: + fake = FakeClient(_gemini_flash_image_response()) + client = GeminiImageGenerationClient(api_key="AIza-test", client=fake) # type: ignore[arg-type] + + await client.generate( + prompt="draw a cat", + model="gemini-3-pro-image", + aspect_ratio="7:5", + image_size="1024x1024", + ) + + assert "responseFormat" not in fake.calls[0]["json"]["generationConfig"] + + @pytest.mark.asyncio async def test_gemini_requires_api_key() -> None: client = GeminiImageGenerationClient(api_key=None) From ef445cc2466e060de3cafc30ad922cebd64f502f Mon Sep 17 00:00:00 2001 From: stupidloud Date: Thu, 2 Jul 2026 23:23:39 +0800 Subject: [PATCH 42/48] fix(image): narrow Gemini Flash aspect-ratio and image-size scoping Address review feedback that the capability checks were broader than the documented per-model matrix: - Drop the extreme aspect ratios (1:4, 4:1, 1:8, 8:1) from the Flash allow-list. They are only documented for 3.1 Flash / Flash Lite, so the global set could send an unsupported ratio to 2.5 Flash Image or 3.1 Pro Image. Keep the ratios common to every Flash image model. - Identify imageSize support positively via "gemini-3" instead of excluding "2.5". The old predicate also matched gemini-2.0-flash-preview-image- generation, which (with the default 1K size) altered that model's request shape even though only Gemini 3+ image models accept a configurable size. Add tests for the gemini-2.0 image-size drop and the extreme-ratio drop. Co-Authored-By: Claude Opus 4.8 (1M context) --- nanobot/providers/image_generation.py | 20 ++++++++++++++------ tests/providers/test_image_generation.py | 22 ++++++++++++++++++++-- 2 files changed, 34 insertions(+), 8 deletions(-) diff --git a/nanobot/providers/image_generation.py b/nanobot/providers/image_generation.py index 0d27f04e0..f9d76e288 100644 --- a/nanobot/providers/image_generation.py +++ b/nanobot/providers/image_generation.py @@ -33,12 +33,15 @@ _AIHUBMIX_ASPECT_RATIO_SIZES = { } _GEMINI_DEFAULT_TIMEOUT_S = 120.0 _GEMINI_IMAGEN_ASPECT_RATIOS = {"1:1", "9:16", "16:9", "3:4", "4:3"} -# Aspect ratios accepted by the Gemini Flash image (generateContent) models. +# Aspect ratios documented for every Gemini Flash image (generateContent) model. +# The extreme ratios (1:4, 4:1, 1:8, 8:1) are only listed for the 3.1 Flash / +# Flash Lite tables, so they are left out to avoid sending an unsupported value +# to 2.5 Flash Image or 3.1 Pro Image. _GEMINI_FLASH_ASPECT_RATIOS = { - "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", - "9:16", "16:9", "21:9", "1:4", "4:1", "1:8", "8:1", + "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9", } -# Image-size tokens accepted by Gemini 3+ image models (2.5 Flash Image ignores it). +# Image-size tokens accepted by Gemini 3+ image models (earlier Flash image +# models expose only a single fixed resolution). _GEMINI_FLASH_IMAGE_SIZES = {"512", "1K", "2K", "4K"} _OLLAMA_DEFAULT_SIDE = 1024 _OLLAMA_SIZE_PRESETS = { @@ -787,8 +790,13 @@ def _gemini_flash_image_config( def _gemini_flash_supports_image_size(model: str) -> bool: - """Return whether the model honors ``imageSize`` (Gemini 3+ image models).""" - return "2.5" not in model.lower() + """Return whether the model honors ``imageSize``. + + Only Gemini 3+ image models expose a configurable image size; earlier Flash + image models (2.0, 2.5) generate at a single fixed resolution, so ``imageSize`` + is identified positively rather than by excluding a single version. + """ + return "gemini-3" in model.lower() async def _aihubmix_images_from_payload( diff --git a/tests/providers/test_image_generation.py b/tests/providers/test_image_generation.py index 8e3110498..19f5eb4d8 100644 --- a/tests/providers/test_image_generation.py +++ b/tests/providers/test_image_generation.py @@ -465,14 +465,32 @@ async def test_gemini_flash_2_5_drops_image_size() -> None: @pytest.mark.asyncio -async def test_gemini_flash_ignores_unsupported_hints() -> None: +async def test_gemini_flash_2_0_drops_image_size() -> None: fake = FakeClient(_gemini_flash_image_response()) client = GeminiImageGenerationClient(api_key="AIza-test", client=fake) # type: ignore[arg-type] + await client.generate( + prompt="draw a cat", + model="gemini-2.0-flash-preview-image-generation", + aspect_ratio="16:9", + image_size="1K", + ) + + image_config = fake.calls[0]["json"]["generationConfig"]["responseFormat"]["image"] + assert image_config == {"aspectRatio": "16:9"} + + +@pytest.mark.asyncio +async def test_gemini_flash_ignores_unsupported_hints() -> None: + fake = FakeClient(_gemini_flash_image_response()) + client = GeminiImageGenerationClient(api_key="AIza-test", client=fake) # type: ignore[arg-type] + + # 7:5 is not a documented ratio; 1:8 is only valid for 3.1 Flash, not Pro; + # 1024x1024 is not a valid Gemini image-size token. All are dropped. await client.generate( prompt="draw a cat", model="gemini-3-pro-image", - aspect_ratio="7:5", + aspect_ratio="1:8", image_size="1024x1024", ) From a8604a31720690ecbb341a7fa5eab34b6f0e695d Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Sun, 26 Jul 2026 17:39:36 +0800 Subject: [PATCH 43/48] fix(image): scope Gemini image sizes by model --- nanobot/providers/image_generation.py | 30 +++++++++++++++--------- tests/providers/test_image_generation.py | 28 ++++++++++++++++++++++ 2 files changed, 47 insertions(+), 11 deletions(-) diff --git a/nanobot/providers/image_generation.py b/nanobot/providers/image_generation.py index f9d76e288..9e9332fd1 100644 --- a/nanobot/providers/image_generation.py +++ b/nanobot/providers/image_generation.py @@ -40,9 +40,11 @@ _GEMINI_IMAGEN_ASPECT_RATIOS = {"1:1", "9:16", "16:9", "3:4", "4:3"} _GEMINI_FLASH_ASPECT_RATIOS = { "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9", } -# Image-size tokens accepted by Gemini 3+ image models (earlier Flash image -# models expose only a single fixed resolution). -_GEMINI_FLASH_IMAGE_SIZES = {"512", "1K", "2K", "4K"} +# Gemini 3 Pro image models accept these sizes. Gemini 3.1 Flash adds 512, +# while Gemini 3.1 Flash Lite supports only 1K. +_GEMINI_3_IMAGE_SIZES = {"1K", "2K", "4K"} +_GEMINI_31_FLASH_IMAGE_SIZES = {"512", *_GEMINI_3_IMAGE_SIZES} +_GEMINI_31_FLASH_LITE_IMAGE_SIZES = {"1K"} _OLLAMA_DEFAULT_SIDE = 1024 _OLLAMA_SIZE_PRESETS = { "1K": 1024, @@ -782,21 +784,27 @@ def _gemini_flash_image_config( config: dict[str, str] = {} if aspect_ratio and aspect_ratio in _GEMINI_FLASH_ASPECT_RATIOS: config["aspectRatio"] = aspect_ratio - if image_size and _gemini_flash_supports_image_size(model): + if image_size: normalized = image_size.strip().upper() - if normalized in _GEMINI_FLASH_IMAGE_SIZES: + if normalized in _gemini_flash_supported_image_sizes(model): config["imageSize"] = normalized return config -def _gemini_flash_supports_image_size(model: str) -> bool: - """Return whether the model honors ``imageSize``. +def _gemini_flash_supported_image_sizes(model: str) -> set[str]: + """Return the ``imageSize`` values documented for a Flash-path model. - Only Gemini 3+ image models expose a configurable image size; earlier Flash - image models (2.0, 2.5) generate at a single fixed resolution, so ``imageSize`` - is identified positively rather than by excluding a single version. + Earlier Flash image models (2.0, 2.5) expose no configurable size. Gemini + 3.1 Flash Lite is intentionally checked before the broader Flash match. """ - return "gemini-3" in model.lower() + normalized = model.lower() + if "gemini-3.1-flash-lite-image" in normalized: + return _GEMINI_31_FLASH_LITE_IMAGE_SIZES + if "gemini-3.1-flash-image" in normalized: + return _GEMINI_31_FLASH_IMAGE_SIZES + if "gemini-3" in normalized: + return _GEMINI_3_IMAGE_SIZES + return set() async def _aihubmix_images_from_payload( diff --git a/tests/providers/test_image_generation.py b/tests/providers/test_image_generation.py index 19f5eb4d8..7fe63e430 100644 --- a/tests/providers/test_image_generation.py +++ b/tests/providers/test_image_generation.py @@ -480,6 +480,34 @@ async def test_gemini_flash_2_0_drops_image_size() -> None: assert image_config == {"aspectRatio": "16:9"} +@pytest.mark.parametrize( + ("model", "image_size", "expected"), + [ + ("gemini-3-pro-image", "512", None), + ("gemini-3.1-flash-lite-image", "2K", None), + ("gemini-3.1-flash-lite-image", "1K", {"imageSize": "1K"}), + ("gemini-3.1-flash-image", "512", {"imageSize": "512"}), + ], +) +@pytest.mark.asyncio +async def test_gemini_flash_scopes_image_size_by_model( + model: str, + image_size: str, + expected: dict[str, str] | None, +) -> None: + fake = FakeClient(_gemini_flash_image_response()) + client = GeminiImageGenerationClient(api_key="AIza-test", client=fake) # type: ignore[arg-type] + + await client.generate( + prompt="draw a cat", + model=model, + image_size=image_size, + ) + + response_format = fake.calls[0]["json"]["generationConfig"].get("responseFormat") + assert response_format == ({"image": expected} if expected else None) + + @pytest.mark.asyncio async def test_gemini_flash_ignores_unsupported_hints() -> None: fake = FakeClient(_gemini_flash_image_response()) From cf1e801a294299fb880a2090df6460534ff03076 Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Mon, 27 Jul 2026 03:00:16 +0800 Subject: [PATCH 44/48] fix(image): align Gemini hints with model capabilities --- nanobot/providers/image_generation.py | 37 ++++++++++++++++++------ tests/providers/test_image_generation.py | 29 +++++++++++++++++++ 2 files changed, 57 insertions(+), 9 deletions(-) diff --git a/nanobot/providers/image_generation.py b/nanobot/providers/image_generation.py index 9e9332fd1..b3843b115 100644 --- a/nanobot/providers/image_generation.py +++ b/nanobot/providers/image_generation.py @@ -33,13 +33,18 @@ _AIHUBMIX_ASPECT_RATIO_SIZES = { } _GEMINI_DEFAULT_TIMEOUT_S = 120.0 _GEMINI_IMAGEN_ASPECT_RATIOS = {"1:1", "9:16", "16:9", "3:4", "4:3"} -# Aspect ratios documented for every Gemini Flash image (generateContent) model. -# The extreme ratios (1:4, 4:1, 1:8, 8:1) are only listed for the 3.1 Flash / -# Flash Lite tables, so they are left out to avoid sending an unsupported value -# to 2.5 Flash Image or 3.1 Pro Image. -_GEMINI_FLASH_ASPECT_RATIOS = { +# Aspect ratios documented for every Gemini image model using generateContent. +_GEMINI_FLASH_COMMON_ASPECT_RATIOS = { "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9", } +# Gemini 3.1 Flash and Flash Lite additionally accept extreme aspect ratios. +_GEMINI_31_FLASH_ASPECT_RATIOS = { + *_GEMINI_FLASH_COMMON_ASPECT_RATIOS, + "1:4", + "4:1", + "1:8", + "8:1", +} # Gemini 3 Pro image models accept these sizes. Gemini 3.1 Flash adds 512, # while Gemini 3.1 Flash Lite supports only 1K. _GEMINI_3_IMAGE_SIZES = {"1K", "2K", "4K"} @@ -778,11 +783,12 @@ def _gemini_flash_image_config( ) -> dict[str, str]: """Build the ``responseFormat.image`` config for Gemini Flash image models. - Aspect ratio applies to all Flash image models; image size is only honored - by Gemini 3+ image models (``gemini-2.5-flash-image`` ignores it). + Capabilities are model-specific: Gemini 3.1 Flash variants support four + additional extreme ratios, while configurable image sizes are limited to + the documented Gemini 3 image model families. """ config: dict[str, str] = {} - if aspect_ratio and aspect_ratio in _GEMINI_FLASH_ASPECT_RATIOS: + if aspect_ratio and aspect_ratio in _gemini_flash_supported_aspect_ratios(model): config["aspectRatio"] = aspect_ratio if image_size: normalized = image_size.strip().upper() @@ -791,6 +797,19 @@ def _gemini_flash_image_config( return config +def _gemini_flash_supported_aspect_ratios(model: str) -> set[str]: + """Return the documented aspect ratios for a generateContent image model.""" + normalized = model.lower() + if ( + "gemini-3.1-flash-lite-image" in normalized + or "gemini-3.1-flash-image" in normalized + ): + return _GEMINI_31_FLASH_ASPECT_RATIOS + if "gemini-" in normalized and "image" in normalized: + return _GEMINI_FLASH_COMMON_ASPECT_RATIOS + return set() + + def _gemini_flash_supported_image_sizes(model: str) -> set[str]: """Return the ``imageSize`` values documented for a Flash-path model. @@ -802,7 +821,7 @@ def _gemini_flash_supported_image_sizes(model: str) -> set[str]: return _GEMINI_31_FLASH_LITE_IMAGE_SIZES if "gemini-3.1-flash-image" in normalized: return _GEMINI_31_FLASH_IMAGE_SIZES - if "gemini-3" in normalized: + if "gemini-3-pro-image" in normalized: return _GEMINI_3_IMAGE_SIZES return set() diff --git a/tests/providers/test_image_generation.py b/tests/providers/test_image_generation.py index 7fe63e430..0ec35c0e4 100644 --- a/tests/providers/test_image_generation.py +++ b/tests/providers/test_image_generation.py @@ -480,10 +480,39 @@ async def test_gemini_flash_2_0_drops_image_size() -> None: assert image_config == {"aspectRatio": "16:9"} +@pytest.mark.parametrize( + ("model", "aspect_ratio", "expected"), + [ + ("gemini-3.1-flash-image", "1:8", {"aspectRatio": "1:8"}), + ("gemini-3.1-flash-lite-image", "4:1", {"aspectRatio": "4:1"}), + ("gemini-3-pro-image", "1:8", None), + ("gemini-2.5-flash-image", "4:1", None), + ], +) +@pytest.mark.asyncio +async def test_gemini_flash_scopes_extreme_aspect_ratios_by_model( + model: str, + aspect_ratio: str, + expected: dict[str, str] | None, +) -> None: + fake = FakeClient(_gemini_flash_image_response()) + client = GeminiImageGenerationClient(api_key="AIza-test", client=fake) # type: ignore[arg-type] + + await client.generate( + prompt="draw a cat", + model=model, + aspect_ratio=aspect_ratio, + ) + + response_format = fake.calls[0]["json"]["generationConfig"].get("responseFormat") + assert response_format == ({"image": expected} if expected else None) + + @pytest.mark.parametrize( ("model", "image_size", "expected"), [ ("gemini-3-pro-image", "512", None), + ("gemini-3-pro", "2K", None), ("gemini-3.1-flash-lite-image", "2K", None), ("gemini-3.1-flash-lite-image", "1K", {"imageSize": "1K"}), ("gemini-3.1-flash-image", "512", {"imageSize": "512"}), From 4408cde019eeba1ebf75ab1f4ebffaa425e278ae Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Sun, 26 Jul 2026 16:50:08 +0800 Subject: [PATCH 45/48] fix(security): harden generated image downloads --- nanobot/providers/image_generation.py | 106 ++++++++++------ tests/providers/test_image_generation.py | 41 +++++-- .../test_image_generation_security.py | 113 ++++++++++++++++++ 3 files changed, 214 insertions(+), 46 deletions(-) create mode 100644 tests/providers/test_image_generation_security.py diff --git a/nanobot/providers/image_generation.py b/nanobot/providers/image_generation.py index b3843b115..7f8246517 100644 --- a/nanobot/providers/image_generation.py +++ b/nanobot/providers/image_generation.py @@ -10,11 +10,13 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from pathlib import Path from typing import Any +from urllib.parse import urljoin import httpx from loguru import logger from nanobot.providers.registry import find_by_name +from nanobot.security.network import PinnedDNSAsyncTransport, UnsafeURLRequestError from nanobot.utils.helpers import detect_image_mime _OPENROUTER_ATTRIBUTION_HEADERS = { @@ -23,6 +25,8 @@ _OPENROUTER_ATTRIBUTION_HEADERS = { "X-OpenRouter-Categories": "cli-agent,personal-agent", } _DEFAULT_TIMEOUT_S = 120.0 +_IMAGE_DOWNLOAD_MAX_BYTES = 32 * 1024 * 1024 +_IMAGE_DOWNLOAD_MAX_REDIRECTS = 5 _AIHUBMIX_TIMEOUT_S = 300.0 _AIHUBMIX_ASPECT_RATIO_SIZES = { "1:1": "1024x1024", @@ -131,16 +135,66 @@ def _aihubmix_model_path(model: str) -> str: async def _download_image_data_url( - client: httpx.AsyncClient, url: str, + *, + transport: httpx.AsyncBaseTransport | None = None, ) -> str: - response = await client.get(url) try: - response.raise_for_status() - except httpx.HTTPStatusError as exc: - detail = response.text[:500] - raise ImageGenerationError(f"failed to download generated image: {detail}") from exc - raw = response.content + safe_transport = PinnedDNSAsyncTransport(inner=transport) + # Proxies resolve the target independently and would defeat DNS pinning. + async with httpx.AsyncClient( + transport=safe_transport, + follow_redirects=False, + timeout=_DEFAULT_TIMEOUT_S, + trust_env=False, + ) as client: + current_url = url + for _ in range(_IMAGE_DOWNLOAD_MAX_REDIRECTS + 1): + async with client.stream("GET", current_url) as response: + if response.is_redirect: + location = response.headers.get("location") + if not location: + raise ImageGenerationError( + "generated image URL redirected without a location" + ) + current_url = urljoin(str(response.url), location) + continue + + try: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + raise ImageGenerationError( + f"failed to download generated image (HTTP {response.status_code})" + ) from exc + + declared_size = response.headers.get("content-length") + if declared_size: + try: + if int(declared_size) > _IMAGE_DOWNLOAD_MAX_BYTES: + raise ImageGenerationError( + "generated image exceeded the 32 MiB download limit" + ) + except ValueError: + pass + + chunks: list[bytes] = [] + total = 0 + async for chunk in response.aiter_bytes(): + total += len(chunk) + if total > _IMAGE_DOWNLOAD_MAX_BYTES: + raise ImageGenerationError( + "generated image exceeded the 32 MiB download limit" + ) + chunks.append(chunk) + raw = b"".join(chunks) + break + else: + raise ImageGenerationError("generated image URL exceeded the redirect limit") + except UnsafeURLRequestError as exc: + raise ImageGenerationError(f"blocked unsafe generated image URL: {exc}") from exc + except httpx.RequestError as exc: + raise ImageGenerationError(f"failed to download generated image: {exc}") from exc + mime = detect_image_mime(raw) if mime is None: raise ImageGenerationError("generated image URL did not return a supported image") @@ -452,7 +506,7 @@ class AIHubMixImageGenerationClient(ImageGenerationProvider): raise ImageGenerationError(f"AIHubMix image generation failed: {detail}") from exc payload = response.json() - images = await _aihubmix_images_from_payload(client, payload) + images = await _aihubmix_images_from_payload(payload) self._require_images(images, payload) @@ -827,7 +881,6 @@ def _gemini_flash_supported_image_sizes(model: str) -> set[str]: async def _aihubmix_images_from_payload( - client: httpx.AsyncClient, payload: dict[str, Any], ) -> list[str]: images: list[str] = [] @@ -846,7 +899,7 @@ async def _aihubmix_images_from_payload( if value.startswith("data:image/"): images.append(value) elif value.startswith(("http://", "https://")): - images.append(await _download_image_data_url(client, value)) + images.append(await _download_image_data_url(value)) return if not isinstance(value, dict): return @@ -1047,15 +1100,7 @@ class OpenAIImageGenerationClient(ImageGenerationProvider): return model async def _parse_images_response(self, payload: dict[str, Any]) -> list[str]: - client = self._client - owns_client = client is None - if owns_client: - client = httpx.AsyncClient(timeout=self.timeout) - try: - return await _openai_images_from_payload(client, payload) - finally: - if owns_client: - await client.aclose() + return await _openai_images_from_payload(payload) async def _post_image_edit( self, @@ -1266,15 +1311,7 @@ class CustomImageGenerationClient(ImageGenerationProvider): logger.info("Custom Images API response ({}): {}", response.status_code, {k: v for k, v in payload.items() if k != "data"}) - client = self._client - owns_client = client is None - if owns_client: - client = httpx.AsyncClient(timeout=self.timeout) - try: - images = await _openai_images_from_payload(client, payload) - finally: - if owns_client: - await client.aclose() + images = await _openai_images_from_payload(payload) self._require_images(images, payload) @@ -1467,7 +1504,6 @@ def _openai_explicit_size_supported( async def _openai_images_from_payload( - client: httpx.AsyncClient, payload: dict[str, Any], ) -> list[str]: """Extract images from OpenAI Images API response. @@ -1484,7 +1520,7 @@ async def _openai_images_from_payload( continue url = item.get("url") if isinstance(url, str) and url: - images.append(await _download_image_data_url(client, url)) + images.append(await _download_image_data_url(url)) return images @@ -1798,7 +1834,7 @@ class ZhipuImageGenerationClient(ImageGenerationProvider): raise ImageGenerationError(f"Zhipu image generation failed: {detail}") from exc payload = response.json() - images = await _zhipu_images_from_payload(client, payload) + images = await _zhipu_images_from_payload(payload) self._require_images(images, payload) @@ -1822,7 +1858,6 @@ def _zhipu_size( async def _zhipu_images_from_payload( - client: httpx.AsyncClient, payload: dict[str, Any], ) -> list[str]: """Extract image data URLs from Zhipu API response. @@ -1836,7 +1871,7 @@ async def _zhipu_images_from_payload( continue url = item.get("url") if isinstance(url, str) and url: - images.append(await _download_image_data_url(client, url)) + images.append(await _download_image_data_url(url)) return images @@ -1999,7 +2034,7 @@ class ModelScopeImageGenerationClient(ImageGenerationProvider): status = data.get("task_status") if status == "SUCCEED": - return await self._collect_images(client, data) + return await self._collect_images(data) if status == "FAILED": raise ImageGenerationError( f"ModelScope image generation task failed: {data}" @@ -2014,7 +2049,6 @@ class ModelScopeImageGenerationClient(ImageGenerationProvider): @staticmethod async def _collect_images( - client: httpx.AsyncClient, data: dict[str, Any], ) -> list[str]: images: list[str] = [] @@ -2023,7 +2057,7 @@ class ModelScopeImageGenerationClient(ImageGenerationProvider): if url.startswith("data:image/"): images.append(url) else: - images.append(await _download_image_data_url(client, url)) + images.append(await _download_image_data_url(url)) return images diff --git a/tests/providers/test_image_generation.py b/tests/providers/test_image_generation.py index 0ec35c0e4..8504d4d69 100644 --- a/tests/providers/test_image_generation.py +++ b/tests/providers/test_image_generation.py @@ -102,6 +102,22 @@ class CodexStreamingCompleteThenErrorResponse(FakeResponse): ) +@pytest.fixture(autouse=True) +def generated_image_downloads(monkeypatch) -> list[str]: + """Keep provider response parsing tests independent from outbound HTTP.""" + urls: list[str] = [] + + async def download(url: str) -> str: + urls.append(url) + return PNG_DATA_URL + + monkeypatch.setattr( + "nanobot.providers.image_generation._download_image_data_url", + download, + ) + return urls + + @pytest.mark.asyncio async def test_openrouter_image_generation_payload_and_response(tmp_path: Path) -> None: ref = tmp_path / "ref.png" @@ -277,7 +293,9 @@ async def test_aihubmix_image_edit_payload_uses_reference_images(tmp_path: Path) @pytest.mark.asyncio -async def test_aihubmix_image_generation_downloads_url_response() -> None: +async def test_aihubmix_image_generation_downloads_url_response( + generated_image_downloads: list[str], +) -> None: fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]})) fake.get_response = FakeResponse({}, content=PNG_BYTES) client = AIHubMixImageGenerationClient( @@ -288,7 +306,7 @@ async def test_aihubmix_image_generation_downloads_url_response() -> None: response = await client.generate(prompt="draw", model="gpt-image-2-free") assert response.images[0].startswith("data:image/png;base64,") - assert fake.get_calls[0]["url"] == "https://cdn.example/image.png" + assert generated_image_downloads == ["https://cdn.example/image.png"] @pytest.mark.asyncio @@ -818,7 +836,7 @@ async def test_openai_b64_json_response_uses_detected_mime() -> None: @pytest.mark.asyncio -async def test_openai_url_download_fallback() -> None: +async def test_openai_url_download_fallback(generated_image_downloads: list[str]) -> None: fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]})) fake.get_response = FakeResponse({}, content=PNG_BYTES) client = OpenAIImageGenerationClient( @@ -829,7 +847,7 @@ async def test_openai_url_download_fallback() -> None: response = await client.generate(prompt="draw", model="dall-e-3") assert response.images[0].startswith("data:image/png;base64,") - assert fake.get_calls[0]["url"] == "https://cdn.example/image.png" + assert generated_image_downloads == ["https://cdn.example/image.png"] @pytest.mark.asyncio @@ -1192,7 +1210,9 @@ async def test_custom_generate_maps_one_k_to_openai_dimension() -> None: @pytest.mark.asyncio -async def test_custom_generate_extra_body_can_override_defaults() -> None: +async def test_custom_generate_extra_body_can_override_defaults( + generated_image_downloads: list[str], +) -> None: fake = FakeClient(FakeResponse({"data": [{"url": "https://images.example/cat.png"}]})) fake.get_response = FakeResponse({}, content=PNG_BYTES) client = CustomImageGenerationClient( @@ -1208,9 +1228,8 @@ async def test_custom_generate_extra_body_can_override_defaults() -> None: image_size="1K", ) - expected_data_url = f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode('ascii')}" - assert response.images == [expected_data_url] - assert fake.get_calls[0]["url"] == "https://images.example/cat.png" + assert response.images == [PNG_DATA_URL] + assert generated_image_downloads == ["https://images.example/cat.png"] body = fake.calls[0]["json"] assert body["response_format"] == "url" assert body["size"] == "2K" @@ -1616,7 +1635,9 @@ async def test_zhipu_image_generation_with_explicit_size() -> None: @pytest.mark.asyncio -async def test_zhipu_image_generation_downloads_url_response() -> None: +async def test_zhipu_image_generation_downloads_url_response( + generated_image_downloads: list[str], +) -> None: fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]})) fake.get_response = FakeResponse({}, content=PNG_BYTES) client = ZhipuImageGenerationClient( @@ -1627,7 +1648,7 @@ async def test_zhipu_image_generation_downloads_url_response() -> None: response = await client.generate(prompt="draw", model="glm-image") assert response.images[0].startswith("data:image/png;base64,") - assert fake.get_calls[0]["url"] == "https://cdn.example/image.png" + assert generated_image_downloads == ["https://cdn.example/image.png"] @pytest.mark.asyncio diff --git a/tests/providers/test_image_generation_security.py b/tests/providers/test_image_generation_security.py new file mode 100644 index 000000000..466b74cd2 --- /dev/null +++ b/tests/providers/test_image_generation_security.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import socket + +import httpx +import pytest + +from nanobot.providers import image_generation +from nanobot.providers.image_generation import ImageGenerationError, _download_image_data_url + +PNG_BYTES = ( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01" + b"\x00\x00\x00\x01\x08\x04\x00\x00\x00\xb5\x1c\x0c\x02" + b"\x00\x00\x00\x0bIDATx\xdacd\xfc\xff\x1f\x00\x03\x03" + b"\x02\x00\xef\xbf\xa7\xdb\x00\x00\x00\x00IEND\xaeB`\x82" +) + + +def _resolve_public(host: str, port: int | None, *args, **kwargs): + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + socket.IPPROTO_TCP, + "", + ("93.184.216.34", port or 0), + ) + ] + + +@pytest.mark.asyncio +async def test_generated_image_download_blocks_private_target() -> None: + requested = False + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal requested + requested = True + return httpx.Response(200, content=PNG_BYTES) + + with pytest.raises(ImageGenerationError, match="blocked unsafe generated image URL"): + await _download_image_data_url( + "http://127.0.0.1/admin", + transport=httpx.MockTransport(handler), + ) + + assert requested is False + + +@pytest.mark.asyncio +async def test_generated_image_download_revalidates_redirects(monkeypatch) -> None: + original_getaddrinfo = socket.getaddrinfo + + def resolve_test_hosts(host: str, port: int | None, *args, **kwargs): + if host == "cdn.example": + return _resolve_public(host, port, *args, **kwargs) + return original_getaddrinfo(host, port, *args, **kwargs) + + monkeypatch.setattr("nanobot.security.network.socket.getaddrinfo", resolve_test_hosts) + requested: list[str] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + requested.append(str(request.url)) + return httpx.Response(302, headers={"location": "http://169.254.169.254/latest"}) + + with pytest.raises(ImageGenerationError, match="blocked unsafe generated image URL"): + await _download_image_data_url( + "https://cdn.example/image.png", + transport=httpx.MockTransport(handler), + ) + + assert requested == ["https://cdn.example/image.png"] + + +@pytest.mark.asyncio +async def test_generated_image_download_returns_valid_data_url(monkeypatch) -> None: + monkeypatch.setattr( + "nanobot.security.network.socket.getaddrinfo", + _resolve_public, + ) + + async def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, content=PNG_BYTES) + + result = await _download_image_data_url( + "https://cdn.example/image.png", + transport=httpx.MockTransport(handler), + ) + + assert result.startswith("data:image/png;base64,") + + +class _OversizedStream(httpx.AsyncByteStream): + async def __aiter__(self): + yield b"12345" + yield b"6789" + + +@pytest.mark.asyncio +async def test_generated_image_download_enforces_streaming_size_limit(monkeypatch) -> None: + monkeypatch.setattr( + "nanobot.security.network.socket.getaddrinfo", + _resolve_public, + ) + monkeypatch.setattr(image_generation, "_IMAGE_DOWNLOAD_MAX_BYTES", 8) + + async def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, stream=_OversizedStream()) + + with pytest.raises(ImageGenerationError, match="download limit"): + await _download_image_data_url( + "https://cdn.example/image.png", + transport=httpx.MockTransport(handler), + ) From cc3dbbe804ba3cc683baa096a923b868b011e75c Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Sun, 26 Jul 2026 22:25:32 +0800 Subject: [PATCH 46/48] fix(security): block IPv6 unspecified SSRF targets --- nanobot/security/network.py | 1 + tests/providers/test_image_generation_security.py | 9 +++++++-- tests/security/test_security_network.py | 1 + 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/nanobot/security/network.py b/nanobot/security/network.py index 95523f3a4..23daf980f 100644 --- a/nanobot/security/network.py +++ b/nanobot/security/network.py @@ -20,6 +20,7 @@ _BLOCKED_NETWORKS = [ ipaddress.ip_network("169.254.0.0/16"), # link-local / cloud metadata ipaddress.ip_network("172.16.0.0/12"), ipaddress.ip_network("192.168.0.0/16"), + ipaddress.ip_network("::/128"), # unspecified; may route to local host ipaddress.ip_network("::1/128"), ipaddress.ip_network("fc00::/7"), # unique local ipaddress.ip_network("fe80::/10"), # link-local v6 diff --git a/tests/providers/test_image_generation_security.py b/tests/providers/test_image_generation_security.py index 466b74cd2..44b586a7f 100644 --- a/tests/providers/test_image_generation_security.py +++ b/tests/providers/test_image_generation_security.py @@ -28,8 +28,13 @@ def _resolve_public(host: str, port: int | None, *args, **kwargs): ] +@pytest.mark.parametrize( + "url", + ["http://127.0.0.1/admin", "http://[::]/admin"], + ids=["ipv4-loopback", "ipv6-unspecified"], +) @pytest.mark.asyncio -async def test_generated_image_download_blocks_private_target() -> None: +async def test_generated_image_download_blocks_unsafe_target(url: str) -> None: requested = False async def handler(request: httpx.Request) -> httpx.Response: @@ -39,7 +44,7 @@ async def test_generated_image_download_blocks_private_target() -> None: with pytest.raises(ImageGenerationError, match="blocked unsafe generated image URL"): await _download_image_data_url( - "http://127.0.0.1/admin", + url, transport=httpx.MockTransport(handler), ) diff --git a/tests/security/test_security_network.py b/tests/security/test_security_network.py index fc4ea767d..9b71022f8 100644 --- a/tests/security/test_security_network.py +++ b/tests/security/test_security_network.py @@ -148,6 +148,7 @@ def test_blocks_sampled_addresses_from_internal_networks(): "169.254.0.0/16", "172.16.0.0/12", "192.168.0.0/16", + "::/128", "::1/128", "fc00::/7", "fe80::/10", From d73794bc688971428c5612bec90e19234ba81f2a Mon Sep 17 00:00:00 2001 From: chengyongru Date: Mon, 27 Jul 2026 00:33:58 +0800 Subject: [PATCH 47/48] fix(image): honor provider proxy for URL downloads --- docs/image-generation.md | 3 + nanobot/providers/image_generation.py | 81 ++++++++----- tests/providers/test_image_generation.py | 58 ++++++--- .../test_image_generation_security.py | 111 +++++++++++++++++- 4 files changed, 211 insertions(+), 42 deletions(-) diff --git a/docs/image-generation.md b/docs/image-generation.md index bcbb025d3..bf727411a 100644 --- a/docs/image-generation.md +++ b/docs/image-generation.md @@ -70,6 +70,9 @@ Provider settings reuse normal provider config fields: | `providers..apiBase` | Optional custom base URL | | `providers..extraHeaders` | Headers merged into provider requests | | `providers..extraBody` | Extra JSON fields merged into provider request bodies | +| `providers..proxy` | Explicit trusted HTTP proxy for provider requests and returned image URL downloads | + +For providers that return image URLs, direct downloads use DNS pinning. When an explicit provider `proxy` is configured, nanobot validates the initial URL and every redirect locally, then relies on that trusted proxy for final DNS resolution and network egress. Process-wide proxy environment variables are not used for these downloads. Both camelCase and snake_case config keys are accepted, but docs use camelCase to match `config.json`. diff --git a/nanobot/providers/image_generation.py b/nanobot/providers/image_generation.py index 7f8246517..d3ad017e3 100644 --- a/nanobot/providers/image_generation.py +++ b/nanobot/providers/image_generation.py @@ -16,7 +16,11 @@ import httpx from loguru import logger from nanobot.providers.registry import find_by_name -from nanobot.security.network import PinnedDNSAsyncTransport, UnsafeURLRequestError +from nanobot.security.network import ( + PinnedDNSAsyncTransport, + UnsafeURLRequestError, + resolve_url_target, +) from nanobot.utils.helpers import detect_image_mime _OPENROUTER_ATTRIBUTION_HEADERS = { @@ -137,19 +141,31 @@ def _aihubmix_model_path(model: str) -> str: async def _download_image_data_url( url: str, *, + proxy: str | None = None, transport: httpx.AsyncBaseTransport | None = None, ) -> str: try: - safe_transport = PinnedDNSAsyncTransport(inner=transport) - # Proxies resolve the target independently and would defeat DNS pinning. - async with httpx.AsyncClient( - transport=safe_transport, - follow_redirects=False, - timeout=_DEFAULT_TIMEOUT_S, - trust_env=False, - ) as client: + client_kwargs: dict[str, Any] = { + "follow_redirects": False, + "timeout": _DEFAULT_TIMEOUT_S, + "trust_env": False, + } + if proxy: + # An explicit provider proxy is a user-selected trusted egress boundary. + # Validate each URL locally, while the proxy owns final DNS resolution. + client_kwargs["proxy"] = proxy + else: + client_kwargs["transport"] = PinnedDNSAsyncTransport(inner=transport) + + async with httpx.AsyncClient(**client_kwargs) as client: current_url = url for _ in range(_IMAGE_DOWNLOAD_MAX_REDIRECTS + 1): + if proxy: + ok, error, _ = resolve_url_target(current_url) + if not ok: + raise ImageGenerationError( + f"blocked unsafe generated image URL: {error}" + ) async with client.stream("GET", current_url) as response: if response.is_redirect: location = response.headers.get("location") @@ -302,6 +318,13 @@ class ImageGenerationProvider(ABC): raise ImageGenerationError(f"{label} returned no images: {provider_error}") raise ImageGenerationError(f"{label} returned no images for this request") + def _http_client_kwargs(self) -> dict[str, Any]: + kwargs: dict[str, Any] = {"timeout": self.timeout} + if self.proxy: + kwargs["proxy"] = self.proxy + kwargs["trust_env"] = False + return kwargs + async def _http_post( self, url: str, @@ -314,11 +337,7 @@ class ImageGenerationProvider(ABC): return await client.post(url, headers=headers, json=body) if self._client is not None: return await self._client.post(url, headers=headers, json=body) - client_kwargs: dict[str, Any] = {"timeout": self.timeout} - if self.proxy: - client_kwargs["proxy"] = self.proxy - client_kwargs["trust_env"] = False - async with httpx.AsyncClient(**client_kwargs) as c: + async with httpx.AsyncClient(**self._http_client_kwargs()) as c: return await c.post(url, headers=headers, json=body) @@ -446,7 +465,7 @@ class AIHubMixImageGenerationClient(ImageGenerationProvider): } size = _aihubmix_size(aspect_ratio, image_size) - client = self._client or httpx.AsyncClient(timeout=self.timeout) + client = self._client or httpx.AsyncClient(**self._http_client_kwargs()) try: return await self._generate_with_client( client, @@ -506,7 +525,7 @@ class AIHubMixImageGenerationClient(ImageGenerationProvider): raise ImageGenerationError(f"AIHubMix image generation failed: {detail}") from exc payload = response.json() - images = await _aihubmix_images_from_payload(payload) + images = await _aihubmix_images_from_payload(payload, proxy=self.proxy) self._require_images(images, payload) @@ -882,6 +901,8 @@ def _gemini_flash_supported_image_sizes(model: str) -> set[str]: async def _aihubmix_images_from_payload( payload: dict[str, Any], + *, + proxy: str | None = None, ) -> list[str]: images: list[str] = [] candidates: list[Any] = [] @@ -899,7 +920,7 @@ async def _aihubmix_images_from_payload( if value.startswith("data:image/"): images.append(value) elif value.startswith(("http://", "https://")): - images.append(await _download_image_data_url(value)) + images.append(await _download_image_data_url(value, proxy=proxy)) return if not isinstance(value, dict): return @@ -1100,7 +1121,7 @@ class OpenAIImageGenerationClient(ImageGenerationProvider): return model async def _parse_images_response(self, payload: dict[str, Any]) -> list[str]: - return await _openai_images_from_payload(payload) + return await _openai_images_from_payload(payload, proxy=self.proxy) async def _post_image_edit( self, @@ -1130,7 +1151,7 @@ class OpenAIImageGenerationClient(ImageGenerationProvider): data=body, files=files, ) - async with httpx.AsyncClient(timeout=self.timeout) as c: + async with httpx.AsyncClient(**self._http_client_kwargs()) as c: return await c.post( f"{self.api_base}/images/edits", headers=headers, @@ -1311,7 +1332,7 @@ class CustomImageGenerationClient(ImageGenerationProvider): logger.info("Custom Images API response ({}): {}", response.status_code, {k: v for k, v in payload.items() if k != "data"}) - images = await _openai_images_from_payload(payload) + images = await _openai_images_from_payload(payload, proxy=self.proxy) self._require_images(images, payload) @@ -1505,6 +1526,8 @@ def _openai_explicit_size_supported( async def _openai_images_from_payload( payload: dict[str, Any], + *, + proxy: str | None = None, ) -> list[str]: """Extract images from OpenAI Images API response. @@ -1520,7 +1543,7 @@ async def _openai_images_from_payload( continue url = item.get("url") if isinstance(url, str) and url: - images.append(await _download_image_data_url(url)) + images.append(await _download_image_data_url(url, proxy=proxy)) return images @@ -1800,7 +1823,7 @@ class ZhipuImageGenerationClient(ImageGenerationProvider): url = f"{self.api_base}/images/generations" - client = self._client or httpx.AsyncClient(timeout=self.timeout) + client = self._client or httpx.AsyncClient(**self._http_client_kwargs()) try: return await self._generate_with_client( client, @@ -1834,7 +1857,7 @@ class ZhipuImageGenerationClient(ImageGenerationProvider): raise ImageGenerationError(f"Zhipu image generation failed: {detail}") from exc payload = response.json() - images = await _zhipu_images_from_payload(payload) + images = await _zhipu_images_from_payload(payload, proxy=self.proxy) self._require_images(images, payload) @@ -1859,6 +1882,8 @@ def _zhipu_size( async def _zhipu_images_from_payload( payload: dict[str, Any], + *, + proxy: str | None = None, ) -> list[str]: """Extract image data URLs from Zhipu API response. @@ -1871,7 +1896,7 @@ async def _zhipu_images_from_payload( continue url = item.get("url") if isinstance(url, str) and url: - images.append(await _download_image_data_url(url)) + images.append(await _download_image_data_url(url, proxy=proxy)) return images @@ -1957,7 +1982,7 @@ class ModelScopeImageGenerationClient(ImageGenerationProvider): body.update(self.extra_body) url = f"{self.api_base}/images/generations" - client = self._client or httpx.AsyncClient(timeout=self.timeout) + client = self._client or httpx.AsyncClient(**self._http_client_kwargs()) try: return await self._generate_with_client( client, @@ -2047,8 +2072,8 @@ class ModelScopeImageGenerationClient(ImageGenerationProvider): f"{_MODELSCOPE_POLL_MAX_ATTEMPTS} polls" ) - @staticmethod async def _collect_images( + self, data: dict[str, Any], ) -> list[str]: images: list[str] = [] @@ -2057,7 +2082,9 @@ class ModelScopeImageGenerationClient(ImageGenerationProvider): if url.startswith("data:image/"): images.append(url) else: - images.append(await _download_image_data_url(url)) + images.append( + await _download_image_data_url(url, proxy=self.proxy) + ) return images diff --git a/tests/providers/test_image_generation.py b/tests/providers/test_image_generation.py index 8504d4d69..8181ab1d5 100644 --- a/tests/providers/test_image_generation.py +++ b/tests/providers/test_image_generation.py @@ -103,19 +103,19 @@ class CodexStreamingCompleteThenErrorResponse(FakeResponse): @pytest.fixture(autouse=True) -def generated_image_downloads(monkeypatch) -> list[str]: +def generated_image_downloads(monkeypatch) -> list[tuple[str, str | None]]: """Keep provider response parsing tests independent from outbound HTTP.""" - urls: list[str] = [] + downloads: list[tuple[str, str | None]] = [] - async def download(url: str) -> str: - urls.append(url) + async def download(url: str, *, proxy: str | None = None) -> str: + downloads.append((url, proxy)) return PNG_DATA_URL monkeypatch.setattr( "nanobot.providers.image_generation._download_image_data_url", download, ) - return urls + return downloads @pytest.mark.asyncio @@ -294,19 +294,21 @@ async def test_aihubmix_image_edit_payload_uses_reference_images(tmp_path: Path) @pytest.mark.asyncio async def test_aihubmix_image_generation_downloads_url_response( - generated_image_downloads: list[str], + generated_image_downloads: list[tuple[str, str | None]], ) -> None: fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]})) fake.get_response = FakeResponse({}, content=PNG_BYTES) + proxy = "http://127.0.0.1:23458" client = AIHubMixImageGenerationClient( api_key="sk-ahm-test", + proxy=proxy, client=fake, # type: ignore[arg-type] ) response = await client.generate(prompt="draw", model="gpt-image-2-free") assert response.images[0].startswith("data:image/png;base64,") - assert generated_image_downloads == ["https://cdn.example/image.png"] + assert generated_image_downloads == [("https://cdn.example/image.png", proxy)] @pytest.mark.asyncio @@ -836,18 +838,22 @@ async def test_openai_b64_json_response_uses_detected_mime() -> None: @pytest.mark.asyncio -async def test_openai_url_download_fallback(generated_image_downloads: list[str]) -> None: +async def test_openai_url_download_fallback( + generated_image_downloads: list[tuple[str, str | None]], +) -> None: fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]})) fake.get_response = FakeResponse({}, content=PNG_BYTES) + proxy = "http://127.0.0.1:23458" client = OpenAIImageGenerationClient( api_key="sk-openai-test", + proxy=proxy, client=fake, # type: ignore[arg-type] ) response = await client.generate(prompt="draw", model="dall-e-3") assert response.images[0].startswith("data:image/png;base64,") - assert generated_image_downloads == ["https://cdn.example/image.png"] + assert generated_image_downloads == [("https://cdn.example/image.png", proxy)] @pytest.mark.asyncio @@ -1211,14 +1217,16 @@ async def test_custom_generate_maps_one_k_to_openai_dimension() -> None: @pytest.mark.asyncio async def test_custom_generate_extra_body_can_override_defaults( - generated_image_downloads: list[str], + generated_image_downloads: list[tuple[str, str | None]], ) -> None: fake = FakeClient(FakeResponse({"data": [{"url": "https://images.example/cat.png"}]})) fake.get_response = FakeResponse({}, content=PNG_BYTES) + proxy = "http://127.0.0.1:23458" client = CustomImageGenerationClient( api_key="sk-custom-test", api_base="https://custom.example/v1", extra_body={"response_format": "url", "size": "2K"}, + proxy=proxy, client=fake, # type: ignore[arg-type] ) @@ -1229,7 +1237,7 @@ async def test_custom_generate_extra_body_can_override_defaults( ) assert response.images == [PNG_DATA_URL] - assert generated_image_downloads == ["https://images.example/cat.png"] + assert generated_image_downloads == [("https://images.example/cat.png", proxy)] body = fake.calls[0]["json"] assert body["response_format"] == "url" assert body["size"] == "2K" @@ -1636,19 +1644,21 @@ async def test_zhipu_image_generation_with_explicit_size() -> None: @pytest.mark.asyncio async def test_zhipu_image_generation_downloads_url_response( - generated_image_downloads: list[str], + generated_image_downloads: list[tuple[str, str | None]], ) -> None: fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]})) fake.get_response = FakeResponse({}, content=PNG_BYTES) + proxy = "http://127.0.0.1:23458" client = ZhipuImageGenerationClient( api_key="sk-zhipu-test", + proxy=proxy, client=fake, # type: ignore[arg-type] ) response = await client.generate(prompt="draw", model="glm-image") assert response.images[0].startswith("data:image/png;base64,") - assert generated_image_downloads == ["https://cdn.example/image.png"] + assert generated_image_downloads == [("https://cdn.example/image.png", proxy)] @pytest.mark.asyncio @@ -1728,7 +1738,9 @@ def _modelscope_fast_poll(monkeypatch) -> None: @pytest.mark.asyncio -async def test_modelscope_image_generation_submit_and_poll() -> None: +async def test_modelscope_image_generation_submit_and_poll( + generated_image_downloads: list[tuple[str, str | None]], +) -> None: submit = FakeResponse({"task_id": "abc123"}) poll_responses = [ FakeResponse({"task_status": "PENDING"}), @@ -1738,9 +1750,11 @@ async def test_modelscope_image_generation_submit_and_poll() -> None: }), ] fake = ModelScopeFakeClient(submit, poll_responses) + proxy = "http://127.0.0.1:23458" client = ModelScopeImageGenerationClient( api_key="ms-token", api_base="https://api-inference.modelscope.cn/v1", + proxy=proxy, client=fake, # type: ignore[arg-type] ) @@ -1750,6 +1764,7 @@ async def test_modelscope_image_generation_submit_and_poll() -> None: ) assert response.images[0].startswith("data:image/png;base64,") + assert generated_image_downloads == [("https://cdn.example/image.png", proxy)] # Verify POST request post_call = fake.calls[0] @@ -1919,3 +1934,18 @@ async def test_modelscope_image_generation_poll_timeout(monkeypatch) -> None: # Should have polled up to the (patched) attempt limit. assert len(fake.get_calls) == 3 + + + +def test_image_provider_http_client_kwargs_include_explicit_proxy() -> None: + proxy = "http://127.0.0.1:23458" + client = AIHubMixImageGenerationClient( + api_key="sk-ahm-test", + proxy=proxy, + ) + + assert client._http_client_kwargs() == { + "timeout": client.timeout, + "proxy": proxy, + "trust_env": False, + } diff --git a/tests/providers/test_image_generation_security.py b/tests/providers/test_image_generation_security.py index 44b586a7f..48bc2e7fa 100644 --- a/tests/providers/test_image_generation_security.py +++ b/tests/providers/test_image_generation_security.py @@ -33,8 +33,16 @@ def _resolve_public(host: str, port: int | None, *args, **kwargs): ["http://127.0.0.1/admin", "http://[::]/admin"], ids=["ipv4-loopback", "ipv6-unspecified"], ) +@pytest.mark.parametrize( + "proxy", + [None, "http://127.0.0.1:23458"], + ids=["direct", "explicit-proxy"], +) @pytest.mark.asyncio -async def test_generated_image_download_blocks_unsafe_target(url: str) -> None: +async def test_generated_image_download_blocks_unsafe_target( + url: str, + proxy: str | None, +) -> None: requested = False async def handler(request: httpx.Request) -> httpx.Response: @@ -45,6 +53,7 @@ async def test_generated_image_download_blocks_unsafe_target(url: str) -> None: with pytest.raises(ImageGenerationError, match="blocked unsafe generated image URL"): await _download_image_data_url( url, + proxy=proxy, transport=httpx.MockTransport(handler), ) @@ -116,3 +125,103 @@ async def test_generated_image_download_enforces_streaming_size_limit(monkeypatc "https://cdn.example/image.png", transport=httpx.MockTransport(handler), ) + + +class _StreamContext: + def __init__(self, response: httpx.Response) -> None: + self.response = response + + async def __aenter__(self) -> httpx.Response: + return self.response + + async def __aexit__(self, exc_type, exc, traceback) -> None: + await self.response.aclose() + + +@pytest.mark.asyncio +async def test_generated_image_download_uses_explicit_provider_proxy( + monkeypatch, +) -> None: + monkeypatch.setattr( + "nanobot.security.network.socket.getaddrinfo", + _resolve_public, + ) + captured: dict[str, object] = {} + + class FakeAsyncClient: + def __init__(self, **kwargs) -> None: + captured["kwargs"] = kwargs + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, traceback) -> None: + return None + + def stream(self, method: str, url: str) -> _StreamContext: + captured["request"] = (method, url) + request = httpx.Request(method, url) + return _StreamContext(httpx.Response(200, content=PNG_BYTES, request=request)) + + monkeypatch.setattr(image_generation.httpx, "AsyncClient", FakeAsyncClient) + proxy = "http://127.0.0.1:23458" + + result = await _download_image_data_url( + "https://cdn.example/image.png", + proxy=proxy, + ) + + assert result.startswith("data:image/png;base64,") + assert captured["request"] == ("GET", "https://cdn.example/image.png") + assert captured["kwargs"] == { + "follow_redirects": False, + "timeout": image_generation._DEFAULT_TIMEOUT_S, + "trust_env": False, + "proxy": proxy, + } + + +@pytest.mark.asyncio +async def test_proxied_generated_image_download_revalidates_redirects( + monkeypatch, +) -> None: + original_getaddrinfo = socket.getaddrinfo + + def resolve_test_hosts(host: str, port: int | None, *args, **kwargs): + if host == "cdn.example": + return _resolve_public(host, port, *args, **kwargs) + return original_getaddrinfo(host, port, *args, **kwargs) + + monkeypatch.setattr("nanobot.security.network.socket.getaddrinfo", resolve_test_hosts) + requested: list[str] = [] + + class FakeAsyncClient: + def __init__(self, **kwargs) -> None: + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, traceback) -> None: + return None + + def stream(self, method: str, url: str) -> _StreamContext: + requested.append(url) + request = httpx.Request(method, url) + return _StreamContext( + httpx.Response( + 302, + headers={"location": "http://169.254.169.254/latest"}, + request=request, + ) + ) + + monkeypatch.setattr(image_generation.httpx, "AsyncClient", FakeAsyncClient) + + with pytest.raises(ImageGenerationError, match="blocked unsafe generated image URL"): + await _download_image_data_url( + "https://cdn.example/image.png", + proxy="http://127.0.0.1:23458", + ) + + assert requested == ["https://cdn.example/image.png"] From b3d3a3e6c35496551e47a6681d12a0afc385f728 Mon Sep 17 00:00:00 2001 From: chengyongru <2755839590@qq.com> Date: Mon, 27 Jul 2026 01:11:31 +0800 Subject: [PATCH 48/48] fix(image): delegate DNS to explicit proxy --- .agent/security.md | 4 +- docs/image-generation.md | 2 +- nanobot/config/schema.py | 2 +- nanobot/providers/image_generation.py | 5 ++- nanobot/security/network.py | 30 ++++++++++++-- .../test_image_generation_security.py | 14 +++---- tests/security/test_security_network.py | 41 +++++++++++++++++++ 7 files changed, 83 insertions(+), 15 deletions(-) diff --git a/.agent/security.md b/.agent/security.md index ca9612669..91f2f41b1 100644 --- a/.agent/security.md +++ b/.agent/security.md @@ -14,9 +14,9 @@ Shell execution (`ExecTool`, `agent/tools/shell.py`) also respects `restrict_to_ ## SSRF Protection -All outbound HTTP requests from agent tools must pass through `validate_url_target` (`security/network.py`). By default it blocks loopback, RFC1918 private addresses, CGNAT ranges, link-local ranges, and cloud metadata endpoints (including `169.254.169.254`). +All outbound HTTP requests from agent tools must pass through the shared URL guards in `security/network.py` (`validate_url_target` or `resolve_url_target`). By default they block loopback, RFC1918 private addresses, CGNAT ranges, link-local ranges, and cloud metadata endpoints (including `169.254.169.254`). -The only escape hatch is `configure_ssrf_whitelist(cidrs)`, which reads from `config.tools.ssrf_whitelist` at load time. +For direct requests, the only escape hatch is `configure_ssrf_whitelist(cidrs)`, which reads from `config.tools.ssrf_whitelist` at load time. An explicitly configured `providers..proxy` is a separate user-authorized trust boundary for provider requests and provider-returned image URL downloads. Those downloads still reject malformed URLs and locally identifiable private/internal targets on every redirect, but hostnames unavailable to local DNS are delegated to the trusted proxy. The user-selected proxy owns final DNS resolution and network egress policy. HTTP/SSE MCP transports are part of this boundary: validate configured MCP URLs before probing or constructing clients, and validate each outgoing HTTP request before redirects are followed. Local/private HTTP MCP endpoints are allowed only through the explicit SSRF whitelist. Stdio MCP servers are not part of the HTTP SSRF path. diff --git a/docs/image-generation.md b/docs/image-generation.md index bf727411a..763bd775e 100644 --- a/docs/image-generation.md +++ b/docs/image-generation.md @@ -72,7 +72,7 @@ Provider settings reuse normal provider config fields: | `providers..extraBody` | Extra JSON fields merged into provider request bodies | | `providers..proxy` | Explicit trusted HTTP proxy for provider requests and returned image URL downloads | -For providers that return image URLs, direct downloads use DNS pinning. When an explicit provider `proxy` is configured, nanobot validates the initial URL and every redirect locally, then relies on that trusted proxy for final DNS resolution and network egress. Process-wide proxy environment variables are not used for these downloads. +For providers that return image URLs, direct downloads use DNS pinning. When an explicit provider `proxy` is configured, nanobot rejects malformed URLs and locally identifiable private/internal targets on the initial URL and every redirect. Hostnames unavailable to local DNS are delegated to that trusted proxy, which owns final DNS resolution and network egress. Process-wide proxy environment variables are not used for these downloads. Both camelCase and snake_case config keys are accepted, but docs use camelCase to match `config.json`. diff --git a/nanobot/config/schema.py b/nanobot/config/schema.py index a36ded81a..8b51404e0 100644 --- a/nanobot/config/schema.py +++ b/nanobot/config/schema.py @@ -199,7 +199,7 @@ class ProviderConfig(Base): extra_headers: dict[str, str] | None = None # Custom headers (e.g. APP-Code for AiHubMix) extra_body: dict[str, Any] | None = None # Extra provider request fields; shape depends on provider/API surface extra_query: dict[str, str] | None = None # Extra query params (e.g. api-version for Azure-style gateways) - proxy: str | None = None # OpenAI-compatible/Codex HTTP proxy URL + proxy: str | None = None # Explicit HTTP proxy; image downloads trust its DNS and egress thinking_style: str | None = None # Thinking/reasoning style for custom providers # Valid values mirror the keys of _THINKING_STYLE_MAP in diff --git a/nanobot/providers/image_generation.py b/nanobot/providers/image_generation.py index d3ad017e3..06170bac0 100644 --- a/nanobot/providers/image_generation.py +++ b/nanobot/providers/image_generation.py @@ -161,7 +161,10 @@ async def _download_image_data_url( current_url = url for _ in range(_IMAGE_DOWNLOAD_MAX_REDIRECTS + 1): if proxy: - ok, error, _ = resolve_url_target(current_url) + ok, error, _ = resolve_url_target( + current_url, + trust_remote_dns=True, + ) if not ok: raise ImageGenerationError( f"blocked unsafe generated image URL: {error}" diff --git a/nanobot/security/network.py b/nanobot/security/network.py index 23daf980f..dba5e14ad 100644 --- a/nanobot/security/network.py +++ b/nanobot/security/network.py @@ -74,7 +74,12 @@ def _is_private(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool: return any(normalized in net for net in _BLOCKED_NETWORKS) -def resolve_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool, str, tuple[str, ...]]: +def resolve_url_target( + url: str, + *, + allow_loopback: bool = False, + trust_remote_dns: bool = False, +) -> tuple[bool, str, tuple[str, ...]]: """Validate a URL is safe to fetch: scheme, hostname, and resolved IPs. ``allow_loopback`` is intentionally narrow: it only permits literal @@ -82,8 +87,14 @@ def resolve_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool, loopback. It does not allow RFC1918, link-local, metadata, or public DNS names that happen to resolve to loopback. + ``trust_remote_dns`` accepts ordinary hostnames unavailable to local DNS. + This is only safe when a user-configured trusted proxy owns final DNS + resolution and network egress. Localhost names and private/internal IP + literals remain blocked. + Returns (ok, error_message, resolved_ips). When ok is True, - resolved_ips contains the public IPs that were validated for this URL. + resolved_ips contains the public IPs that were validated for this URL, or + is empty when an unresolved hostname is delegated to a trusted proxy. """ try: p = urlparse(url) @@ -102,7 +113,20 @@ def resolve_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool, try: infos = socket.getaddrinfo(hostname, None, socket.AF_UNSPEC, socket.SOCK_STREAM) except socket.gaierror: - return False, f"Cannot resolve hostname: {hostname}", () + if not trust_remote_dns: + return False, f"Cannot resolve hostname: {hostname}", () + + normalized_hostname = hostname.rstrip(".").lower() + if normalized_hostname == "localhost" or normalized_hostname.endswith(".localhost"): + return False, f"Blocked local/internal hostname: {hostname}", () + + try: + literal_addr = ipaddress.ip_address(normalized_hostname) + except ValueError: + return True, "", () + if _is_private(literal_addr): + return False, f"Blocked private/internal address: {literal_addr}", () + return True, "", (str(_normalize_addr(literal_addr)),) addrs: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = [] for info in infos: diff --git a/tests/providers/test_image_generation_security.py b/tests/providers/test_image_generation_security.py index 48bc2e7fa..5f0ca6088 100644 --- a/tests/providers/test_image_generation_security.py +++ b/tests/providers/test_image_generation_security.py @@ -139,13 +139,13 @@ class _StreamContext: @pytest.mark.asyncio -async def test_generated_image_download_uses_explicit_provider_proxy( +async def test_generated_image_download_delegates_unresolved_host_to_provider_proxy( monkeypatch, ) -> None: - monkeypatch.setattr( - "nanobot.security.network.socket.getaddrinfo", - _resolve_public, - ) + def fail_local_dns(host: str, port: int | None, *args, **kwargs): + raise socket.gaierror(f"cannot resolve {host}") + + monkeypatch.setattr("nanobot.security.network.socket.getaddrinfo", fail_local_dns) captured: dict[str, object] = {} class FakeAsyncClient: @@ -167,12 +167,12 @@ async def test_generated_image_download_uses_explicit_provider_proxy( proxy = "http://127.0.0.1:23458" result = await _download_image_data_url( - "https://cdn.example/image.png", + "https://proxy-only.example/image.png", proxy=proxy, ) assert result.startswith("data:image/png;base64,") - assert captured["request"] == ("GET", "https://cdn.example/image.png") + assert captured["request"] == ("GET", "https://proxy-only.example/image.png") assert captured["kwargs"] == { "follow_redirects": False, "timeout": image_generation._DEFAULT_TIMEOUT_S, diff --git a/tests/security/test_security_network.py b/tests/security/test_security_network.py index 9b71022f8..25b64631b 100644 --- a/tests/security/test_security_network.py +++ b/tests/security/test_security_network.py @@ -195,6 +195,47 @@ def test_resolve_url_target_returns_validated_public_ips(): assert resolved_ips == ("93.184.216.34",) +@pytest.mark.parametrize( + ("trust_remote_dns", "expected_ok"), + [(False, False), (True, True)], +) +def test_resolve_url_target_only_delegates_dns_to_trusted_proxy( + trust_remote_dns: bool, + expected_ok: bool, +): + with patch( + "nanobot.security.network.socket.getaddrinfo", + side_effect=socket.gaierror("local DNS unavailable"), + ): + ok, err, resolved_ips = resolve_url_target( + "https://proxy-only.example/image.png", + trust_remote_dns=trust_remote_dns, + ) + + assert ok is expected_ok, err + assert resolved_ips == () + + +@pytest.mark.parametrize( + "url", + [ + "http://localhost/secret", + "http://service.localhost/secret", + "http://127.0.0.1/secret", + "http://169.254.169.254/latest", + "http://[::1]/secret", + ], +) +def test_resolve_url_target_does_not_delegate_local_targets(url: str): + with patch( + "nanobot.security.network.socket.getaddrinfo", + side_effect=socket.gaierror("local DNS unavailable"), + ): + ok, _, _ = resolve_url_target(url, trust_remote_dns=True) + + assert not ok + + def test_pin_resolved_url_dns_prevents_second_resolution_rebind(): def _rebinding_resolver(hostname, port, family=0, type_=0): return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("169.254.169.254", 0))]