diff --git a/nanobot/webui/transcript.py b/nanobot/webui/transcript.py index bd10f70e0..2d9b6da2f 100644 --- a/nanobot/webui/transcript.py +++ b/nanobot/webui/transcript.py @@ -246,7 +246,7 @@ class WebUITranscriptRecorder: try: dup = json.loads(json.dumps(event, ensure_ascii=False)) append_transcript_object(f"websocket:{chat_id}", dup) - except (ValueError, TypeError) as e: + except (OSError, ValueError, TypeError) as e: self._log.warning("webui transcript append failed: {}", e) def _next_turn_seq(self, chat_id: str, turn_id: str) -> int: @@ -1278,84 +1278,6 @@ def replay_transcript_to_ui_messages( return messages -def _session_content_text(content: Any) -> str: - if isinstance(content, str): - return content - if isinstance(content, list): - parts: list[str] = [] - for item in content: - if isinstance(item, Mapping): - if item.get("type") == "text" and isinstance(item.get("text"), str): - parts.append(str(item["text"])) - elif isinstance(item.get("content"), str): - parts.append(str(item["content"])) - return "\n".join(parts) - return "" - - -def _session_user_to_transcript_event(message: Mapping[str, Any]) -> dict[str, Any]: - event: dict[str, Any] = { - "event": "user", - "text": _session_content_text(message.get("content")), - } - media = message.get("media") - if isinstance(media, list) and media: - event["media_paths"] = [str(path) for path in media if path] - cli_apps = message.get("cli_apps") - if isinstance(cli_apps, list) and cli_apps: - event["cli_apps"] = [dict(app) for app in cli_apps if isinstance(app, dict)] - mcp_presets = message.get("mcp_presets") - if isinstance(mcp_presets, list) and mcp_presets: - event["mcp_presets"] = [dict(preset) for preset in mcp_presets if isinstance(preset, dict)] - return event - - -def _restore_session_user_events( - lines: list[dict[str, Any]], - session_messages: list[dict[str, Any]] | None, -) -> list[dict[str, Any]]: - """Interleave missing user events for legacy transcripts that only persisted replies.""" - if not session_messages: - return lines - session_user_count = sum(1 for m in session_messages if m.get("role") == "user") - transcript_user_count = sum(1 for line in lines if line.get("event") == "user") - if session_user_count == 0 or transcript_user_count >= session_user_count: - return lines - - non_user_lines = [line for line in lines if line.get("event") != "user"] - line_index = 0 - - def pop_assistant_turn() -> list[dict[str, Any]]: - nonlocal line_index - turn: list[dict[str, Any]] = [] - while line_index < len(non_user_lines): - current = non_user_lines[line_index] - line_index += 1 - turn.append(current) - ev = current.get("event") - if ev == "turn_end": - break - if ev in {"message", "stream_end"} and ( - line_index >= len(non_user_lines) - or non_user_lines[line_index].get("event") != "turn_end" - ): - break - return turn - - restored: list[dict[str, Any]] = [] - for session_message in session_messages: - role = session_message.get("role") - if role == "user": - restored.append(_session_user_to_transcript_event(session_message)) - continue - if role == "assistant": - restored.extend(pop_assistant_turn()) - - if line_index < len(non_user_lines): - restored.extend(non_user_lines[line_index:]) - return restored - - def build_webui_thread_response( session_key: str, *, diff --git a/tests/channels/test_websocket_channel.py b/tests/channels/test_websocket_channel.py index 6afeec154..85530da65 100644 --- a/tests/channels/test_websocket_channel.py +++ b/tests/channels/test_websocket_channel.py @@ -27,7 +27,7 @@ from nanobot.channels.websocket import ( from nanobot.config.loader import load_config, save_config from nanobot.config.schema import Config, ModelPresetConfig from nanobot.session import webui_turns as wth -from nanobot.session.manager import Session, SessionManager +from nanobot.session.manager import SessionManager from nanobot.webui.gateway_services import GatewayServices, build_gateway_services from nanobot.webui.http_utils import ( issue_route_secret_matches as _issue_route_secret_matches, @@ -378,7 +378,7 @@ async def test_webui_user_transcript_append_failure_does_not_block_inbound( def fail_append(_session_key: str, _obj: dict[str, Any]) -> None: raise OSError("disk full") - monkeypatch.setattr("nanobot.channels.websocket.append_transcript_object", fail_append) + monkeypatch.setattr("nanobot.webui.transcript.append_transcript_object", fail_append) channel = _ch(bus) conn = AsyncMock() conn.remote_address = ("127.0.0.1", 50123) diff --git a/tests/cli/test_restart_command.py b/tests/cli/test_restart_command.py index e35d52d89..d0ff8ca91 100644 --- a/tests/cli/test_restart_command.py +++ b/tests/cli/test_restart_command.py @@ -214,8 +214,16 @@ class TestRestartCommand: assert "Tasks: 3 active" in response.content @pytest.mark.asyncio - async def test_run_agent_loop_resets_usage_when_provider_omits_it(self): + async def test_run_agent_loop_estimates_usage_when_provider_omits_it(self, monkeypatch): loop, _bus = _make_loop() + monkeypatch.setattr( + "nanobot.agent.runner.estimate_prompt_tokens_chain", + lambda *_args, **_kwargs: (123, "test"), + ) + monkeypatch.setattr( + "nanobot.agent.runner.estimate_message_tokens", + lambda _message: 7, + ) loop.provider.chat_with_retry = AsyncMock(side_effect=[ LLMResponse(content="first", usage={"prompt_tokens": 9, "completion_tokens": 4}), LLMResponse(content="second", usage={}), @@ -226,8 +234,9 @@ class TestRestartCommand: assert loop._last_usage["completion_tokens"] == 4 await loop._run_agent_loop([]) - assert loop._last_usage["prompt_tokens"] == 0 - assert loop._last_usage["completion_tokens"] == 0 + assert loop._last_usage["prompt_tokens"] == 123 + assert loop._last_usage["completion_tokens"] == 7 + assert loop._last_usage["estimated_tokens"] == 130 @pytest.mark.asyncio async def test_status_falls_back_to_last_usage_when_context_estimate_missing(self): diff --git a/tests/utils/test_webui_transcript.py b/tests/utils/test_webui_transcript.py index 0cfe82bdb..5b0e35b17 100644 --- a/tests/utils/test_webui_transcript.py +++ b/tests/utils/test_webui_transcript.py @@ -151,18 +151,20 @@ def test_build_response_restores_session_users_for_legacy_transcript( key, {"event": "message", "chat_id": "legacy-users", "text": "assistant one"}, ) + append_transcript_object(key, {"event": "turn_end", "chat_id": "legacy-users"}) append_transcript_object( key, {"event": "message", "chat_id": "legacy-users", "text": "assistant two"}, ) + append_transcript_object(key, {"event": "turn_end", "chat_id": "legacy-users"}) out = build_webui_thread_response( key, session_messages=[ {"role": "user", "content": "prompt one", "timestamp": "2026-06-02T10:00:00"}, - {"role": "assistant", "content": "session one"}, + {"role": "assistant", "content": "assistant one"}, {"role": "user", "content": "prompt two", "timestamp": "2026-06-02T10:01:00"}, - {"role": "assistant", "content": "session two"}, + {"role": "assistant", "content": "assistant two"}, ], ) @@ -185,19 +187,21 @@ def test_build_response_restores_session_users_without_duplicating_new_transcrip key, {"event": "message", "chat_id": "mixed-users", "text": "old assistant"}, ) + append_transcript_object(key, {"event": "turn_end", "chat_id": "mixed-users"}) append_transcript_object(key, {"event": "user", "chat_id": "mixed-users", "text": "new prompt"}) append_transcript_object( key, {"event": "message", "chat_id": "mixed-users", "text": "new assistant"}, ) + append_transcript_object(key, {"event": "turn_end", "chat_id": "mixed-users"}) out = build_webui_thread_response( key, session_messages=[ {"role": "user", "content": "old prompt"}, - {"role": "assistant", "content": "old session assistant"}, + {"role": "assistant", "content": "old assistant"}, {"role": "user", "content": "new prompt"}, - {"role": "assistant", "content": "new session assistant"}, + {"role": "assistant", "content": "new assistant"}, ], )