"""Unit tests for textual_adapter functions.""" import asyncio import sys from asyncio import Future from collections.abc import AsyncIterator, Generator from datetime import datetime from importlib.metadata import PackageNotFoundError from io import StringIO from pathlib import Path from types import SimpleNamespace from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest from langchain_core.messages import AIMessage, HumanMessage, ToolMessage from langgraph.types import Command from pydantic import ValidationError from rich.console import Console from deepagents_code import config as config_module from deepagents_code._ask_user_types import AskUserWidgetResult, Question from deepagents_code.config import build_stream_config from deepagents_code.textual_adapter import ( ModelStats, SessionStats, TextualUIAdapter, _build_interrupted_ai_message, _handle_interrupt_cleanup, _is_summarization_chunk, execute_task_textual, format_token_count, print_usage_table, ) from deepagents_code.widgets.messages import ( AppMessage, SummarizationMessage, ToolCallMessage, ) async def _mock_mount(widget: object) -> None: """Mock mount function for tests.""" def _mock_approval() -> Future[object]: """Mock approval function for tests.""" future: Future[object] = Future() return future def _noop_status(_: str) -> None: """No-op status callback for tests.""" class TestTextualUIAdapterInit: """Tests for `TextualUIAdapter` initialization.""" def test_set_spinner_callback_stored(self) -> None: """Verify `set_spinner` callback is properly stored.""" async def mock_spinner(status: str | None) -> None: pass adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, set_spinner=mock_spinner, ) assert adapter._set_spinner is mock_spinner def test_set_spinner_defaults_to_none(self) -> None: """Verify `set_spinner` is optional and defaults to `None`.""" adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, ) assert adapter._set_spinner is None def test_current_tool_messages_initialized_empty(self) -> None: """Verify `_current_tool_messages` is initialized as empty dict.""" adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, ) assert adapter._current_tool_messages == {} def test_token_callbacks_initialized_none(self) -> None: """Verify token callbacks are initialized as `None`.""" adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, ) assert adapter._on_tokens_update is None assert adapter._on_tokens_pending is None assert adapter._on_tokens_show is None def test_on_tool_complete_defaults_to_none_and_accepts_callback(self) -> None: """Verify `on_tool_complete` is optional and can be assigned via init.""" adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, ) assert adapter._on_tool_complete is None callback = MagicMock() adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, on_tool_complete=callback, ) assert adapter._on_tool_complete is callback def test_set_token_callbacks(self) -> None: """Verify token callbacks can be assigned.""" adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, ) def update_cb(count: int, *, approximate: bool = False) -> None: pass def pending_cb() -> None: pass def show_cb(*, approximate: bool = False) -> None: pass adapter._on_tokens_update = update_cb adapter._on_tokens_pending = pending_cb adapter._on_tokens_show = show_cb assert adapter._on_tokens_update is update_cb assert adapter._on_tokens_pending is pending_cb assert adapter._on_tokens_show is show_cb def test_finalize_pending_tools_with_error_marks_and_clears(self) -> None: """Pending tool widgets should be marked error and then cleared.""" set_active = MagicMock() adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, set_active_message=set_active, ) tool_1 = MagicMock() tool_2 = MagicMock() adapter._current_tool_messages = {"a": tool_1, "b": tool_2} adapter.finalize_pending_tools_with_error("Agent error: boom") tool_1.set_error.assert_called_once_with("Agent error: boom") tool_2.set_error.assert_called_once_with("Agent error: boom") assert adapter._current_tool_messages == {} set_active.assert_called_once_with(None) class TestInterruptCleanup: """Tests for interrupt cleanup token handling.""" async def test_tool_only_interrupt_marks_tokens_approximate(self) -> None: """Tool-only interrupted turns should keep the stale-token marker.""" mounted: list[object] = [] async def mount_message(widget: object) -> None: mounted.append(widget) await asyncio.sleep(0) set_spinner = AsyncMock() set_active = MagicMock() adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=_mock_approval, set_spinner=set_spinner, set_active_message=set_active, ) tool_widget = MagicMock() tool_widget._tool_name = "read_file" tool_widget._args = {"path": "notes.txt"} adapter._current_tool_messages = {"call-1": tool_widget} show_calls: list[bool] = [] def show_cb(*, approximate: bool = False) -> None: show_calls.append(approximate) adapter._on_tokens_show = show_cb agent = SimpleNamespace(aupdate_state=AsyncMock()) turn_stats = SessionStats() config = {"configurable": {"thread_id": "t-1"}} with patch( "deepagents_code.textual_adapter.time.monotonic", return_value=101.0 ): await _handle_interrupt_cleanup( adapter=adapter, agent=agent, config=config, # ty: ignore pending_text_by_namespace={}, captured_input_tokens=0, captured_output_tokens=0, turn_stats=turn_stats, start_time=100.0, ) assert mounted assert show_calls == [True] assert turn_stats.wall_time_seconds == pytest.approx(1.0) set_active.assert_called_once_with(None) set_spinner.assert_awaited_once_with(None) tool_widget.set_rejected.assert_called_once_with() assert adapter._current_tool_messages == {} interrupted_payload = agent.aupdate_state.await_args_list[0].args[1] interrupted_msg = interrupted_payload["messages"][0] assert interrupted_msg.tool_calls[0]["id"] == "call-1" assert interrupted_msg.tool_calls[0]["name"] == "read_file" async def test_interrupt_stops_active_assistant_streams(self) -> None: """Interrupted streaming messages should not leave flush timers running.""" sync_message_content = MagicMock() assistant_msg = SimpleNamespace( id="asst-1", _content="partial response", stop_stream=AsyncMock(), ) assistant_messages = {(): assistant_msg} adapter = TextualUIAdapter( mount_message=AsyncMock(), update_status=_noop_status, request_approval=_mock_approval, set_spinner=AsyncMock(), set_active_message=MagicMock(), sync_message_content=sync_message_content, ) agent = SimpleNamespace(aupdate_state=AsyncMock()) await _handle_interrupt_cleanup( adapter=adapter, agent=agent, config={"configurable": {"thread_id": "t-1"}}, pending_text_by_namespace={(): "partial response"}, assistant_message_by_namespace=assistant_messages, captured_input_tokens=0, captured_output_tokens=0, turn_stats=SessionStats(), start_time=0.0, ) assistant_msg.stop_stream.assert_awaited_once_with() sync_message_content.assert_called_once_with("asst-1", "partial response") assert assistant_messages == {} async def test_disables_tracing_during_state_save(self) -> None: """Interrupt-cleanup `aupdate_state` calls must run with tracing disabled. Interrupt state writes (partial AI message + cancellation notice) are internal recovery mechanics. Surfacing them as standalone `UpdateState` runs in LangSmith would add noise unrelated to user-visible agent activity. """ from langsmith import get_tracing_context captured: list[object] = [] async def _capture(*_args: object, **_kwargs: object) -> None: # noqa: RUF029 captured.append(get_tracing_context().get("enabled")) agent = SimpleNamespace(aupdate_state=AsyncMock(side_effect=_capture)) adapter = TextualUIAdapter( mount_message=AsyncMock(), update_status=_noop_status, request_approval=_mock_approval, set_spinner=AsyncMock(), set_active_message=MagicMock(), ) await _handle_interrupt_cleanup( adapter=adapter, agent=agent, config={"configurable": {"thread_id": "t-1"}}, pending_text_by_namespace={}, captured_input_tokens=0, captured_output_tokens=0, turn_stats=SessionStats(), start_time=0.0, ) assert captured, "aupdate_state was never called" assert all(v is False for v in captured), ( f"tracing was not disabled: {captured}" ) async def test_disables_tracing_when_interrupted_msg_present(self) -> None: """Both `aupdate_state` calls disable tracing when interrupted_msg is set. When there is a partial AI message to save, both writes (interrupted AI message and cancellation notice) must be suppressed from LangSmith traces. """ from langsmith import get_tracing_context captured: list[object] = [] async def _capture(*_args: object, **_kwargs: object) -> None: # noqa: RUF029 captured.append(get_tracing_context().get("enabled")) tool_widget = MagicMock() tool_widget._tool_name = "read_file" tool_widget._args = {"path": "notes.txt"} agent = SimpleNamespace(aupdate_state=AsyncMock(side_effect=_capture)) adapter = TextualUIAdapter( mount_message=AsyncMock(), update_status=_noop_status, request_approval=_mock_approval, set_spinner=AsyncMock(), set_active_message=MagicMock(), ) adapter._current_tool_messages = {"call-1": tool_widget} await _handle_interrupt_cleanup( adapter=adapter, agent=agent, config={"configurable": {"thread_id": "t-1"}}, pending_text_by_namespace={}, captured_input_tokens=0, captured_output_tokens=0, turn_stats=SessionStats(), start_time=0.0, ) assert len(captured) == 2, ( f"expected 2 aupdate_state calls, got {len(captured)}" ) assert all(v is False for v in captured), ( f"tracing was not disabled: {captured}" ) class TestInterruptCleanupTokenPersist: """`_context_tokens` rides on the cancellation `aupdate_state` write.""" async def test_includes_context_tokens_in_cancellation_update(self) -> None: """The cancellation HumanMessage write carries the latest token count.""" captured: list[dict[str, Any]] = [] async def _capture(_config: object, values: dict[str, Any]) -> None: # noqa: RUF029 captured.append(values) agent = SimpleNamespace(aupdate_state=AsyncMock(side_effect=_capture)) adapter = TextualUIAdapter( mount_message=AsyncMock(), update_status=_noop_status, request_approval=_mock_approval, set_spinner=AsyncMock(), set_active_message=MagicMock(), ) await _handle_interrupt_cleanup( adapter=adapter, agent=agent, config={"configurable": {"thread_id": "t-1"}}, pending_text_by_namespace={}, captured_input_tokens=4321, captured_output_tokens=0, turn_stats=SessionStats(), start_time=0.0, ) # Only the cancellation write happens (no partial AI message in this test); # it carries both `messages` and `_context_tokens`. assert len(captured) == 1 assert captured[0]["_context_tokens"] == 4321 assert "messages" in captured[0] async def test_omits_context_tokens_when_no_usage_captured(self) -> None: """Zero tokens means we never saw `usage_metadata`; preserve the prior value.""" captured: list[dict[str, Any]] = [] async def _capture(_config: object, values: dict[str, Any]) -> None: # noqa: RUF029 captured.append(values) agent = SimpleNamespace(aupdate_state=AsyncMock(side_effect=_capture)) adapter = TextualUIAdapter( mount_message=AsyncMock(), update_status=_noop_status, request_approval=_mock_approval, set_spinner=AsyncMock(), set_active_message=MagicMock(), ) await _handle_interrupt_cleanup( adapter=adapter, agent=agent, config={"configurable": {"thread_id": "t-1"}}, pending_text_by_namespace={}, captured_input_tokens=0, captured_output_tokens=0, turn_stats=SessionStats(), start_time=0.0, ) assert len(captured) == 1 assert "_context_tokens" not in captured[0] async def test_includes_context_tokens_for_output_only_turn(self) -> None: """Output-only AI turns (no input usage) still persist a count.""" captured: list[dict[str, Any]] = [] async def _capture(_config: object, values: dict[str, Any]) -> None: # noqa: RUF029 captured.append(values) agent = SimpleNamespace(aupdate_state=AsyncMock(side_effect=_capture)) adapter = TextualUIAdapter( mount_message=AsyncMock(), update_status=_noop_status, request_approval=_mock_approval, set_spinner=AsyncMock(), set_active_message=MagicMock(), ) await _handle_interrupt_cleanup( adapter=adapter, agent=agent, config={"configurable": {"thread_id": "t-1"}}, pending_text_by_namespace={}, captured_input_tokens=0, captured_output_tokens=500, turn_stats=SessionStats(), start_time=0.0, ) assert len(captured) == 1 assert captured[0]["_context_tokens"] == 500 async def test_remote_agent_interrupt_write_carries_context_tokens(self) -> None: """Remote agents are not skipped on the interrupt-cleanup write. Locks in the deletion of the old `_persist_context_tokens` `RemoteAgent` short-circuit so a future refactor cannot silently re-introduce it. """ from deepagents_code.remote_client import RemoteAgent captured: list[dict[str, Any]] = [] async def _capture(_config: object, values: dict[str, Any]) -> None: # noqa: RUF029 captured.append(values) agent = MagicMock(spec=RemoteAgent) agent.aupdate_state = AsyncMock(side_effect=_capture) adapter = TextualUIAdapter( mount_message=AsyncMock(), update_status=_noop_status, request_approval=_mock_approval, set_spinner=AsyncMock(), set_active_message=MagicMock(), ) await _handle_interrupt_cleanup( adapter=adapter, agent=agent, config={"configurable": {"thread_id": "t-1"}}, pending_text_by_namespace={}, captured_input_tokens=1234, captured_output_tokens=88, turn_stats=SessionStats(), start_time=0.0, ) assert isinstance(agent, RemoteAgent) assert len(captured) == 1 assert captured[0]["_context_tokens"] == 1322 async def test_partial_ai_message_write_does_not_carry_tokens(self) -> None: """Only the cancellation write carries `_context_tokens`.""" captured: list[dict[str, Any]] = [] async def _capture(_config: object, values: dict[str, Any]) -> None: # noqa: RUF029 captured.append(values) tool_widget = MagicMock() tool_widget._tool_name = "read_file" tool_widget._args = {"path": "notes.txt"} agent = SimpleNamespace(aupdate_state=AsyncMock(side_effect=_capture)) adapter = TextualUIAdapter( mount_message=AsyncMock(), update_status=_noop_status, request_approval=_mock_approval, set_spinner=AsyncMock(), set_active_message=MagicMock(), ) adapter._current_tool_messages = {"call-1": tool_widget} await _handle_interrupt_cleanup( adapter=adapter, agent=agent, config={"configurable": {"thread_id": "t-1"}}, pending_text_by_namespace={}, captured_input_tokens=7777, captured_output_tokens=0, turn_stats=SessionStats(), start_time=0.0, ) assert len(captured) == 2 # First write is the interrupted AI message; should not be polluted. assert "_context_tokens" not in captured[0] # Second write is the cancellation HumanMessage; carries the token count. assert captured[1]["_context_tokens"] == 7777 class TestBuildStreamConfig: """Tests for `build_stream_config` metadata construction.""" def setup_method(self) -> None: """Clear the git-branch cache between tests.""" config_module._git_branch_cache.clear() def test_dcode_agent_fields_present(self) -> None: """Selected dcode agent metadata should be present.""" config = build_stream_config("t-456", assistant_id="my-agent") assert "assistant_id" not in config["metadata"] assert config["metadata"]["dcode_agent_name"] == "my-agent" assert config["metadata"]["agent_name"] == "my-agent" assert "updated_at" in config["metadata"] assert "cwd" in config["metadata"] def test_updated_at_is_valid_iso_timestamp(self) -> None: """`updated_at` should be a valid timezone-aware ISO 8601 timestamp.""" config = build_stream_config("t-456", assistant_id="my-agent") raw = config["metadata"]["updated_at"] assert isinstance(raw, str) parsed = datetime.fromisoformat(raw) assert parsed.tzinfo is not None def test_no_dcode_agent_fields_when_none(self) -> None: """Selected dcode agent fields should be absent when unset.""" config = build_stream_config("t-789", assistant_id=None) metadata = config["metadata"] assert "assistant_id" not in metadata assert "dcode_agent_name" not in metadata assert "agent_name" not in metadata assert "updated_at" not in metadata assert "cwd" in metadata def test_no_dcode_agent_fields_when_empty_string(self) -> None: """Empty-string `assistant_id` should be treated as absent.""" config = build_stream_config("t-000", assistant_id="") metadata = config["metadata"] assert "assistant_id" not in metadata assert "dcode_agent_name" not in metadata assert "agent_name" not in metadata assert "updated_at" not in metadata assert "cwd" in metadata def test_git_branch_included_when_available(self) -> None: """Git branch should be included in metadata when in a git repo.""" with patch( "deepagents_code.config._get_git_branch", return_value="feature-branch", ): config = build_stream_config("t-git", assistant_id="agent") assert config["metadata"]["git_branch"] == "feature-branch" def test_git_branch_absent_when_not_in_repo(self) -> None: """Git branch should be absent when not in a git repo.""" with patch( "deepagents_code.config._get_git_branch", return_value=None, ): config = build_stream_config("t-nogit", assistant_id="agent") assert "git_branch" not in config["metadata"] def test_configurable_thread_id(self) -> None: """`configurable.thread_id` should match the provided thread ID.""" config = build_stream_config("t-abc", assistant_id=None) assert config["configurable"]["thread_id"] == "t-abc" def test_sandbox_type_included_when_set(self) -> None: """Sandbox type should appear in metadata when provided.""" config = build_stream_config("t-sb", assistant_id=None, sandbox_type="daytona") assert config["metadata"]["sandbox_type"] == "daytona" def test_sandbox_type_absent_when_none(self) -> None: """Sandbox type should be absent from metadata when not provided.""" config = build_stream_config("t-nosb", assistant_id=None) assert "sandbox_type" not in config["metadata"] def test_sandbox_type_none_string_excluded(self) -> None: """The argparse sentinel `"none"` should not leak into metadata.""" config = build_stream_config("t-none", assistant_id=None, sandbox_type="none") assert "sandbox_type" not in config["metadata"] def test_no_model_keys_in_configurable(self) -> None: """Model/model_params should not be in configurable.""" config = build_stream_config("t-no-model", assistant_id=None) assert "model" not in config["configurable"] assert "model_params" not in config["configurable"] def test_versions_contains_cli_version(self) -> None: """CLI version should always be present in metadata.lc_versions.""" from deepagents_code._version import __version__ config = build_stream_config("t-ver", assistant_id=None) assert config["metadata"]["lc_versions"] == {"deepagents-code": __version__} def test_dcode_client_deepagents_version_is_diagnostic_metadata(self) -> None: """Client-side SDK version should not be reported as graph instrumentation.""" with patch("deepagents_code.config.version", return_value="1.2.3"): config = build_stream_config("t-sdk", assistant_id=None) assert config["metadata"]["dcode_client_deepagents_version"] == "1.2.3" assert "deepagents" not in config["metadata"]["lc_versions"] def test_dcode_client_deepagents_version_absent_when_metadata_missing( self, ) -> None: """Missing SDK metadata should not prevent stream config construction.""" with patch( "deepagents_code.config.version", side_effect=PackageNotFoundError("deepagents"), ): config = build_stream_config("t-missing-sdk", assistant_id=None) assert "dcode_client_deepagents_version" not in config["metadata"] def test_dcode_client_deepagents_version_does_not_import_sdk( self, monkeypatch: pytest.MonkeyPatch, ) -> None: """SDK version metadata lookup should not import the SDK package.""" for module in list(sys.modules): if module == "deepagents" or module.startswith("deepagents."): monkeypatch.delitem(sys.modules, module, raising=False) with patch("deepagents_code.config.version", return_value="1.2.3"): build_stream_config("t-no-sdk-import", assistant_id=None) assert not any( module == "deepagents" or module.startswith("deepagents.") for module in sys.modules ) def test_user_id_included_when_set(self) -> None: """DEEPAGENTS_CODE_USER_ID should appear in metadata when set.""" with patch.dict("os.environ", {"DEEPAGENTS_CODE_USER_ID": "mason"}): config = build_stream_config("t-uid", assistant_id=None) assert config["metadata"]["user_id"] == "mason" def test_user_id_absent_when_unset(self) -> None: """user_id should be absent from metadata when env var is not set.""" with patch.dict("os.environ", {"DEEPAGENTS_CODE_USER_ID": ""}): config = build_stream_config("t-nouid", assistant_id=None) assert "user_id" not in config["metadata"] class TestGetGitBranch: """Tests for `_get_git_branch` caching.""" def setup_method(self) -> None: """Clear the git-branch cache between tests.""" config_module._git_branch_cache.clear() def test_reuses_cached_branch_for_same_working_directory(self) -> None: """Repeated lookups in one repo should only resolve the branch once.""" with ( patch( "deepagents_code.config.Path.cwd", return_value=Path("/tmp/repo"), ), patch( "deepagents_code.config.resolve_git_branch", return_value="feature-branch", ) as mock_resolve, ): assert config_module._get_git_branch() == "feature-branch" assert config_module._get_git_branch() == "feature-branch" mock_resolve.assert_called_once_with("/tmp/repo") class TestGetGitBranchOSError: """Tests for _get_git_branch when Path.cwd() raises OSError.""" def setup_method(self) -> None: """Clear the git-branch cache between tests.""" config_module._git_branch_cache.clear() def test_returns_none_on_cwd_oserror(self) -> None: """_get_git_branch should return None when cwd is inaccessible.""" with patch( "deepagents_code.config.Path.cwd", side_effect=OSError("deleted"), ): assert config_module._get_git_branch() is None class TestBuildStreamConfigOSError: """Tests for build_stream_config when Path.cwd() raises OSError.""" def setup_method(self) -> None: """Clear the git-branch cache between tests.""" config_module._git_branch_cache.clear() def test_cwd_absent_on_oserror(self) -> None: """Cwd should be absent from metadata when Path.cwd() raises.""" with patch( "deepagents_code.config.Path.cwd", side_effect=OSError("deleted"), ): config = build_stream_config("t-err", assistant_id="agent") assert "cwd" not in config["metadata"] class TestIsSummarizationChunk: """Tests for `_is_summarization_chunk` detection.""" def test_returns_true_for_summarization_source(self) -> None: """Should return `True` when `lc_source` is `'summarization'`.""" metadata = {"lc_source": "summarization"} assert _is_summarization_chunk(metadata) is True def test_returns_false_for_none_metadata(self) -> None: """Should return `False` when `metadata` is `None`.""" assert _is_summarization_chunk(None) is False assert _is_summarization_chunk({}) is False def test_returns_false_for_none_lc_source(self) -> None: """Should return `False` when `lc_source` is not `'summarization'`.""" metadata_none = {"lc_source": None} assert _is_summarization_chunk(metadata_none) is False metadata_other = {"lc_source": "other"} assert _is_summarization_chunk(metadata_other) is False metadata_missing = {"other_key": "value"} assert _is_summarization_chunk(metadata_missing) is False def test_returns_false_for_unrelated_metadata(self) -> None: """Should return `False` when only unrelated keys are present.""" assert _is_summarization_chunk({"langgraph_node": "model"}) is False assert _is_summarization_chunk({"langgraph_node": None}) is False class _FakeAgent: """Minimal async stream agent used for adapter execution tests.""" def __init__(self, chunks: list[tuple]) -> None: self._chunks = chunks async def astream(self, *_: Any, **__: Any) -> AsyncIterator[tuple[Any, ...]]: """Yield preconfigured stream chunks.""" for chunk in self._chunks: yield chunk class _SequencedAgent: """Agent test double that returns a different stream per call.""" def __init__(self, streams_by_call: list[list[tuple[Any, ...]]]) -> None: self._streams_by_call = streams_by_call self.stream_inputs: list[dict | Command] = [] async def astream( self, stream_input: dict | Command, *_: Any, **__: Any, ) -> AsyncIterator[tuple[Any, ...]]: """Yield chunks for this invocation and record stream inputs.""" self.stream_inputs.append(stream_input) chunks = self._streams_by_call.pop(0) if self._streams_by_call else [] for chunk in chunks: yield chunk def _ask_user_interrupt_chunk(payload: dict[str, Any]) -> tuple[Any, ...]: """Build an updates-stream chunk containing one ask_user interrupt.""" interrupt = SimpleNamespace(id="interrupt-1", value=payload) return ((), "updates", {"__interrupt__": [interrupt]}) def _hitl_interrupt_chunk(payload: dict[str, Any]) -> tuple[Any, ...]: """Build an updates-stream chunk containing one HITL interrupt.""" interrupt = SimpleNamespace(id="interrupt-1", value=payload) return ((), "updates", {"__interrupt__": [interrupt]}) def _tool_chunk( *, name: str | None, args: str, chunk_id: str | None, index: int = 0, ) -> tuple[Any, ...]: """Build a `messages`-stream chunk carrying one streamed tool-call fragment.""" from langchain_core.messages import AIMessageChunk message = AIMessageChunk( content="", tool_call_chunks=[ { "name": name, "args": args, "id": chunk_id, "index": index, "type": "tool_call_chunk", } ], ) return ((), "messages", (message, {})) def _usage_chunk(*, input_tokens: int, output_tokens: int) -> tuple[Any, ...]: """Build a `messages`-stream chunk carrying only `usage_metadata`.""" from langchain_core.messages import AIMessageChunk message = AIMessageChunk( content="", usage_metadata={ "input_tokens": input_tokens, "output_tokens": output_tokens, "total_tokens": input_tokens + output_tokens, }, ) return ((), "messages", (message, {})) class TestExecuteTaskTextualUsageStats: """`execute_task_textual` forwards the active provider into usage stats. The per-model recording API is unit-tested directly elsewhere; this guards the call site actually reading `settings.model_provider` and threading it through `record_request`. """ async def test_records_provider_from_settings(self) -> None: """A usage chunk records the configured provider on `turn_stats`.""" async def mount_message(_: object) -> None: await asyncio.sleep(0) turn_stats = SessionStats() adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=_mock_approval, ) with patch("deepagents_code.config.settings") as mock_settings: mock_settings.model_name = "gpt-5.5" mock_settings.model_provider = "openai" await execute_task_textual( user_input="hello", agent=_FakeAgent([_usage_chunk(input_tokens=100, output_tokens=50)]), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, turn_stats=turn_stats, ) assert turn_stats.per_model["openai", "gpt-5.5"].input_tokens == 100 assert turn_stats.per_model["openai", "gpt-5.5"].output_tokens == 50 class TestExecuteTaskTextualToolCallStreaming: """Tests for incremental tool-call argument accumulation.""" async def test_fragmented_args_mount_once_when_json_completes(self) -> None: """Args streamed across many fragments parse once the JSON is whole. The tool row mounts a single time with fully accumulated args, even though the JSON arrives split across several chunks. """ mounted: list[object] = [] async def mount_message(widget: object) -> None: await asyncio.sleep(0) mounted.append(widget) # Split a JSON object across fragments; only the last one closes it. chunks = [ _tool_chunk(name="edit_file", args='{"path": ', chunk_id="t1"), _tool_chunk(name=None, args='"a.py", ', chunk_id=None), _tool_chunk(name=None, args='"content": "x"}', chunk_id=None), ] adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=_mock_approval, ) await execute_task_textual( user_input="hello", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) tool_msgs = [m for m in mounted if isinstance(m, ToolCallMessage)] assert len(tool_msgs) == 1 assert tool_msgs[0]._tool_name == "edit_file" assert tool_msgs[0]._args == {"path": "a.py", "content": "x"} async def test_incomplete_args_do_not_mount(self) -> None: """A tool row stays unmounted while its JSON args are still partial.""" mounted: list[object] = [] async def mount_message(widget: object) -> None: await asyncio.sleep(0) mounted.append(widget) # JSON never closes — the row must not mount with partial args. chunks = [ _tool_chunk(name="edit_file", args='{"path": ', chunk_id="t1"), _tool_chunk(name=None, args='"a.py"', chunk_id=None), ] adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=_mock_approval, ) await execute_task_textual( user_input="hello", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) assert not [m for m in mounted if isinstance(m, ToolCallMessage)] async def test_scalar_args_mount_eagerly_when_complete(self) -> None: """Non-object JSON args parse as soon as the scalar is whole. Scalars never close with `}`/`]`, so the bracket heuristic that defers large objects never fires for them; they must still mount (wrapped as `{"value": ...}`) once the accumulated fragment is valid JSON. """ mounted: list[object] = [] async def mount_message(widget: object) -> None: await asyncio.sleep(0) mounted.append(widget) # A JSON string split so the first fragment is not yet valid JSON. chunks = [ _tool_chunk(name="echo", args='"hel', chunk_id="t1"), _tool_chunk(name=None, args='lo"', chunk_id=None), ] adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=_mock_approval, ) await execute_task_textual( user_input="hello", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) tool_msgs = [m for m in mounted if isinstance(m, ToolCallMessage)] assert len(tool_msgs) == 1 assert tool_msgs[0]._args == {"value": "hello"} async def test_dict_args_resolve_without_reparsing(self) -> None: """A complete `tool_call` block mounts with its dict args verbatim.""" mounted: list[object] = [] async def mount_message(widget: object) -> None: await asyncio.sleep(0) mounted.append(widget) chunks = [ ( (), "messages", (_tool_call_message("read_file", {"path": "a.py"}, "t1"), {}), ), ] adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=_mock_approval, ) await execute_task_textual( user_input="hello", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) tool_msgs = [m for m in mounted if isinstance(m, ToolCallMessage)] assert len(tool_msgs) == 1 assert tool_msgs[0]._args == {"path": "a.py"} async def test_interleaved_fragments_accumulate_per_tool(self) -> None: """Fragments for two concurrent tool calls accumulate independently. Each tool call carries a distinct stream index, so interleaved argument fragments must not bleed across buffers. """ mounted: list[object] = [] async def mount_message(widget: object) -> None: await asyncio.sleep(0) mounted.append(widget) # Two tools (index 0 and 1) with interleaved argument fragments. chunks = [ _tool_chunk(name="read_file", args='{"path": ', chunk_id="t0", index=0), _tool_chunk(name="grep", args='{"pattern": ', chunk_id="t1", index=1), _tool_chunk(name=None, args='"a.py"}', chunk_id=None, index=0), _tool_chunk(name=None, args='"x"}', chunk_id=None, index=1), ] adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=_mock_approval, ) await execute_task_textual( user_input="hello", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) tool_msgs = [m for m in mounted if isinstance(m, ToolCallMessage)] by_name = {m._tool_name: m._args for m in tool_msgs} assert by_name == { "read_file": {"path": "a.py"}, "grep": {"pattern": "x"}, } class TestExecuteTaskTextualSummarizationFeedback: """Tests for summarization spinner and notification feedback.""" async def test_spinner_transitions_for_summarization_stream(self) -> None: """Spinner should move Thinking -> Offloading -> Thinking.""" statuses: list[str | None] = [] async def record_spinner(status: str | None) -> None: await asyncio.sleep(0) statuses.append(status) async def mount_message(_widget: object) -> None: await asyncio.sleep(0) chunks = [ ( (), "messages", (AIMessage(content="summary chunk"), {"lc_source": "summarization"}), ), ((), "messages", (HumanMessage(content="regular chunk"), {})), ] adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=_mock_approval, set_spinner=record_spinner, ) await execute_task_textual( user_input="hello", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) assert statuses[0] == "Thinking" assert "Offloading" in statuses assert statuses[-1] == "Thinking" async def test_mounts_summarization_notification_on_regular_chunk(self) -> None: """Notification should render when regular chunks resume after summarization.""" statuses: list[str | None] = [] mounted_widgets: list[object] = [] async def record_spinner(status: str | None) -> None: await asyncio.sleep(0) statuses.append(status) async def mount_message(widget: object) -> None: await asyncio.sleep(0) mounted_widgets.append(widget) chunks = [ ( (), "messages", (AIMessage(content="summary chunk"), {"lc_source": "summarization"}), ), # Regular chunk from the actual model — signals summarization ended. ((), "messages", (HumanMessage(content="regular"), {})), ] adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=_mock_approval, set_spinner=record_spinner, ) await execute_task_textual( user_input="hello", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) assert any( isinstance(widget, SummarizationMessage) for widget in mounted_widgets ) async def test_mounts_notification_when_stream_ends_mid_summarization(self) -> None: """Notification should still render if stream exhausts during summarization.""" mounted_widgets: list[object] = [] async def record_spinner(_status: str | None) -> None: await asyncio.sleep(0) async def mount_message(widget: object) -> None: await asyncio.sleep(0) mounted_widgets.append(widget) # Only summarization chunks, no regular chunks follow. chunks = [ ( (), "messages", (AIMessage(content="summary chunk"), {"lc_source": "summarization"}), ), ] adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=_mock_approval, set_spinner=record_spinner, ) await execute_task_textual( user_input="hello", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) assert any( isinstance(widget, SummarizationMessage) for widget in mounted_widgets ) def _tool_call_message( name: str, args: dict[str, Any], tool_id: str ) -> SimpleNamespace: """Build a message-like object with content_blocks containing one tool call.""" return SimpleNamespace( content_blocks=[ {"type": "tool_call", "name": name, "args": args, "id": tool_id} ] ) def _text_message(text: str) -> SimpleNamespace: """Build a message-like object with content_blocks containing one text block.""" return SimpleNamespace(content_blocks=[{"type": "text", "text": text}]) class TestExecuteTaskTextualParallelToolSpinner: """Regression tests for #1796: premature spinner with parallel tools.""" async def test_spinner_not_shown_until_all_parallel_tools_complete(self) -> None: """With two parallel tools, Thinking appears only at start and after last.""" statuses: list[str | None] = [] async def record_spinner(status: str | None) -> None: await asyncio.sleep(0) statuses.append(status) async def mount_message(_widget: object) -> None: await asyncio.sleep(0) chunks = [ ( (), "messages", ( _tool_call_message("task", {"task": "a"}, "tool-a"), {}, ), ), ( (), "messages", ( _tool_call_message("task", {"task": "b"}, "tool-b"), {}, ), ), ( (), "messages", ( ToolMessage(content="result a", tool_call_id="tool-a"), {}, ), ), ( (), "messages", ( ToolMessage(content="result b", tool_call_id="tool-b"), {}, ), ), ] adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=_mock_approval, set_spinner=record_spinner, ) await execute_task_textual( user_input="hello", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=True), adapter=adapter, ) assert statuses[0] == "Thinking" thinking_count = sum(1 for s in statuses if s == "Thinking") assert thinking_count == 2, ( "Expected exactly 2 Thinking calls (start + after last tool); " f"got {thinking_count}: {statuses}" ) async def test_on_tool_complete_fires_per_tool_message(self) -> None: """`on_tool_complete` should fire once per `ToolMessage`, even in parallel.""" tool_complete = MagicMock() tc = _tool_call_message chunks = [ ((), "messages", (tc("task", {"task": "a"}, "tool-a"), {})), ((), "messages", (tc("task", {"task": "b"}, "tool-b"), {})), ((), "messages", (ToolMessage(content="a", tool_call_id="tool-a"), {})), ((), "messages", (ToolMessage(content="b", tool_call_id="tool-b"), {})), ] adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, on_tool_complete=tool_complete, ) await execute_task_textual( user_input="hi", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=True), adapter=adapter, ) assert tool_complete.call_count == 2 async def test_on_tool_complete_exception_is_swallowed(self) -> None: """A raising `on_tool_complete` must not break agent streaming.""" tc = _tool_call_message chunks = [ ((), "messages", (tc("task", {"task": "a"}, "tool-a"), {})), ((), "messages", (ToolMessage(content="a", tool_call_id="tool-a"), {})), ] adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, on_tool_complete=MagicMock(side_effect=RuntimeError("boom")), ) await execute_task_textual( user_input="hi", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=True), adapter=adapter, ) async def test_spinner_shown_after_single_tool_completes(self) -> None: """Spinner should show Thinking after the only tool completes.""" statuses: list[str | None] = [] async def record_spinner(status: str | None) -> None: await asyncio.sleep(0) statuses.append(status) chunks = [ ( (), "messages", ( _tool_call_message("ls", {"path": "."}, "tool-1"), {}, ), ), ( (), "messages", ( ToolMessage(content="file1.py", tool_call_id="tool-1"), {}, ), ), ] adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, set_spinner=record_spinner, ) await execute_task_textual( user_input="list files", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=True), adapter=adapter, ) assert statuses[-1] == "Thinking" async def test_edit_file_tool_keeps_thinking_spinner_while_pending(self) -> None: """`edit_file` should not leave a visual gap before approval/execution.""" statuses: list[str | None] = [] async def record_spinner(status: str | None) -> None: await asyncio.sleep(0) statuses.append(status) chunks = [ ( (), "messages", ( _tool_call_message( "edit_file", { "file_path": "example.py", "old_string": "old", "new_string": "new", }, "tool-1", ), {}, ), ), ( (), "messages", ( ToolMessage(content="edited", tool_call_id="tool-1"), {}, ), ), ] adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, set_spinner=record_spinner, ) await execute_task_textual( user_input="edit the file", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=True), adapter=adapter, ) assert statuses[:2] == ["Thinking", "Thinking"] assert None not in statuses async def test_auto_executed_tool_shows_running_at_mount(self) -> None: """Auto-executed tools (no approval) spin immediately when mounted. Regression guard: read-only tools such as `grep`/`glob` previously sat visually idle from mount until their result arrived. The stream here ends right after the tool call (no result), so the row is observed in its mount-time state. """ chunks = [ ( (), "messages", (_tool_call_message("grep", {"pattern": "foo"}, "tool-1"), {}), ), ] adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, ) await execute_task_textual( user_input="search", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=True), adapter=adapter, ) tool_msg = adapter._current_tool_messages["tool-1"] assert tool_msg._status == "running" async def test_edit_file_does_not_get_per_tool_spinner_at_mount(self) -> None: """`edit_file` relies on the global Thinking spinner, not a per-tool one. Negative counterpart to the auto-executed case: tools in `_TOOL_CALLS_KEEP_THINKING_SPINNER` must NOT be flipped to "running" at mount, or they would show a duplicate spinner alongside "Thinking". """ chunks = [ ( (), "messages", ( _tool_call_message( "edit_file", { "file_path": "example.py", "old_string": "old", "new_string": "new", }, "tool-1", ), {}, ), ), ] adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, ) await execute_task_textual( user_input="edit", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=True), adapter=adapter, ) tool_msg = adapter._current_tool_messages["tool-1"] assert tool_msg._status != "running" async def test_spinner_with_three_parallel_tools_out_of_order(self) -> None: """Three parallel tools completed out of order; Thinking after all.""" statuses: list[str | None] = [] async def record_spinner(status: str | None) -> None: await asyncio.sleep(0) statuses.append(status) tc = _tool_call_message chunks = [ ((), "messages", (tc("task", {"task": "a"}, "tool-a"), {})), ((), "messages", (tc("task", {"task": "b"}, "tool-b"), {})), ((), "messages", (tc("task", {"task": "c"}, "tool-c"), {})), # Complete out of dispatch order: B, A, C ( (), "messages", ( ToolMessage( content="result b", tool_call_id="tool-b", ), {}, ), ), ( (), "messages", ( ToolMessage( content="result a", tool_call_id="tool-a", ), {}, ), ), ( (), "messages", ( ToolMessage( content="result c", tool_call_id="tool-c", ), {}, ), ), ] adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, set_spinner=record_spinner, ) await execute_task_textual( user_input="hello", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=True), adapter=adapter, ) thinking_count = sum(1 for s in statuses if s == "Thinking") assert thinking_count == 2, ( "Expected exactly 2 Thinking calls (start + after last tool); " f"got {thinking_count}: {statuses}" ) async def test_spinner_recovers_with_untracked_tool_id(self) -> None: """Spinner still shows Thinking with an untracked tool_call_id.""" statuses: list[str | None] = [] async def record_spinner(status: str | None) -> None: await asyncio.sleep(0) statuses.append(status) tc = _tool_call_message chunks = [ ((), "messages", (tc("task", {"task": "a"}, "tool-a"), {})), # Result with a tool_call_id that was never dispatched ( (), "messages", ( ToolMessage( content="result a", tool_call_id="tool-a", ), {}, ), ), ( (), "messages", ( ToolMessage( content="unknown", tool_call_id="tool-unknown", ), {}, ), ), ] adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, set_spinner=record_spinner, ) await execute_task_textual( user_input="hello", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=True), adapter=adapter, ) # After the tracked tool completes, dict is empty so spinner should show. # The untracked ToolMessage should not break spinner recovery. thinking_calls = [i for i, s in enumerate(statuses) if s == "Thinking"] assert len(thinking_calls) >= 2, ( f"Expected at least 2 Thinking calls; got {len(thinking_calls)}: {statuses}" ) class TestExecuteTaskTextualTextThenToolSpinner: """Regression tests: spinner must stay visible between text and tool call. When the assistant streams explanatory text and then emits a tool call, the model often pauses between finishing the text and producing the tool call. The spinner should remain visible during that pause rather than disappearing as soon as the first text chunk arrives. """ async def test_spinner_not_hidden_when_text_chunk_arrives(self) -> None: """Streaming a text block must not hide the Thinking spinner.""" statuses: list[str | None] = [] async def record_spinner(status: str | None) -> None: await asyncio.sleep(0) statuses.append(status) chunks = [ ((), "messages", (_text_message("Now I'll call a tool..."), {})), ((), "messages", (_tool_call_message("ls", {"path": "."}, "tool-1"), {})), ((), "messages", (ToolMessage(content="ok", tool_call_id="tool-1"), {})), ] adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, set_spinner=record_spinner, ) # Patch AssistantMessage so it doesn't require a real Textual DOM. fake_msg = AsyncMock() fake_msg.id = "asst-test" with patch( "deepagents_code.textual_adapter.AssistantMessage", return_value=fake_msg ): await execute_task_textual( user_input="hi", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=True), adapter=adapter, ) # Expected sequence: # 1. "Thinking" before astream # 2. "Thinking" after mounting the streaming AssistantMessage # (re-anchor the spinner below the message so the user still # sees activity if the model pauses before the tool call) # 3. None when the tool call mounts # 4. "Thinking" after the tool result assert statuses[0] == "Thinking" assert statuses[1] == "Thinking" assert None in statuses assert statuses[-1] == "Thinking" # The spinner must never be hidden before the tool call arrives. first_none = statuses.index(None) text_thinking_seen = statuses[:first_none].count("Thinking") >= 2 assert text_thinking_seen, ( f"Spinner was hidden during text streaming before tool call: {statuses}" ) async def test_spinner_reanchors_for_text_after_tool_cycle(self) -> None: """Text -> tool_call -> tool_result -> text must re-anchor the spinner. After a tool cycle completes, the tool_call handler pops the previous AssistantMessage from `assistant_message_by_namespace`, so the next text chunk mounts a fresh widget. The new re-anchor call at `textual_adapter.py:780-784` must fire for that second text burst so the spinner stays visible between it and any follow-up tool call. """ statuses: list[str | None] = [] async def record_spinner(status: str | None) -> None: await asyncio.sleep(0) statuses.append(status) chunks = [ ((), "messages", (_text_message("First, I'll inspect..."), {})), ((), "messages", (_tool_call_message("ls", {"path": "."}, "tool-1"), {})), ((), "messages", (ToolMessage(content="ok", tool_call_id="tool-1"), {})), ((), "messages", (_text_message("Now the second step..."), {})), ] adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, set_spinner=record_spinner, ) fake_msg = AsyncMock() fake_msg.id = "asst-test" with patch( "deepagents_code.textual_adapter.AssistantMessage", return_value=fake_msg ): await execute_task_textual( user_input="hi", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=True), adapter=adapter, ) # Expected Thinking calls: # 1. Before astream (line 517) # 2. After first AssistantMessage mount (re-anchor, line 784) # 3. After tool result (line 705) # 4. After second AssistantMessage mount (re-anchor again) thinking_count = sum(1 for s in statuses if s == "Thinking") assert thinking_count >= 4, ( f"Expected at least 4 Thinking calls including re-anchors after " f"each text mount; got {thinking_count}: {statuses}" ) async def test_spinner_reanchor_skipped_while_tools_pending(self) -> None: """The re-anchor must be gated on `not _current_tool_messages`. Contrived sequence: a tool call mounts (populating `_current_tool_messages`), then a text chunk arrives before the tool result. The new re-anchor logic must NOT call `_set_spinner("Thinking")` in that window — the tool-call widget is the dominant progress indicator. """ statuses: list[str | None] = [] async def record_spinner(status: str | None) -> None: await asyncio.sleep(0) statuses.append(status) chunks = [ ((), "messages", (_tool_call_message("ls", {"path": "."}, "tool-1"), {})), ((), "messages", (_text_message("Meanwhile..."), {})), ((), "messages", (ToolMessage(content="ok", tool_call_id="tool-1"), {})), ] adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, set_spinner=record_spinner, ) fake_msg = AsyncMock() fake_msg.id = "asst-test" with patch( "deepagents_code.textual_adapter.AssistantMessage", return_value=fake_msg ): await execute_task_textual( user_input="hi", agent=_FakeAgent(chunks), assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=True), adapter=adapter, ) # Thinking calls should be: # 1. Before astream # 2. After tool result (guard is back to empty) # The re-anchor must NOT fire while the tool is in flight. thinking_count = sum(1 for s in statuses if s == "Thinking") assert thinking_count == 2, ( f"Expected 2 Thinking calls (start + after tool); got " f"{thinking_count}: {statuses}" ) class TestExecuteTaskTextualHITLShellSuppression: """Tests for shell-tool widget suppression during HITL approval.""" async def _run_with_decision( self, *, tool_call_name: str, tool_call_id: str, approval_decision: dict[str, Any], extra_tool_calls: list[tuple[str, dict[str, Any], str]] | None = None, ) -> tuple[ TextualUIAdapter, list[object], dict[str, tuple[bool, bool, str]], ]: """Drive a HITL flow and snapshot widget visibility during the await. Returns the adapter, the mounted widgets, and a mapping of `tool_call_id -> (display, _awaiting_approval, _status)` captured while the approval future is pending. The status entry locks in the pause behavior: tools start their spinner at mount but are reverted to `pending` while blocked on the approval decision. """ mounted: list[object] = [] snapshots: dict[str, tuple[bool, bool, str]] = {} async def mount_message(widget: object) -> None: await asyncio.sleep(0) mounted.append(widget) future: asyncio.Future[object] = asyncio.Future() async def request_approval( _action_requests: list[dict[str, Any]], _assistant_id: str | None, ) -> asyncio.Future[object]: await asyncio.sleep(0) for tid, tool_msg in adapter._current_tool_messages.items(): snapshots[tid] = ( bool(tool_msg.display), tool_msg._awaiting_approval, tool_msg._status, ) future.set_result(approval_decision) return future message_chunks: list[tuple[Any, ...]] = [ ( (), "messages", ( _tool_call_message( tool_call_name, {"command": "echo hi"}, tool_call_id ), {}, ), ) ] for name, args, tid in extra_tool_calls or []: message_chunks.append( ((), "messages", (_tool_call_message(name, args, tid), {})) ) action_requests = [{"name": tool_call_name, "args": {"command": "echo hi"}}] for name, args, _tid in extra_tool_calls or []: action_requests.append({"name": name, "args": args}) agent = _SequencedAgent( streams_by_call=[ [ *message_chunks, _hitl_interrupt_chunk( { "action_requests": action_requests, "review_configs": [ { "action_name": req["name"], "allowed_decisions": ["approve", "reject"], } for req in action_requests ], } ), ], [], ] ) adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=request_approval, ) await execute_task_textual( user_input="hello", agent=agent, assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) return adapter, mounted, snapshots async def test_shell_tool_widget_suppressed_during_approval(self) -> None: """`execute` widget should be hidden during the await and restored after.""" _adapter, mounted, snapshots = await self._run_with_decision( tool_call_name="execute", tool_call_id="tool-shell", approval_decision={"type": "approve"}, ) tool_rows = [w for w in mounted if isinstance(w, ToolCallMessage)] assert len(tool_rows) == 1 # While the future was pending, the widget was hidden and its spinner # paused (reverted from the mount-time "running" to "pending"). assert snapshots["tool-shell"] == (False, True, "pending") # After the finally block, it was restored and the spinner resumed # (the resumed stream is empty, so the row never reaches a result). assert tool_rows[0].display is True assert tool_rows[0]._awaiting_approval is False assert tool_rows[0]._status == "running" async def test_non_shell_tool_widget_not_suppressed(self) -> None: """`read_file` widget should stay visible — only shell tools are hidden.""" _adapter, mounted, snapshots = await self._run_with_decision( tool_call_name="read_file", tool_call_id="tool-read", approval_decision={"type": "approve"}, ) tool_rows = [w for w in mounted if isinstance(w, ToolCallMessage)] assert len(tool_rows) == 1 # Visible the whole time, never marked as awaiting approval, but the # spinner is paused to "pending" while the decision is outstanding. assert snapshots["tool-read"] == (True, False, "pending") assert tool_rows[0].display is True assert tool_rows[0]._awaiting_approval is False # Resumed to "running" after approval (resumed stream yields no result). assert tool_rows[0]._status == "running" async def test_batch_approval_keeps_all_widgets_visible(self) -> None: """Batched approvals (>1 request) must not hide any tool widget. The approval dialog only renders a per-tool command preview for single-tool approvals. For batches it shows just a count header, so suppressing the streamed rows would leave the user with no preview of what's being approved. """ _adapter, _mounted, snapshots = await self._run_with_decision( tool_call_name="execute", tool_call_id="tool-shell", approval_decision={"type": "approve"}, extra_tool_calls=[("read_file", {"path": "notes.txt"}, "tool-read")], ) assert snapshots["tool-shell"] == (True, False, "pending") assert snapshots["tool-read"] == (True, False, "pending") async def test_batch_of_shell_tools_keeps_all_widgets_visible(self) -> None: """Multiple parallel `execute` calls: all rows stay visible. Regression guard: the batch approval dialog does not render per-tool commands, so hiding every `execute` row left users with only a generic "N Tool Calls Require Approval" header. """ _adapter, _mounted, snapshots = await self._run_with_decision( tool_call_name="execute", tool_call_id="tool-shell-1", approval_decision={"type": "approve"}, extra_tool_calls=[ ("execute", {"command": "echo bye"}, "tool-shell-2"), ], ) assert snapshots["tool-shell-1"] == (True, False, "pending") assert snapshots["tool-shell-2"] == (True, False, "pending") async def test_shell_widget_restored_when_approval_raises(self) -> None: """`finally` must restore the widget even if approval raises.""" mounted: list[object] = [] async def mount_message(widget: object) -> None: await asyncio.sleep(0) mounted.append(widget) async def request_approval( _action_requests: list[dict[str, Any]], _assistant_id: str | None, ) -> asyncio.Future[object]: await asyncio.sleep(0) msg = "boom" raise RuntimeError(msg) agent = _SequencedAgent( streams_by_call=[ [ ( (), "messages", ( _tool_call_message( "execute", {"command": "echo hi"}, "tool-shell" ), {}, ), ), _hitl_interrupt_chunk( { "action_requests": [ {"name": "execute", "args": {"command": "echo hi"}} ], "review_configs": [ { "action_name": "execute", "allowed_decisions": ["approve", "reject"], } ], } ), ], [], ] ) adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=request_approval, ) with pytest.raises(RuntimeError, match="boom"): await execute_task_textual( user_input="hello", agent=agent, assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) tool_rows = [w for w in mounted if isinstance(w, ToolCallMessage)] assert len(tool_rows) == 1 assert tool_rows[0].display is True assert tool_rows[0]._awaiting_approval is False class TestExecuteTaskTextualAskUser: """Tests for ask_user interrupt handling in the Textual adapter.""" async def test_ask_user_interrupt_mounts_tool_call_row(self) -> None: """ask_user interrupts should mount the tool row before the prompt.""" mounted: list[object] = [] future: asyncio.Future[AskUserWidgetResult] = asyncio.Future() future.set_result({"type": "answered", "answers": ["Alice"]}) async def mount_message(widget: object) -> None: await asyncio.sleep(0) mounted.append(widget) async def request_ask_user( _questions: list[Question], ) -> asyncio.Future[AskUserWidgetResult] | None: await asyncio.sleep(0) return future agent = _SequencedAgent( streams_by_call=[ [ _ask_user_interrupt_chunk( { "type": "ask_user", "questions": [{"question": "Name?", "type": "text"}], "tool_call_id": "tool-1", } ) ], [], ] ) adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=_mock_approval, request_ask_user=request_ask_user, ) await execute_task_textual( user_input="hello", agent=agent, assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) tool_rows = [ widget for widget in mounted if isinstance(widget, ToolCallMessage) ] assert len(tool_rows) == 1 tool_row = tool_rows[0] assert tool_row.tool_name == "ask_user" assert tool_row.has_expandable_args is True # Answered cleanup pops the row from `_current_tool_messages`. assert "tool-1" not in adapter._current_tool_messages async def test_ask_user_mount_failure_does_not_register_tool_id(self) -> None: """Mount failure should not poison `displayed_tool_ids` on the adapter.""" async def mount_message(_widget: object) -> None: await asyncio.sleep(0) msg = "mount failed" raise RuntimeError(msg) future: asyncio.Future[AskUserWidgetResult] = asyncio.Future() future.set_result({"type": "answered", "answers": ["Alice"]}) async def request_ask_user( _questions: list[Question], ) -> asyncio.Future[AskUserWidgetResult] | None: await asyncio.sleep(0) return future agent = _SequencedAgent( streams_by_call=[ [ _ask_user_interrupt_chunk( { "type": "ask_user", "questions": [{"question": "Name?", "type": "text"}], "tool_call_id": "tool-1", } ) ], [], ] ) adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=_mock_approval, request_ask_user=request_ask_user, ) await execute_task_textual( user_input="hello", agent=agent, assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) # The flow continued, resumed with the answer, and never registered the # broken tool row. assert "tool-1" not in adapter._current_tool_messages async def test_ask_user_duplicate_interrupt_only_mounts_once(self) -> None: """Re-emitting the same `tool_call_id` should not double-mount.""" mounted: list[object] = [] future: asyncio.Future[AskUserWidgetResult] = asyncio.Future() future.set_result({"type": "answered", "answers": ["Alice"]}) async def mount_message(widget: object) -> None: await asyncio.sleep(0) mounted.append(widget) async def request_ask_user( _questions: list[Question], ) -> asyncio.Future[AskUserWidgetResult] | None: await asyncio.sleep(0) return future payload = { "type": "ask_user", "questions": [{"question": "Name?", "type": "text"}], "tool_call_id": "tool-dedup", } agent = _SequencedAgent( streams_by_call=[ [ _ask_user_interrupt_chunk(payload), _ask_user_interrupt_chunk(payload), ], [], ] ) adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=_mock_approval, request_ask_user=request_ask_user, ) await execute_task_textual( user_input="hello", agent=agent, assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) tool_rows = [w for w in mounted if isinstance(w, ToolCallMessage)] assert len(tool_rows) == 1 async def test_ask_user_cancelled_marks_row_rejected_and_halts(self) -> None: """Cancelled result should reject the row and not resume generation.""" mounted: list[object] = [] token_events: list[str] = [] future: asyncio.Future[AskUserWidgetResult] = asyncio.Future() future.set_result({"type": "cancelled"}) async def mount_message(widget: object) -> None: await asyncio.sleep(0) mounted.append(widget) async def request_ask_user( _questions: list[Question], ) -> asyncio.Future[AskUserWidgetResult] | None: await asyncio.sleep(0) return future agent = _SequencedAgent( streams_by_call=[ [ _ask_user_interrupt_chunk( { "type": "ask_user", "questions": [{"question": "Name?", "type": "text"}], "tool_call_id": "tool-1", } ) ], [], ] ) adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=_mock_approval, request_ask_user=request_ask_user, ) adapter._on_tokens_pending = lambda: token_events.append("pending") adapter._on_tokens_show = lambda *, approximate=False: token_events.append( f"show:{approximate}" ) await execute_task_textual( user_input="hello", agent=agent, assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) assert len(agent.stream_inputs) == 1 assert "tool-1" not in adapter._current_tool_messages app_messages = [widget for widget in mounted if isinstance(widget, AppMessage)] assert len(app_messages) == 1 assert "Question cancelled" in str(app_messages[0]._content) assert token_events == ["pending", "show:False"] async def test_hitl_rejection_restores_token_display_before_halt(self) -> None: """Rejected approval should restore tokens before returning early.""" mounted: list[object] = [] token_events: list[str] = [] future: asyncio.Future[object] = asyncio.Future() future.set_result({"type": "reject"}) async def mount_message(widget: object) -> None: await asyncio.sleep(0) mounted.append(widget) async def request_approval( _action_requests: list[dict[str, Any]], _assistant_id: str | None, ) -> asyncio.Future[object]: await asyncio.sleep(0) return future agent = _SequencedAgent( streams_by_call=[ [ _hitl_interrupt_chunk( { "action_requests": [ {"name": "read_file", "args": {"path": "notes.txt"}} ], "review_configs": [ { "action_name": "read_file", "allowed_decisions": ["approve", "reject"], } ], } ) ], [], ] ) adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=request_approval, ) adapter._on_tokens_pending = lambda: token_events.append("pending") adapter._on_tokens_show = lambda *, approximate=False: token_events.append( f"show:{approximate}" ) await execute_task_textual( user_input="hello", agent=agent, assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) assert len(agent.stream_inputs) == 1 app_messages = [widget for widget in mounted if isinstance(widget, AppMessage)] assert len(app_messages) == 1 assert "Command rejected" in str(app_messages[0]._content) assert token_events == ["pending", "show:False"] async def test_hitl_rejection_with_reason_resumes_agent(self) -> None: """Rejected approval with a reason should resume so the agent can react.""" mounted: list[object] = [] future: asyncio.Future[object] = asyncio.Future() future.set_result({"type": "reject", "message": "use a safer command"}) async def mount_message(widget: object) -> None: await asyncio.sleep(0) mounted.append(widget) async def request_approval( _action_requests: list[dict[str, Any]], _assistant_id: str | None, ) -> asyncio.Future[object]: await asyncio.sleep(0) return future agent = _SequencedAgent( streams_by_call=[ [ _hitl_interrupt_chunk( { "action_requests": [ {"name": "execute", "args": {"command": "rm file"}} ], "review_configs": [ { "action_name": "execute", "allowed_decisions": ["approve", "reject"], } ], } ) ], [], ] ) adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=request_approval, ) await execute_task_textual( user_input="hello", agent=agent, assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) assert len(agent.stream_inputs) == 2 resume_cmd = agent.stream_inputs[1] assert isinstance(resume_cmd, Command) resume_payload = cast("dict[str, dict[str, Any]]", resume_cmd.resume) decisions = resume_payload["interrupt-1"]["decisions"] assert decisions == [{"type": "reject", "message": "use a safer command"}] app_messages = [widget for widget in mounted if isinstance(widget, AppMessage)] assert not any("Command rejected" in str(msg._content) for msg in app_messages) async def test_ask_user_invalid_answers_payload_marks_row_error(self) -> None: """Non-list answers should mark row as error and pop it.""" mounted: list[ToolCallMessage] = [] error_calls: list[str] = [] future: asyncio.Future[object] = asyncio.Future() future.set_result({"type": "answered", "answers": "not-a-list"}) async def mount_message(widget: object) -> None: await asyncio.sleep(0) if isinstance(widget, ToolCallMessage): original = widget.set_error def _capture(error: str) -> None: error_calls.append(error) original(error) widget.set_error = _capture # ty: ignore mounted.append(widget) async def request_ask_user( _questions: list[Question], ) -> asyncio.Future[object] | None: await asyncio.sleep(0) return future agent = _SequencedAgent( streams_by_call=[ [ _ask_user_interrupt_chunk( { "type": "ask_user", "questions": [{"question": "Name?", "type": "text"}], "tool_call_id": "tool-1", } ) ], [], ] ) adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=_mock_approval, # This test intentionally returns a malformed widget payload. request_ask_user=cast("Any", request_ask_user), ) await execute_task_textual( user_input="hello", agent=agent, assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) resume_cmd = agent.stream_inputs[1] assert isinstance(resume_cmd, Command) resume_payload = cast("dict[str, dict[str, Any]]", resume_cmd.resume) assert resume_payload["interrupt-1"]["status"] == "error" assert ( resume_payload["interrupt-1"]["error"] == "invalid ask_user answers payload" ) assert len(mounted) == 1 assert "invalid ask_user answers payload" in error_calls assert "tool-1" not in adapter._current_tool_messages async def test_ask_user_unsupported_marks_row_error(self) -> None: """When no callback is registered, the mounted row gets an error.""" mounted: list[ToolCallMessage] = [] error_calls: list[str] = [] async def mount_message(widget: object) -> None: await asyncio.sleep(0) if isinstance(widget, ToolCallMessage): original = widget.set_error def _capture(error: str) -> None: error_calls.append(error) original(error) widget.set_error = _capture # ty: ignore mounted.append(widget) agent = _SequencedAgent( streams_by_call=[ [ _ask_user_interrupt_chunk( { "type": "ask_user", "questions": [{"question": "Name?", "type": "text"}], "tool_call_id": "tool-1", } ) ], [], ] ) adapter = TextualUIAdapter( mount_message=mount_message, update_status=_noop_status, request_approval=_mock_approval, request_ask_user=None, ) await execute_task_textual( user_input="hello", agent=agent, assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) assert len(mounted) == 1 assert "ask_user not supported by this UI" in error_calls assert "tool-1" not in adapter._current_tool_messages async def test_request_ask_user_returning_none_is_reported_as_error(self) -> None: """A `None` callback result should resume with explicit error status.""" async def request_ask_user( _questions: list[Question], ) -> asyncio.Future[AskUserWidgetResult] | None: await asyncio.sleep(0) return None agent = _SequencedAgent( streams_by_call=[ [ _ask_user_interrupt_chunk( { "type": "ask_user", "questions": [{"question": "Name?", "type": "text"}], "tool_call_id": "tool-1", } ) ], [], ] ) adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, request_ask_user=request_ask_user, ) await execute_task_textual( user_input="hello", agent=agent, assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) assert len(agent.stream_inputs) >= 2 resume_cmd = agent.stream_inputs[1] assert isinstance(resume_cmd, Command) resume_payload = cast("dict[str, dict[str, Any]]", resume_cmd.resume) ask_user_resume = resume_payload["interrupt-1"] assert ask_user_resume["status"] == "error" assert ask_user_resume["error"] == "ask_user callback returned no response" assert ask_user_resume["answers"] == [""] async def test_request_ask_user_mount_error_is_not_treated_as_cancel(self) -> None: """UI mount failures should resume with explicit error status.""" async def request_ask_user( _questions: list[Question], ) -> asyncio.Future[AskUserWidgetResult] | None: await asyncio.sleep(0) msg = "boom" raise RuntimeError(msg) agent = _SequencedAgent( streams_by_call=[ [ _ask_user_interrupt_chunk( { "type": "ask_user", "questions": [{"question": "Name?", "type": "text"}], "tool_call_id": "tool-1", } ) ], [], ] ) adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, request_ask_user=request_ask_user, ) await execute_task_textual( user_input="hello", agent=agent, assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) resume_cmd = agent.stream_inputs[1] assert isinstance(resume_cmd, Command) resume_payload = cast("dict[str, dict[str, Any]]", resume_cmd.resume) ask_user_resume = resume_payload["interrupt-1"] assert ask_user_resume["status"] == "error" assert ask_user_resume["error"] == "failed to display ask_user prompt" assert ask_user_resume["answers"] == [""] async def test_request_ask_user_missing_callback_is_reported_as_error(self) -> None: """ask_user interrupts without a UI callback should resume with error.""" agent = _SequencedAgent( streams_by_call=[ [ _ask_user_interrupt_chunk( { "type": "ask_user", "questions": [{"question": "Name?", "type": "text"}], "tool_call_id": "tool-1", } ) ], [], ] ) adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, request_ask_user=None, ) await execute_task_textual( user_input="hello", agent=agent, assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) resume_cmd = agent.stream_inputs[1] assert isinstance(resume_cmd, Command) resume_payload = cast("dict[str, dict[str, Any]]", resume_cmd.resume) ask_user_resume = resume_payload["interrupt-1"] assert ask_user_resume["status"] == "error" assert ask_user_resume["error"] == "ask_user not supported by this UI" assert ask_user_resume["answers"] == [""] async def test_spinner_reappears_after_ask_user_resume(self) -> None: """Spinner should re-show Thinking on each astream iteration. Regression for a gap where the model was working on the resume payload after an ask_user response but no spinner was visible. """ statuses: list[str | None] = [] async def record_spinner(status: str | None) -> None: await asyncio.sleep(0) statuses.append(status) async def request_ask_user( _questions: list[Question], ) -> asyncio.Future[AskUserWidgetResult] | None: await asyncio.sleep(0) return None agent = _SequencedAgent( streams_by_call=[ [ _ask_user_interrupt_chunk( { "type": "ask_user", "questions": [{"question": "Name?", "type": "text"}], "tool_call_id": "tool-1", } ) ], [], ] ) adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, request_ask_user=request_ask_user, set_spinner=record_spinner, ) await execute_task_textual( user_input="hello", agent=agent, assistant_id="assistant", session_state=SimpleNamespace(thread_id="thread-1", auto_approve=False), adapter=adapter, ) # Two astream iterations (interrupt, then resume) -> expect # Thinking set before each, and nothing above that count since # no tool calls stream in this test. assert len(agent.stream_inputs) == 2 thinking_count = sum(1 for s in statuses if s == "Thinking") assert thinking_count == 2, ( f"Expected Thinking spinner on each iteration; got {statuses}" ) async def test_invalid_ask_user_interrupt_payload_raises_validation_error( self, ) -> None: """Missing required ask_user keys should fail validation at ingestion.""" agent = _SequencedAgent( streams_by_call=[ [ _ask_user_interrupt_chunk( { "type": "ask_user", # Missing required keys: `questions` and `tool_call_id`. } ) ] ] ) adapter = TextualUIAdapter( mount_message=_mock_mount, update_status=_noop_status, request_approval=_mock_approval, ) with pytest.raises(ValidationError): await execute_task_textual( user_input="hello", agent=agent, assistant_id="assistant", session_state=SimpleNamespace( thread_id="thread-1", auto_approve=False, ), adapter=adapter, ) # --------------------------------------------------------------------------- # Helpers for dict-iteration safety tests # --------------------------------------------------------------------------- def _make_tool_widget(name: str = "tool", args: dict | None = None) -> MagicMock: """Create a MagicMock that mimics a ToolCallMessage widget.""" widget = MagicMock() widget._tool_name = name widget._args = args or {} return widget class _MutatingItemsDict(dict): # noqa: FURB189 # must subclass dict to override C-level iteration """Dict whose `.items()` deletes another key mid-iteration. This deterministically reproduces the `RuntimeError: dictionary changed size during iteration` that occurs when async tool-result callbacks mutate `_current_tool_messages` while the HITL approval loop is iterating over it. We intentionally subclass `dict` (not `UserDict`) because we need to override the C-level iteration that triggers the error. """ def items(self) -> Generator[tuple[str, Any], None, None]: # ty: ignore """Yield items while mutating the dict mid-iteration.""" it = iter(dict.items(self)) first = next(it) # Remove a *different* key while iteration is in progress. remaining = [k for k in self if k != first[0]] if remaining: del self[remaining[0]] yield first yield from it class _MutatingValuesDict(dict): # noqa: FURB189 # must subclass dict to override C-level iteration """Dict whose `.values()` deletes a key mid-iteration. We intentionally subclass `dict` (not `UserDict`) because we need to override the C-level iteration that triggers the error. """ def values(self) -> Generator[Any, None, None]: # ty: ignore """Yield values while mutating the dict mid-iteration.""" it = iter(dict.values(self)) first = next(it) # Remove the first key to trigger size-change error. first_key = next(iter(self)) del self[first_key] yield first yield from it class TestDictIterationSafety: """Regression tests for #956. Parallel tool calls can modify `adapter._current_tool_messages` while another coroutine iterates over it, raising `RuntimeError: dictionary changed size during iteration`. The fix wraps every iteration with `list()` so a snapshot is taken before the loop body runs. These tests prove the fix is necessary and sufficient. """ # -- Test A: bare iteration over a mutating dict raises ---- def test_items_iteration_fails_without_list(self) -> None: """Iterating .items() on a concurrently-mutated dict raises.""" d = _MutatingItemsDict( {f"id_{i}": _make_tool_widget(f"t{i}") for i in range(3)} ) with pytest.raises(RuntimeError, match="changed size"): for _ in d.items(): pass def test_values_iteration_fails_without_list(self) -> None: """Iterating .values() on a concurrently-mutated dict raises.""" d = _MutatingValuesDict( {f"id_{i}": _make_tool_widget(f"t{i}") for i in range(3)} ) with pytest.raises(RuntimeError, match="changed size"): for _ in d.values(): pass # -- Test B: list() snapshot protects iteration ---- def test_items_iteration_safe_with_list(self) -> None: """`list(d.items())` snapshots before mutation can occur.""" d: dict = {f"id_{i}": _make_tool_widget(f"t{i}") for i in range(5)} collected = [] for key, _val in list(d.items()): collected.append(key) d.pop(key, None) # mutate during loop body assert len(collected) == 5 assert len(d) == 0 def test_values_iteration_safe_with_list(self) -> None: """`list(d.values())` snapshots before mutation.""" d: dict = {f"id_{i}": _make_tool_widget(f"t{i}") for i in range(5)} collected = [] keys = list(d.keys()) for val in list(d.values()): collected.append(val) if keys: d.pop(keys.pop(0), None) assert len(collected) == 5 # -- Test C: _build_interrupted_ai_message uses list() ---- def test_build_interrupted_ai_message_safe(self) -> None: """_build_interrupted_ai_message correctly builds an AIMessage. Verifies the function reconstructs tool calls and content from the provided widget dict. The `list()` snapshot inside the production code protects against external async mutation at `await` boundaries, which cannot be deterministically simulated in a synchronous unit test. """ widgets = { f"id_{i}": _make_tool_widget(f"tool_{i}", {"k": i}) for i in range(4) } pending_text: dict[tuple, str] = {(): "hello"} result = _build_interrupted_ai_message(pending_text, widgets) assert result is not None assert result.content == "hello" assert len(result.tool_calls) == 4 names = {tc["name"] for tc in result.tool_calls} assert names == {"tool_0", "tool_1", "tool_2", "tool_3"} def test_build_interrupted_ai_message_empty(self) -> None: """Returns None when there is no text and no tool calls.""" result = _build_interrupted_ai_message({}, {}) assert result is None # --------------------------------------------------------------------------- # SessionStats tests # --------------------------------------------------------------------------- class TestSessionStats: """Tests for `SessionStats` recording and merging.""" def test_record_request_named_model(self) -> None: """record_request updates totals and per_model for a named model.""" stats = SessionStats() stats.record_request("gpt-4", 100, 50) assert stats.request_count == 1 assert stats.input_tokens == 100 assert stats.output_tokens == 50 assert ("", "gpt-4") in stats.per_model assert stats.per_model["", "gpt-4"].request_count == 1 assert stats.per_model["", "gpt-4"].input_tokens == 100 assert stats.per_model["", "gpt-4"].output_tokens == 50 def test_record_request_empty_model(self) -> None: """record_request with empty model skips per_model entry.""" stats = SessionStats() stats.record_request("", 200, 80) assert stats.request_count == 1 assert stats.input_tokens == 200 assert stats.output_tokens == 80 assert stats.per_model == {} def test_record_request_multiple_models(self) -> None: """Multiple models create separate per_model entries.""" stats = SessionStats() stats.record_request("gpt-4", 100, 50) stats.record_request("claude-opus-4-6", 200, 80) assert stats.request_count == 2 assert stats.input_tokens == 300 assert stats.output_tokens == 130 assert len(stats.per_model) == 2 assert stats.per_model["", "gpt-4"].request_count == 1 assert stats.per_model["", "claude-opus-4-6"].request_count == 1 def test_record_request_splits_same_model_by_provider(self) -> None: """Provider-specific model names should not collapse into one entry.""" stats = SessionStats() stats.record_request("gpt-4", 100, 50, provider="openai") stats.record_request("gpt-4", 200, 80, provider="azure") assert len(stats.per_model) == 2 assert stats.per_model["openai", "gpt-4"].request_count == 1 assert stats.per_model["azure", "gpt-4"].request_count == 1 def test_merge(self) -> None: """merge() folds another SessionStats into self.""" a = SessionStats( request_count=1, input_tokens=100, output_tokens=50, wall_time_seconds=1.0 ) a.per_model["", "gpt-4"] = ModelStats( request_count=1, input_tokens=100, output_tokens=50, model_name="gpt-4" ) b = SessionStats( request_count=2, input_tokens=300, output_tokens=120, wall_time_seconds=2.5 ) b.per_model["", "claude-opus-4-6"] = ModelStats( request_count=2, input_tokens=300, output_tokens=120, model_name="claude-opus-4-6", ) a.merge(b) assert a.request_count == 3 assert a.input_tokens == 400 assert a.output_tokens == 170 assert a.wall_time_seconds == pytest.approx(3.5) assert len(a.per_model) == 2 assert a.per_model["", "claude-opus-4-6"].request_count == 2 def test_merge_overlapping_models(self) -> None: """merge() combines per_model entries for the same model.""" a = SessionStats() a.record_request("gpt-4", 100, 50) b = SessionStats() b.record_request("gpt-4", 200, 80) a.merge(b) assert a.request_count == 2 assert a.input_tokens == 300 assert a.output_tokens == 130 assert a.per_model["", "gpt-4"].request_count == 2 assert a.per_model["", "gpt-4"].input_tokens == 300 assert a.per_model["", "gpt-4"].output_tokens == 130 def test_merge_splits_same_model_by_provider(self) -> None: """merge() preserves provider-specific entries for the same model.""" a = SessionStats() a.record_request("gpt-4", 100, 50, provider="openai") b = SessionStats() b.record_request("gpt-4", 200, 80, provider="azure") a.merge(b) assert len(a.per_model) == 2 assert a.per_model["openai", "gpt-4"].input_tokens == 100 assert a.per_model["azure", "gpt-4"].input_tokens == 200 # --------------------------------------------------------------------------- # format_token_count tests # --------------------------------------------------------------------------- class TestFormatTokenCount: """Tests for `format_token_count` shared formatter.""" def test_small_count(self) -> None: assert format_token_count(500) == "500" def test_thousands(self) -> None: assert format_token_count(12_500) == "12.5K" def test_millions(self) -> None: assert format_token_count(1_200_000) == "1.2M" def test_exact_thousand(self) -> None: assert format_token_count(1000) == "1.0K" def test_zero(self) -> None: assert format_token_count(0) == "0" # --------------------------------------------------------------------------- # print_usage_table tests # --------------------------------------------------------------------------- class TestPrintUsageTable: """Tests for `print_usage_table` output.""" def test_no_model_called_skips_unknown_row(self) -> None: """When no model was called, the table should not show 'unknown'.""" stats = SessionStats() buf = StringIO() console = Console(file=buf, force_terminal=True) print_usage_table(stats, wall_time=1.5, console=console) output = buf.getvalue() assert "unknown" not in output assert "Usage Stats" not in output assert "Agent active" in output def test_single_model_shows_name(self) -> None: """Single-model session should display the model name.""" stats = SessionStats() stats.record_request("gpt-4", 100, 50) buf = StringIO() console = Console(file=buf, force_terminal=True) print_usage_table(stats, wall_time=2.0, console=console) output = buf.getvalue() assert "gpt-4" in output assert "unknown" not in output def test_shows_provider_name(self) -> None: """The table should include the provider for each model.""" stats = SessionStats() stats.record_request("gpt-4", 100, 50, provider="openai") buf = StringIO() console = Console(file=buf, force_terminal=True) print_usage_table(stats, wall_time=2.0, console=console) output = buf.getvalue() assert "Provider" in output assert "openai" in output assert "gpt-4" in output def test_multi_model_shows_all_names_and_total(self) -> None: """Multi-model session should show each model and a Total row.""" stats = SessionStats() stats.record_request("gpt-4", 100, 50) stats.record_request("claude-opus-4-6", 200, 80) buf = StringIO() console = Console(file=buf, force_terminal=True) print_usage_table(stats, wall_time=2.0, console=console) output = buf.getvalue() assert "gpt-4" in output assert "claude-opus-4-6" in output assert "Total" in output assert "unknown" not in output def test_same_model_with_different_providers_shows_separate_rows(self) -> None: """Same-name models from different providers should render separately.""" stats = SessionStats() stats.record_request("gpt-4", 100, 50, provider="openai") stats.record_request("gpt-4", 200, 80, provider="azure") buf = StringIO() console = Console(file=buf, force_terminal=True) print_usage_table(stats, wall_time=2.0, console=console) output = buf.getvalue() assert "openai" in output assert "azure" in output assert "Total" in output # Two distinct rows, not a collapsed one: each provider's per-row token # counts must appear (100/50 and 200/80), alongside the 300/130 totals. assert "100" in output assert "50" in output assert "200" in output assert "80" in output def test_tokens_with_no_wall_time_omits_timing_line(self) -> None: """Token table should print but timing line should be absent.""" stats = SessionStats() stats.record_request("gpt-4", 100, 50) buf = StringIO() console = Console(file=buf, force_terminal=True) print_usage_table(stats, wall_time=0.0, console=console) output = buf.getvalue() assert "gpt-4" in output assert "Agent active" not in output def test_no_requests_no_time_prints_nothing(self) -> None: """Empty stats with negligible wall time should print nothing.""" stats = SessionStats() buf = StringIO() console = Console(file=buf, force_terminal=True) print_usage_table(stats, wall_time=0.01, console=console) output = buf.getvalue() assert output.strip() == ""