From 1f5619cc97dbb2c5e9d7b1f42c07e672a7224835 Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 21 Jul 2026 00:55:15 -0400 Subject: [PATCH 1/4] feat(tui): preserve media and nested agent events --- src/pythinker_code/ui/shell/replay.py | 3 +- .../ui/shell/visualize/_blocks.py | 6 + .../ui/shell/visualize/_live_view.py | 128 +++++++++++- tests/ui_and_conv/test_live_content_parts.py | 127 ++++++++++++ .../test_nested_subagent_events.py | 194 ++++++++++++++++++ tests/ui_and_conv/test_replay.py | 40 +++- .../ui_and_conv/test_subagent_live_stream.py | 23 ++- 7 files changed, 508 insertions(+), 13 deletions(-) create mode 100644 tests/ui_and_conv/test_live_content_parts.py create mode 100644 tests/ui_and_conv/test_nested_subagent_events.py diff --git a/src/pythinker_code/ui/shell/replay.py b/src/pythinker_code/ui/shell/replay.py index 8759ac1a..3e9407f2 100644 --- a/src/pythinker_code/ui/shell/replay.py +++ b/src/pythinker_code/ui/shell/replay.py @@ -162,9 +162,10 @@ def _is_clear_command_input(user_input: str | list[ContentPart]) -> bool: def _is_user_message(message: Message) -> bool: - # FIXME: should consider non-text tool call results which are sent as user messages if message.role != "user": return False + if message.tool_call_id is not None: + return False if message.extract_text().startswith("CHECKPOINT"): return False if is_notification_message(message): diff --git a/src/pythinker_code/ui/shell/visualize/_blocks.py b/src/pythinker_code/ui/shell/visualize/_blocks.py index 18cef4cb..12ee1db8 100644 --- a/src/pythinker_code/ui/shell/visualize/_blocks.py +++ b/src/pythinker_code/ui/shell/visualize/_blocks.py @@ -1252,6 +1252,7 @@ def __init__(self, tool_call: ToolCall): self._subagent_type: str | None = None self._ongoing_subagent_tool_calls: dict[str, ToolCall] = {} + self._finished_subagent_tool_call_ids: set[str] = set() self._last_subagent_tool_call: ToolCall | None = None self._n_finished_subagent_tool_calls = 0 self._finished_subagent_tool_counts: Counter[str] = Counter() @@ -1385,6 +1386,8 @@ def finish(self, result: ToolReturnValue): self._renderable = self._compose() def append_sub_tool_call(self, tool_call: ToolCall): + if tool_call.id in self._finished_subagent_tool_call_ids: + return self._ongoing_subagent_tool_calls[tool_call.id] = tool_call self._last_subagent_tool_call = tool_call self._renderable = self._compose() @@ -1401,10 +1404,13 @@ def append_sub_tool_call_part(self, tool_call_part: ToolCallPart): self._renderable = self._compose() def finish_sub_tool_call(self, tool_result: ToolResult): + if tool_result.tool_call_id in self._finished_subagent_tool_call_ids: + return self._last_subagent_tool_call = None sub_tool_call = self._ongoing_subagent_tool_calls.pop(tool_result.tool_call_id, None) if sub_tool_call is None: return + self._finished_subagent_tool_call_ids.add(tool_result.tool_call_id) self._subagent_output_parts.pop(tool_result.tool_call_id, None) self._subagent_output_had_stderr.pop(tool_result.tool_call_id, None) self._subagent_execution_started.discard(tool_result.tool_call_id) diff --git a/src/pythinker_code/ui/shell/visualize/_live_view.py b/src/pythinker_code/ui/shell/visualize/_live_view.py index c21a47f3..c0f2f174 100644 --- a/src/pythinker_code/ui/shell/visualize/_live_view.py +++ b/src/pythinker_code/ui/shell/visualize/_live_view.py @@ -42,7 +42,11 @@ from pythinker_code.ui.shell.console import console, current_console_width from pythinker_code.ui.shell.echo import render_user_echo from pythinker_code.ui.shell.focus_model import FocusTuiModel -from pythinker_code.ui.shell.glyphs import TRANSCRIPT_ACTIVE_MARKER, TRANSCRIPT_TOOL_GUTTER +from pythinker_code.ui.shell.glyphs import ( + TRANSCRIPT_ACTIVE_MARKER, + TRANSCRIPT_ASSISTANT_MARKER, + TRANSCRIPT_TOOL_GUTTER, +) from pythinker_code.ui.shell.keyboard import KeyboardListener, KeyEvent from pythinker_code.ui.shell.mcp_status import render_mcp_startup_text from pythinker_code.ui.shell.motion import ( @@ -91,10 +95,12 @@ from pythinker_code.utils.aioqueue import Queue, QueueShutDown from pythinker_code.utils.datetime import format_elapsed from pythinker_code.utils.logging import logger +from pythinker_code.utils.rich.columns import BulletColumns from pythinker_code.wire import WireUISide from pythinker_code.wire.types import ( ApprovalRequest, ApprovalResponse, + AudioURLPart, BtwBegin, BtwEnd, CompactionBegin, @@ -102,6 +108,7 @@ ContentPart, HookResolved, HookTriggered, + ImageURLPart, MCPLoadingBegin, MCPLoadingEnd, Notification, @@ -126,6 +133,7 @@ ToolResult, TurnBegin, TurnEnd, + VideoURLPart, WireMessage, ) @@ -143,6 +151,8 @@ # running long enough that a quick turn won't flash it. _WORKING_TIP_MIN_ELAPSED_S = 4.0 _MAX_PINNED_TODO_ROWS = 5 +_MAX_SUBAGENT_EVENT_DEPTH = 16 +_SEEN_UNKNOWN_CONTENT_PART_TYPES: set[str] = set() def _todo_activity_label(label: str) -> str: @@ -252,6 +262,7 @@ def __init__( self._current_content_block: _ContentBlock | None = None self._tool_call_blocks: dict[str, _ToolCallBlock] = {} + self._subagent_tool_call_ancestry: dict[str, tuple[_ToolCallBlock, int]] = {} self._last_tool_call_block: _ToolCallBlock | None = None self._held_tool_search_block: _ToolCallBlock | None = None self._completed_expandable_tool_blocks = deque[_ToolCallBlock](maxlen=20) @@ -1538,6 +1549,7 @@ def cleanup(self, is_interrupt: bool) -> None: self._mcp_loading_spinner = None self._btw_spinner = None self._hook_blocks.clear() + self._subagent_tool_call_ancestry.clear() self._current_step_retry = None if is_interrupt: @@ -1563,6 +1575,7 @@ def discard_retry_attempt(self, retry: StepRetry) -> None: """ self._current_content_block = None self._tool_call_blocks.clear() + self._subagent_tool_call_ancestry.clear() self._last_tool_call_block = None self._held_tool_search_block = None self._current_step_retry = retry @@ -1684,9 +1697,36 @@ def append_content(self, part: ContentPart) -> None: if text: self._current_content_block.append(text) self.refresh_soon() + case ImageURLPart(): + self._append_content_label("[image]") + case AudioURLPart(audio_url=audio): + suffix = f":{sanitize_ansi(audio.id)}" if audio.id else "" + self._append_content_label(f"[audio{suffix}]") + case VideoURLPart(): + self._append_content_label("[video]") case _: - # TODO: support more content part types - pass + part_type = part.type + if part_type not in _SEEN_UNKNOWN_CONTENT_PART_TYPES: + _SEEN_UNKNOWN_CONTENT_PART_TYPES.add(part_type) + logger.debug( + "Rendering unknown content part type in live view: {part_type}", + part_type=part_type, + ) + self._append_content_label(f"[{sanitize_ansi(part_type)}]", unknown=True) + + def _append_content_label(self, label: str, *, unknown: bool = False) -> None: + """Render a payload-free media or future-content placeholder.""" + self._current_step_retry = None + self.flush_content(FlushReason.TOOL_START) + muted = tui_rich_style("muted") + marker_style = muted if unknown else tui_rich_style("success") + self._emit_final_scrollback( + BulletColumns( + Text(label, style=muted), + bullet=Text(TRANSCRIPT_ASSISTANT_MARKER, style=marker_style), + ) + ) + self.refresh_soon() def append_tool_call(self, tool_call: ToolCall) -> None: self._current_step_retry = None @@ -1877,16 +1917,65 @@ def show_next_question_request(self) -> None: self._on_question_panel_state_changed() def handle_subagent_event(self, event: SubagentEvent) -> None: - if event.parent_tool_call_id is None: + self._dispatch_subagent_event(event, recursive_depth=1) + + def _dispatch_subagent_event(self, event: SubagentEvent, *, recursive_depth: int) -> None: + if recursive_depth > _MAX_SUBAGENT_EVENT_DEPTH: + self._render_subagent_event_fallback( + event, + reason="depth limit reached", + depth=recursive_depth, + ) return - block = self._tool_call_blocks.get(event.parent_tool_call_id) - if block is None: + + parent_tool_call_id = event.parent_tool_call_id + if parent_tool_call_id is None: + self._render_subagent_event_fallback( + event, + reason="missing ancestry", + depth=recursive_depth, + ) return + + block = self._tool_call_blocks.get(parent_tool_call_id) + parent_depth = 0 + if block is None: + ancestry = self._subagent_tool_call_ancestry.get(parent_tool_call_id) + if ancestry is None: + self._render_subagent_event_fallback( + event, + reason="missing ancestry", + depth=recursive_depth, + ) + return + block, parent_depth = ancestry + if event.agent_id is not None and event.subagent_type is not None: block.set_subagent_metadata(event.agent_id, event.subagent_type) match event.event: + case SubagentEvent() as nested_event: + if recursive_depth >= _MAX_SUBAGENT_EVENT_DEPTH: + self._render_subagent_event_fallback( + nested_event, + reason="depth limit reached", + depth=recursive_depth + 1, + ) + return + self._dispatch_subagent_event( + nested_event, + recursive_depth=recursive_depth + 1, + ) case ToolCall() as tool_call: + tool_depth = parent_depth + 1 + if tool_depth > _MAX_SUBAGENT_EVENT_DEPTH: + self._render_subagent_event_fallback( + event, + reason="depth limit reached", + depth=tool_depth, + ) + return + self._subagent_tool_call_ancestry[tool_call.id] = (block, tool_depth) block.append_sub_tool_call(tool_call) self.refresh_soon() case ToolCallPart() as tool_call_part: @@ -1906,6 +1995,29 @@ def handle_subagent_event(self, event: SubagentEvent) -> None: ) self.refresh_soon() case _: - # ignore other events for now - # TODO: may need to handle multi-level nested subagents pass + + def _render_subagent_event_fallback( + self, + event: SubagentEvent, + *, + reason: str, + depth: int, + ) -> None: + event_type = type(event.event).__name__ + logger.debug( + "Unable to render nested subagent event: reason={reason} parent={parent} " + "depth={depth} event_type={event_type}", + reason=reason, + parent=event.parent_tool_call_id, + depth=depth, + event_type=event_type, + ) + muted = tui_rich_style("muted") + self._emit_action_block( + BulletColumns( + Text(f"Nested subagent activity unavailable · {reason}", style=muted), + bullet=Text(TRANSCRIPT_ASSISTANT_MARKER, style=muted), + ) + ) + self.refresh_soon() diff --git a/tests/ui_and_conv/test_live_content_parts.py b/tests/ui_and_conv/test_live_content_parts.py new file mode 100644 index 00000000..1fa409a0 --- /dev/null +++ b/tests/ui_and_conv/test_live_content_parts.py @@ -0,0 +1,127 @@ +"""Lossless live-view rendering for media and future content parts.""" + +from __future__ import annotations + +from collections.abc import Iterable + +import pytest +from pythinker_core.message import ContentPart +from rich.console import Console, Group, RenderableType + +from pythinker_code.ui.shell.visualize import _live_view as live_view_module +from pythinker_code.ui.shell.visualize import _LiveView +from pythinker_code.wire.types import ( + AudioURLPart, + ImageURLPart, + StatusUpdate, + TextPart, + VideoURLPart, +) + + +class FuturePayloadPart(ContentPart): + type: str = "future_payload" + payload: str + + +def _render(renderables: Iterable[RenderableType], *, color: bool = False) -> str: + console = Console( + width=100, + record=True, + highlight=False, + color_system="standard" if color else None, + ) + console.print(Group(*renderables)) + return console.export_text() + + +def _capture_scrollback( + monkeypatch: pytest.MonkeyPatch, +) -> list[RenderableType]: + emitted: list[RenderableType] = [] + monkeypatch.setattr( + live_view_module, + "emit_scrollback_block", + lambda _console, renderable: emitted.append(renderable), + ) + return emitted + + +def test_text_media_text_flushes_at_stable_boundaries( + monkeypatch: pytest.MonkeyPatch, +) -> None: + emitted = _capture_scrollback(monkeypatch) + view = _LiveView(StatusUpdate(context_tokens=1000)) + + view.append_content(TextPart(text="before")) + view.append_content( + ImageURLPart(image_url=ImageURLPart.ImageURL(url="data:image/png;base64,SECRET_IMAGE")) + ) + view.append_content(TextPart(text="after")) + view.flush_content() + + assert len(emitted) == 3 + assert "before" in _render([emitted[0]]) + assert "[image]" in _render([emitted[1]]) + assert "after" in _render([emitted[2]]) + output = _render(emitted) + assert output.index("before") < output.index("[image]") < output.index("after") + assert "data:image" not in output + assert "SECRET_IMAGE" not in output + + +@pytest.mark.parametrize("ascii_mode", [False, True]) +def test_all_media_labels_survive_no_color_and_ascii_modes( + monkeypatch: pytest.MonkeyPatch, + ascii_mode: bool, +) -> None: + emitted = _capture_scrollback(monkeypatch) + if ascii_mode: + monkeypatch.setenv("PYTHINKER_ASCII_UI", "1") + view = _LiveView(StatusUpdate(context_tokens=1000)) + + view.append_content( + ImageURLPart(image_url=ImageURLPart.ImageURL(url="https://media.invalid/image-secret")) + ) + view.append_content( + AudioURLPart( + audio_url=AudioURLPart.AudioURL( + url="data:audio/aac;base64,SECRET_AUDIO", + id="clip-7", + ) + ) + ) + view.append_content( + VideoURLPart(video_url=VideoURLPart.VideoURL(url="data:video/mp4;base64,SECRET_VIDEO")) + ) + + output = _render(emitted) + assert "[image]" in output + assert "[audio:clip-7]" in output + assert "[video]" in output + assert "media.invalid" not in output + assert "SECRET_AUDIO" not in output + assert "SECRET_VIDEO" not in output + + +def test_unknown_content_uses_muted_label_and_logs_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + emitted = _capture_scrollback(monkeypatch) + debug_calls: list[tuple[tuple[object, ...], dict[str, object]]] = [] + + def capture_debug(*args: object, **kwargs: object) -> None: + debug_calls.append((args, kwargs)) + + monkeypatch.setattr(live_view_module.logger, "debug", capture_debug) + view = _LiveView(StatusUpdate(context_tokens=1000)) + + view.append_content(FuturePayloadPart(payload="SECRET_FUTURE_PAYLOAD")) + view.append_content(FuturePayloadPart(payload="ANOTHER_SECRET_PAYLOAD")) + + output = _render(emitted, color=True) + assert output.count("[future_payload]") == 2 + assert "SECRET_FUTURE_PAYLOAD" not in output + assert "ANOTHER_SECRET_PAYLOAD" not in output + assert len(debug_calls) == 1 + assert debug_calls[0][1]["part_type"] == "future_payload" diff --git a/tests/ui_and_conv/test_nested_subagent_events.py b/tests/ui_and_conv/test_nested_subagent_events.py new file mode 100644 index 00000000..9fd39064 --- /dev/null +++ b/tests/ui_and_conv/test_nested_subagent_events.py @@ -0,0 +1,194 @@ +"""Nested SubagentEvent ancestry and bounded dispatch tests.""" + +from __future__ import annotations + +from collections.abc import Iterable + +import pytest +from pythinker_core.message import ToolCall +from pythinker_core.tooling import ToolOk +from rich.console import Console, Group, RenderableType + +from pythinker_code.ui.shell.visualize import _live_view as live_view_module +from pythinker_code.ui.shell.visualize import _LiveView +from pythinker_code.wire.types import ( + Event, + StatusUpdate, + SubagentEvent, + ToolCallPart, + ToolExecutionStarted, + ToolOutputPart, + ToolResult, + TurnBegin, +) +from pythinker_code.wire.types import ToolCall as WireToolCall + + +def _render(renderables: Iterable[RenderableType]) -> str: + console = Console(width=100, record=True, highlight=False, color_system=None) + console.print(Group(*renderables)) + return console.export_text() + + +def _tool_call(call_id: str, name: str, arguments: str) -> ToolCall: + return ToolCall( + id=call_id, + function=ToolCall.FunctionBody(name=name, arguments=arguments), + ) + + +def _root_agent_call() -> WireToolCall: + return WireToolCall( + id="root-agent", + function=WireToolCall.FunctionBody( + name="Agent", + arguments='{"description":"nested task","subagent_type":"coder","prompt":"work"}', + ), + ) + + +def _nested_event(ancestor_ids: list[str], event: Event) -> SubagentEvent: + nested: Event = event + for parent_id in reversed(ancestor_ids): + nested = SubagentEvent( + parent_tool_call_id=parent_id, + agent_id=f"agent-{parent_id}", + subagent_type="coder", + event=nested, + ) + return SubagentEvent( + parent_tool_call_id="root-agent", + agent_id="agent-root", + subagent_type="coder", + event=nested, + ) + + +def _view() -> _LiveView: + view = _LiveView(StatusUpdate(context_tokens=1000)) + view.dispatch_wire_message(TurnBegin(user_input="work")) + view.dispatch_wire_message(_root_agent_call()) + return view + + +@pytest.mark.parametrize("nesting_depth", [2, 3]) +def test_nested_tool_lifecycle_rolls_up_under_root_and_first_result_wins( + nesting_depth: int, +) -> None: + view = _view() + ancestor_ids: list[str] = [] + + for depth in range(1, nesting_depth + 1): + call_id = f"nested-{depth}" + is_leaf = depth == nesting_depth + call = _tool_call( + call_id, + "Read" if is_leaf else "Agent", + '{"file_path":"src/' if is_leaf else '{"description":"deeper"}', + ) + view.dispatch_wire_message(_nested_event(ancestor_ids, call)) + ancestor_ids.append(call_id) + + leaf_id = ancestor_ids[-1] + leaf_parent_ids = ancestor_ids[:-1] + view.dispatch_wire_message( + _nested_event(leaf_parent_ids, ToolCallPart(arguments_part='module.py"}')) + ) + view.dispatch_wire_message( + _nested_event(leaf_parent_ids, ToolExecutionStarted(tool_call_id=leaf_id)) + ) + view.dispatch_wire_message( + _nested_event( + leaf_parent_ids, + ToolOutputPart(tool_call_id=leaf_id, text="STREAMED_NESTED_OUTPUT\n"), + ) + ) + + root_block = view._tool_call_blocks["root-agent"] + for expected_depth, call_id in enumerate(ancestor_ids, start=1): + indexed_block, indexed_depth = view._subagent_tool_call_ancestry[call_id] + assert indexed_block is root_block + assert indexed_depth == expected_depth + assert leaf_id in root_block._subagent_execution_started + assert "STREAMED_NESTED_OUTPUT" in _render([view.compose()]) + assert "src/module.py" in _render([view.compose()]) + + view.dispatch_wire_message( + _nested_event( + leaf_parent_ids, + ToolResult(tool_call_id=leaf_id, return_value=ToolOk(output="FIRST_RESULT")), + ) + ) + view.dispatch_wire_message( + _nested_event(leaf_parent_ids, ToolExecutionStarted(tool_call_id=leaf_id)) + ) + view.dispatch_wire_message( + _nested_event( + leaf_parent_ids, + ToolOutputPart(tool_call_id=leaf_id, text="LATE_OUTPUT"), + ) + ) + view.dispatch_wire_message( + _nested_event( + leaf_parent_ids, + ToolResult(tool_call_id=leaf_id, return_value=ToolOk(output="LATE_RESULT")), + ) + ) + + finished = [ + item for item in root_block._finished_subagent_tool_calls if item.call.id == leaf_id + ] + assert len(finished) == 1 + assert finished[0].result.output == "FIRST_RESULT" + assert root_block._n_finished_subagent_tool_calls == 1 + assert leaf_id not in root_block._subagent_execution_started + assert "LATE_OUTPUT" not in root_block._subagent_output_parts + + +def test_missing_ancestry_renders_one_muted_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + emitted: list[RenderableType] = [] + monkeypatch.setattr( + live_view_module, + "emit_scrollback_block", + lambda _console, renderable: emitted.append(renderable), + ) + view = _view() + + view.dispatch_wire_message( + SubagentEvent( + parent_tool_call_id="missing-parent", + event=ToolOutputPart(tool_call_id="missing-tool", text="SECRET_PAYLOAD"), + ) + ) + + assert len(emitted) == 1 + output = _render(emitted) + assert "Nested subagent activity unavailable" in output + assert "missing ancestry" in output + assert "SECRET_PAYLOAD" not in output + + +def test_recursive_depth_overflow_renders_one_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + emitted: list[RenderableType] = [] + monkeypatch.setattr( + live_view_module, + "emit_scrollback_block", + lambda _console, renderable: emitted.append(renderable), + ) + view = _view() + nested: Event = ToolOutputPart(tool_call_id="never-dispatched", text="TOO_DEEP") + for _ in range(17): + nested = SubagentEvent(parent_tool_call_id="root-agent", event=nested) + assert isinstance(nested, SubagentEvent) + + view.dispatch_wire_message(nested) + + assert len(emitted) == 1 + output = _render(emitted) + assert "Nested subagent activity unavailable" in output + assert "depth limit reached" in output + assert "TOO_DEEP" not in output diff --git a/tests/ui_and_conv/test_replay.py b/tests/ui_and_conv/test_replay.py index 9674ceeb..94d324fd 100644 --- a/tests/ui_and_conv/test_replay.py +++ b/tests/ui_and_conv/test_replay.py @@ -14,7 +14,7 @@ ) from pythinker_code.utils.aioqueue import QueueShutDown from pythinker_code.wire.file import WireFile -from pythinker_code.wire.types import SteerInput, StepBegin, TextPart, TurnBegin +from pythinker_code.wire.types import ImageURLPart, SteerInput, StepBegin, TextPart, TurnBegin @pytest.fixture(autouse=True) @@ -131,6 +131,44 @@ def test_build_replay_turns_from_history_keeps_plain_steer_as_user_turn() -> Non assert turns[1].user_message.extract_text(" ") == "A steer follow-up" +def test_real_user_media_message_starts_replay_turn() -> None: + media = ImageURLPart( + image_url=ImageURLPart.ImageURL(url="data:image/png;base64,REAL_USER_MEDIA") + ) + history = [ + Message(role="user", content=[media]), + Message(role="assistant", content=[TextPart(text="I can inspect that image.")]), + ] + + turns = _build_replay_turns_from_history(history) + + assert len(turns) == 1 + assert turns[0].user_message.content == [media] + assert turns[0].user_message.tool_call_id is None + + +def test_user_role_media_tool_result_does_not_start_replay_turn() -> None: + remapped_tool_media = ImageURLPart( + image_url=ImageURLPart.ImageURL(url="data:image/png;base64,TOOL_RESULT_MEDIA") + ) + history = [ + Message(role="user", content=[TextPart(text="Inspect the generated image")]), + Message(role="assistant", content=[TextPart(text="Starting inspection")]), + Message( + role="user", + content=[remapped_tool_media], + tool_call_id="provider-remapped-tool-result", + ), + Message(role="assistant", content=[TextPart(text="Inspection complete")]), + ] + + turns = _build_replay_turns_from_history(history) + + assert len(turns) == 1 + assert turns[0].user_message.extract_text(" ") == "Inspect the generated image" + assert turns[0].n_steps == 2 + + @pytest.mark.asyncio async def test_build_replay_turns_from_wire_keeps_steer_as_user_turn(tmp_path: Path) -> None: wire_file = WireFile(tmp_path / "wire.jsonl") diff --git a/tests/ui_and_conv/test_subagent_live_stream.py b/tests/ui_and_conv/test_subagent_live_stream.py index c7d01540..d626d6b7 100644 --- a/tests/ui_and_conv/test_subagent_live_stream.py +++ b/tests/ui_and_conv/test_subagent_live_stream.py @@ -2,9 +2,10 @@ from __future__ import annotations +import pytest from pythinker_core.message import ToolCall from pythinker_core.tooling import ToolOk -from rich.console import Console +from rich.console import Console, RenderableType from pythinker_code.ui.shell.visualize import _LiveView from pythinker_code.wire.types import ( @@ -127,9 +128,19 @@ def test_subagent_tool_call_and_args_request_live_refresh(): assert "src/app.py" in _render(view) -def test_output_part_for_unknown_parent_is_silently_ignored(): +def test_output_part_for_unknown_parent_renders_fallback_without_payload( + monkeypatch: pytest.MonkeyPatch, +): view = _LiveView(StatusUpdate(context_tokens=1000)) view.dispatch_wire_message(TurnBegin(user_input="scan")) + emitted: list[RenderableType] = [] + from pythinker_code.ui.shell.visualize import _live_view as live_view_module + + monkeypatch.setattr( + live_view_module, + "emit_scrollback_block", + lambda _console, renderable: emitted.append(renderable), + ) # No agent tool call dispatched — parent_tool_call_id won't resolve view.dispatch_wire_message( SubagentEvent( @@ -139,9 +150,15 @@ def test_output_part_for_unknown_parent_is_silently_ignored(): event=ToolOutputPart(tool_call_id="sub-1", text="should be ignored\n"), ) ) - # Must not raise; compose must still work + # Must not raise or expose the nested payload; a muted fallback is emitted. output = _render(view) assert "should be ignored" not in output + assert len(emitted) == 1 + console = Console(width=100, record=True, highlight=False, color_system=None) + console.print(emitted[0]) + fallback = console.export_text() + assert "Nested subagent activity unavailable" in fallback + assert "should be ignored" not in fallback def test_output_cleared_after_sub_tool_call_finishes(): From f1eb37f52883fa2c9689fc817fc0915ca6c69d25 Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 21 Jul 2026 01:22:22 -0400 Subject: [PATCH 2/4] refactor(shell): harden command and idle event execution --- src/pythinker_code/ui/shell/__init__.py | 143 +++++++------- src/pythinker_code/ui/shell/command_runner.py | 134 +++++++++++++ .../ui/shell/visualize/_blocks.py | 24 ++- .../test_background_completion_watcher.py | 82 ++++++-- tests/ui_and_conv/test_shell_bang_spawn.py | 49 +++-- .../ui_and_conv/test_shell_command_runner.py | 182 ++++++++++++++++++ tests/ui_and_conv/test_tool_call_block.py | 28 +++ 7 files changed, 530 insertions(+), 112 deletions(-) create mode 100644 src/pythinker_code/ui/shell/command_runner.py create mode 100644 tests/ui_and_conv/test_shell_command_runner.py diff --git a/src/pythinker_code/ui/shell/__init__.py b/src/pythinker_code/ui/shell/__init__.py index eec1ecd7..3939b3f2 100644 --- a/src/pythinker_code/ui/shell/__init__.py +++ b/src/pythinker_code/ui/shell/__init__.py @@ -4,7 +4,6 @@ import asyncio import contextlib import json -import os import re import shlex import textwrap @@ -12,7 +11,7 @@ from collections import deque from collections.abc import Awaitable, Callable, Coroutine from dataclasses import dataclass -from enum import Enum +from enum import Enum, StrEnum from typing import TYPE_CHECKING, Any, Protocol, cast if TYPE_CHECKING: @@ -25,7 +24,6 @@ APITimeoutError, ChatProviderError, ) -from pythinker_host.windows import windows_console_detach_flags from rich import box from rich.align import Align from rich.cells import cell_len @@ -49,6 +47,7 @@ run_soul, ) from pythinker_code.soul.pythinkersoul import FLOW_COMMAND_PREFIX, PythinkerSoul +from pythinker_code.ui.shell.command_runner import ShellCommandRunner from pythinker_code.ui.shell.components.render_utils import ( cell_width, render_message_response, @@ -106,7 +105,6 @@ from pythinker_code.utils.logging import logger from pythinker_code.utils.signals import install_sigint_handler from pythinker_code.utils.slashcmd import SlashCommand, SlashCommandCall, parse_slash_command_call -from pythinker_code.utils.subprocess_env import get_clean_env from pythinker_code.utils.term import ensure_new_line, ensure_tty_sane from pythinker_code.wire.types import ( ApprovalRequest, @@ -117,9 +115,20 @@ ) +class PromptEventKind(StrEnum): + INPUT = "input" + INPUT_ACTIVITY = "input_activity" + BACKGROUND_GRACE_EXPIRED = "background_grace_expired" + BACKGROUND_NOOP = "bg_noop" + INTERRUPT = "interrupt" + EOF = "eof" + CWD_LOST = "cwd_lost" + ERROR = "error" + + @dataclass(slots=True) class _PromptEvent: - kind: str + kind: PromptEventKind user_input: UserInput | None = None @@ -235,15 +244,19 @@ async def wait_for_next(self, idle_events: asyncio.Queue[_PromptEvent]) -> _Prom assert self._event is not None bg_wait_task = asyncio.create_task(self._event.wait()) - done, _ = await asyncio.wait( - [idle_task, bg_wait_task], - return_when=asyncio.FIRST_COMPLETED, - ) - for t in (idle_task, bg_wait_task): - if t not in done: - t.cancel() + done: set[asyncio.Task[Any]] = set() + try: + done, _ = await asyncio.wait( + [idle_task, bg_wait_task], + return_when=asyncio.FIRST_COMPLETED, + ) + finally: + for task in (idle_task, bg_wait_task): + if task.done(): + continue + task.cancel() with contextlib.suppress(asyncio.CancelledError): - await t + await task if idle_task in done: if bg_wait_task in done: @@ -255,8 +268,8 @@ async def wait_for_next(self, idle_events: asyncio.Queue[_PromptEvent]) -> _Prom if self._has_pending_llm_notifications(): if self._can_auto_trigger_pending(): return None - return _PromptEvent(kind="bg_noop") - return _PromptEvent(kind="bg_noop") + return _PromptEvent(kind=PromptEventKind.BACKGROUND_NOOP) + return _PromptEvent(kind=PromptEventKind.BACKGROUND_NOOP) def _has_pending_llm_notifications(self) -> bool: if self._notifications is None: @@ -735,7 +748,7 @@ async def _route_prompt_events( self._running_interrupt_handler() continue resume_prompt.clear() - await idle_events.put(_PromptEvent(kind="interrupt")) + await idle_events.put(_PromptEvent(kind=PromptEventKind.INTERRUPT)) continue except EOFError: logger.debug("Prompt router got EOF") @@ -748,17 +761,17 @@ async def _route_prompt_events( self._running_interrupt_handler() return resume_prompt.clear() - await idle_events.put(_PromptEvent(kind="eof")) + await idle_events.put(_PromptEvent(kind=PromptEventKind.EOF)) return except CwdLostError: logger.error("Working directory no longer exists") resume_prompt.clear() - await idle_events.put(_PromptEvent(kind="cwd_lost")) + await idle_events.put(_PromptEvent(kind=PromptEventKind.CWD_LOST)) return except Exception: logger.exception("Prompt router crashed") resume_prompt.clear() - await idle_events.put(_PromptEvent(kind="error")) + await idle_events.put(_PromptEvent(kind=PromptEventKind.ERROR)) return if prompt_session.last_submission_was_running: # noqa: SIM102 @@ -769,7 +782,7 @@ async def _route_prompt_events( # Handler already unbound — fall through to idle path. resume_prompt.clear() - await idle_events.put(_PromptEvent(kind="input", user_input=user_input)) + await idle_events.put(_PromptEvent(kind=PromptEventKind.INPUT, user_input=user_input)) def _register_task_label_resolver(self) -> None: """Let TaskOutput/TaskStop headers show a task's friendly description @@ -1042,6 +1055,13 @@ def _can_auto_trigger_pending() -> bool: else: result = await bg_watcher.wait_for_next(idle_events) + if ( + result is not None + and result.kind is PromptEventKind.BACKGROUND_GRACE_EXPIRED + ): + logger.debug("Background auto-trigger input grace elapsed") + result = None + if result is None: if self._should_defer_background_auto_trigger(prompt_session): deferred_bg_trigger = True @@ -1082,28 +1102,29 @@ def _can_auto_trigger_pending() -> bool: event = result - if event.kind == "input_activity": + if event.kind is PromptEventKind.INPUT_ACTIVITY: + logger.debug("Deferring background auto-trigger for local input activity") continue - if event.kind == "bg_noop": + if event.kind is PromptEventKind.BACKGROUND_NOOP: continue - if event.kind == "interrupt": + if event.kind is PromptEventKind.INTERRUPT: _t = _get_tui_tokens() console.print(f"[{_t.muted}]Tip: press Ctrl-D or send 'exit' to quit[/]") resume_prompt.set() continue - if event.kind == "eof": + if event.kind is PromptEventKind.EOF: console.print("Bye!") break - if event.kind == "cwd_lost": + if event.kind is PromptEventKind.CWD_LOST: self._print_cwd_lost_crash() shell_ok = False break - if event.kind == "error": + if event.kind is PromptEventKind.ERROR: shell_ok = False break @@ -1251,67 +1272,31 @@ async def _run_shell_command(self, command: str) -> None: track("input_bash") - proc: asyncio.subprocess.Process | None = None - max_output_bytes = 1_000_000 + runner_task = asyncio.create_task(ShellCommandRunner().run(command)) + interrupted = False - async def _read_stream_limited(stream: asyncio.StreamReader | None, limit: int) -> bytes: - if stream is None: - return b"" - chunks: list[bytes] = [] - total = 0 - truncated = False - while True: - chunk = await stream.read(65536) - if not chunk: - break - remaining = limit - total - if remaining > 0: - chunks.append(chunk[:remaining]) - total += min(len(chunk), remaining) - if len(chunk) > remaining: - truncated = True - if truncated: - chunks.append(b"\n... output truncated ...\n") - return b"".join(chunks) - - def _handler(): + def _handler() -> None: + nonlocal interrupted logger.debug("SIGINT received.") - if proc: - proc.terminate() + interrupted = True + runner_task.cancel() loop = asyncio.get_running_loop() remove_sigint = install_sigint_handler(loop, _handler) try: - # TODO: For the sake of simplicity, we now use `create_subprocess_shell`. - # Later we should consider making this behave like a real shell. - spawn_kwargs: dict[str, Any] = {} - if os.name == "nt": - # CREATE_NO_WINDOW: don't share the interactive console — a child - # touching it via the Win32 console API bypasses the pipes and - # can blank the TUI until terminal restart. - spawn_kwargs["creationflags"] = windows_console_detach_flags( - new_process_group=False - ) - proc = await asyncio.create_subprocess_shell( - command, - env=get_clean_env(), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - **spawn_kwargs, - ) - stdout_task = asyncio.create_task(_read_stream_limited(proc.stdout, max_output_bytes)) - stderr_task = asyncio.create_task(_read_stream_limited(proc.stderr, max_output_bytes)) - await proc.wait() - stdout_bytes, stderr_bytes = await asyncio.gather(stdout_task, stderr_task) - stdout = stdout_bytes.decode("utf-8", errors="replace") if stdout_bytes else "" - stderr = stderr_bytes.decode("utf-8", errors="replace") if stderr_bytes else "" + # Commands intentionally run in the detected configured shell. Each + # execution remains isolated, so state such as `cd` is not persistent. + result = await runner_task output = _format_local_shell_output( - stdout=stdout, - stderr=stderr, - returncode=proc.returncode, + stdout=result.stdout, + stderr=result.stderr, + returncode=result.returncode, ) if output is not None: console.print(render_message_response(output)) + except asyncio.CancelledError: + if not interrupted: + raise except Exception as e: logger.exception("Failed to run shell command:") console.print( @@ -1794,7 +1779,9 @@ async def _wait_for_input_or_activity( if idle_task in done: return idle_task.result() - return _PromptEvent(kind="input_activity") + if activity_task in done: + return _PromptEvent(kind=PromptEventKind.INPUT_ACTIVITY) + return _PromptEvent(kind=PromptEventKind.BACKGROUND_GRACE_EXPIRED) async def _watch_root_wire_hub(self) -> None: if not isinstance(self.soul, PythinkerSoul): diff --git a/src/pythinker_code/ui/shell/command_runner.py b/src/pythinker_code/ui/shell/command_runner.py new file mode 100644 index 00000000..832ebaaf --- /dev/null +++ b/src/pythinker_code/ui/shell/command_runner.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import asyncio +from dataclasses import dataclass + +from pythinker_host import AsyncReadable, Host, HostProcess, get_current_host + +from pythinker_code.utils.environment import Environment +from pythinker_code.utils.logging import logger +from pythinker_code.utils.subprocess_env import get_clean_env + +_MAX_STREAM_BYTES = 1024 * 1024 +_TRUNCATION_NOTICE = b"\n... output truncated ...\n" + + +@dataclass(frozen=True, slots=True) +class ShellCommandResult: + stdout: str + stderr: str + returncode: int + truncated_stdout: bool + truncated_stderr: bool + + +class ShellCommandRunner: + """Execute isolated foreground commands through the detected host shell.""" + + def __init__( + self, + host: Host | None = None, + *, + environment: Environment | None = None, + output_limit_bytes: int = _MAX_STREAM_BYTES, + ) -> None: + if output_limit_bytes < 0: + raise ValueError("output_limit_bytes must be non-negative") + self._host = host or get_current_host() + self._environment = environment + self._output_limit_bytes = output_limit_bytes + + @staticmethod + def build_argv(command: str, *, shell_name: str, shell_path: str) -> tuple[str, ...]: + """Build the explicit argv used for POSIX shells and PowerShell.""" + option = "-command" if shell_name == "Windows PowerShell" else "-c" + return (shell_path, option, command) + + async def run(self, command: str) -> ShellCommandResult: + environment = self._environment or await Environment.detect() + argv = self.build_argv( + command, + shell_name=environment.shell_name, + shell_path=str(environment.shell_path), + ) + process = await self._host.exec(*argv, env=get_clean_env()) + process.stdin.close() + + stdout_task = asyncio.create_task( + self._read_stream_limited(process.stdout, self._output_limit_bytes) + ) + stderr_task = asyncio.create_task( + self._read_stream_limited(process.stderr, self._output_limit_bytes) + ) + wait_task = asyncio.create_task(process.wait()) + tasks = (stdout_task, stderr_task, wait_task) + try: + await asyncio.gather(*tasks) + except asyncio.CancelledError: + await self._terminate_and_reap(process, tasks) + raise + except Exception: + await self._terminate_and_reap(process, tasks) + raise + + stdout_bytes, truncated_stdout = stdout_task.result() + stderr_bytes, truncated_stderr = stderr_task.result() + returncode = wait_task.result() + + encoding = "utf-8" + return ShellCommandResult( + stdout=stdout_bytes.decode(encoding=encoding, errors="replace"), + stderr=stderr_bytes.decode(encoding=encoding, errors="replace"), + returncode=returncode, + truncated_stdout=truncated_stdout, + truncated_stderr=truncated_stderr, + ) + + @staticmethod + async def _read_stream_limited(stream: AsyncReadable, limit: int) -> tuple[bytes, bool]: + chunks: list[bytes] = [] + stored = 0 + truncated = False + while True: + chunk = await stream.read(65536) + if not chunk: + break + remaining = limit - stored + if remaining > 0: + kept = chunk[:remaining] + chunks.append(kept) + stored += len(kept) + if len(chunk) > max(remaining, 0): + truncated = True + if truncated: + chunks.append(_TRUNCATION_NOTICE) + return b"".join(chunks), truncated + + @staticmethod + async def _terminate_and_reap( + process: HostProcess, + tasks: tuple[ + asyncio.Task[tuple[bytes, bool]], + asyncio.Task[tuple[bytes, bool]], + asyncio.Task[int], + ], + ) -> None: + try: + await process.kill() + except Exception: + logger.exception("Failed to terminate cancelled shell command") + + try: + await process.wait() + except Exception: + logger.exception("Failed to reap cancelled shell command") + + for task in tasks: + if not task.done(): + task.cancel() + outcomes = await asyncio.gather(*tasks, return_exceptions=True) + for outcome in outcomes: + if isinstance(outcome, BaseException) and not isinstance( + outcome, asyncio.CancelledError + ): + logger.debug("Shell command cleanup task failed: {error}", error=outcome) diff --git a/src/pythinker_code/ui/shell/visualize/_blocks.py b/src/pythinker_code/ui/shell/visualize/_blocks.py index 12ee1db8..4bf46657 100644 --- a/src/pythinker_code/ui/shell/visualize/_blocks.py +++ b/src/pythinker_code/ui/shell/visualize/_blocks.py @@ -89,6 +89,7 @@ _ELLIPSIS = "..." _THINKING_PREVIEW_LINES = 6 _COMPOSING_PREVIEW_LINES = 12 +_ARGUMENT_SCAN_GROWTH_BYTES = 1024 # Smooth-streaming reveal pacing (composing text only). Deltas arrive bursty; # instead of revealing each whole chunk at once, a paced reveal cursor advances @@ -1247,6 +1248,7 @@ def __init__(self, tool_call: ToolCall): self._argument = self._extract_worklog_argument( tool_call.function.arguments, self._tool_name ) + self._argument_scan_bytes = 0 self._result: ToolReturnValue | None = None self._subagent_id: str | None = None self._subagent_type: str | None = None @@ -1347,22 +1349,34 @@ def append_args_part(self, args_part: str): if self.finished: return self._lexer.append_string(args_part) - # TODO: maybe don't extract detail if it's already stable - argument = self._extract_worklog_argument(self._lexer.complete_json(), self._tool_name) - if argument and argument != self._argument: - self._argument = argument - self._renderable = self._compose() + encoding = "utf-8" + self._argument_scan_bytes += len(args_part.encode(encoding=encoding)) + terminal_delimiter = args_part.rstrip().endswith(("}", "]")) + if ( + not self._argument + or self._argument_scan_bytes >= _ARGUMENT_SCAN_GROWTH_BYTES + or terminal_delimiter + ): + self._scan_worklog_argument() def mark_execution_started(self) -> None: # Terminal states are monotonic: a late ToolExecutionStarted (event # reordering, duplicate delivery) must not restyle a finished row. if self._execution_started or self.finished: return + self._scan_worklog_argument() self._execution_started = True if self._tui_card is not None: self._tui_card.mark_execution_started() self._renderable = self._compose() + def _scan_worklog_argument(self) -> None: + argument = self._extract_worklog_argument(self._lexer.complete_json(), self._tool_name) + self._argument_scan_bytes = 0 + if argument and argument != self._argument: + self._argument = argument + self._renderable = self._compose() + def append_output_part(self, text: str, *, stream: str = "output") -> None: if self.finished or not text: return diff --git a/tests/ui_and_conv/test_background_completion_watcher.py b/tests/ui_and_conv/test_background_completion_watcher.py index 9fef54d5..e06967cc 100644 --- a/tests/ui_and_conv/test_background_completion_watcher.py +++ b/tests/ui_and_conv/test_background_completion_watcher.py @@ -10,7 +10,12 @@ import pytest from pythinker_code.soul import Soul -from pythinker_code.ui.shell import Shell, _BackgroundCompletionWatcher, _PromptEvent +from pythinker_code.ui.shell import ( + PromptEventKind, + Shell, + _BackgroundCompletionWatcher, + _PromptEvent, +) def _make_watcher( @@ -42,7 +47,7 @@ async def test_pending_notification_and_empty_queue_waits_for_user_input(): await asyncio.sleep(0) assert task.done() is False - event = _PromptEvent(kind="input") + event = _PromptEvent(kind=PromptEventKind.INPUT) await queue.put(event) result = await task assert result is event @@ -53,7 +58,7 @@ async def test_pending_notification_but_user_input_queued_returns_event(): """Pending LLM notification + queued user input → user input wins.""" watcher = _make_watcher(has_pending=True, can_auto_trigger_pending=False) queue: asyncio.Queue[_PromptEvent] = asyncio.Queue() - event = _PromptEvent(kind="input") + event = _PromptEvent(kind=PromptEventKind.INPUT) await queue.put(event) result = await watcher.wait_for_next(queue) @@ -65,7 +70,7 @@ async def test_pending_notification_but_eof_queued_returns_eof(): """Pending notification + queued EOF → user can still exit.""" watcher = _make_watcher(has_pending=True, can_auto_trigger_pending=False) queue: asyncio.Queue[_PromptEvent] = asyncio.Queue() - eof = _PromptEvent(kind="eof") + eof = _PromptEvent(kind=PromptEventKind.EOF) await queue.put(eof) result = await watcher.wait_for_next(queue) @@ -123,7 +128,7 @@ async def _set_event(): asyncio.create_task(_set_event()) result = await watcher.wait_for_next(queue) assert result is not None - assert result.kind == "bg_noop" + assert result.kind is PromptEventKind.BACKGROUND_NOOP @pytest.mark.asyncio @@ -140,7 +145,7 @@ async def _set_event(): asyncio.create_task(_set_event()) result = await watcher.wait_for_next(queue) assert result is not None - assert result.kind == "bg_noop" + assert result.kind is PromptEventKind.BACKGROUND_NOOP @pytest.mark.asyncio @@ -148,7 +153,7 @@ async def test_user_input_wins_over_simultaneous_bg_event(): """Both idle and bg fire simultaneously → user input takes priority.""" watcher = _make_watcher() queue: asyncio.Queue[_PromptEvent] = asyncio.Queue() - event = _PromptEvent(kind="input") + event = _PromptEvent(kind=PromptEventKind.INPUT) # Both ready before await await queue.put(event) @@ -173,7 +178,7 @@ async def test_disabled_watcher_just_awaits_idle(): assert not watcher.enabled queue: asyncio.Queue[_PromptEvent] = asyncio.Queue() - event = _PromptEvent(kind="input") + event = _PromptEvent(kind=PromptEventKind.INPUT) await queue.put(event) result = await watcher.wait_for_next(queue) @@ -192,6 +197,7 @@ def __init__( self._recent = recent self._remaining = remaining self._event = asyncio.Event() + self.wait_cancelled = False def has_pending_input(self) -> bool: return self._pending @@ -203,7 +209,11 @@ def recent_input_activity_remaining(self, *, within_s: float) -> float: return self._remaining async def wait_for_input_activity(self) -> None: - await self._event.wait() + try: + await self._event.wait() + except asyncio.CancelledError: + self.wait_cancelled = True + raise self._event.clear() @@ -236,7 +246,7 @@ async def test_shell_wait_for_input_or_activity_returns_activity_event() -> None prompt._event.set() result = await task - assert result.kind == "input_activity" + assert result.kind is PromptEventKind.INPUT_ACTIVITY @pytest.mark.asyncio @@ -244,7 +254,7 @@ async def test_shell_wait_for_input_or_activity_returns_idle_event() -> None: shell = Shell(cast(Soul, SimpleNamespace(available_slash_commands=[], name="x")), None) prompt = _FakePromptActivity() queue: asyncio.Queue[_PromptEvent] = asyncio.Queue() - expected = _PromptEvent(kind="input") + expected = _PromptEvent(kind=PromptEventKind.INPUT) task = asyncio.create_task(shell._wait_for_input_or_activity(prompt, queue)) await queue.put(expected) @@ -263,5 +273,53 @@ async def test_shell_wait_for_input_or_activity_times_out_for_recent_activity_on result = await shell._wait_for_input_or_activity(prompt, queue, timeout_s=0.05) elapsed = asyncio.get_running_loop().time() - started - assert result.kind == "input_activity" + assert result.kind is PromptEventKind.BACKGROUND_GRACE_EXPIRED assert elapsed >= 0.04 + + +@pytest.mark.asyncio +async def test_shell_wait_for_input_or_activity_user_input_wins_activity_tie() -> None: + shell = Shell(cast(Soul, SimpleNamespace(available_slash_commands=[], name="x")), None) + prompt = _FakePromptActivity() + queue: asyncio.Queue[_PromptEvent] = asyncio.Queue() + expected = _PromptEvent(kind=PromptEventKind.INPUT) + + await queue.put(expected) + prompt._event.set() + + result = await shell._wait_for_input_or_activity(prompt, queue, timeout_s=0) + assert result is expected + + +@pytest.mark.asyncio +async def test_shell_wait_for_input_or_activity_cleans_up_tasks_when_cancelled() -> None: + shell = Shell(cast(Soul, SimpleNamespace(available_slash_commands=[], name="x")), None) + prompt = _FakePromptActivity() + queue: asyncio.Queue[_PromptEvent] = asyncio.Queue() + + task = asyncio.create_task(shell._wait_for_input_or_activity(prompt, queue, timeout_s=30)) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert prompt.wait_cancelled is True + expected = _PromptEvent(kind=PromptEventKind.INPUT) + await queue.put(expected) + assert queue.get_nowait() is expected + + +@pytest.mark.asyncio +async def test_background_watcher_cleans_up_waiters_when_cancelled() -> None: + watcher = _make_watcher() + queue: asyncio.Queue[_PromptEvent] = asyncio.Queue() + + task = asyncio.create_task(watcher.wait_for_next(queue)) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + expected = _PromptEvent(kind=PromptEventKind.INPUT) + await queue.put(expected) + assert queue.get_nowait() is expected diff --git a/tests/ui_and_conv/test_shell_bang_spawn.py b/tests/ui_and_conv/test_shell_bang_spawn.py index e0d7eda1..8b2c66d1 100644 --- a/tests/ui_and_conv/test_shell_bang_spawn.py +++ b/tests/ui_and_conv/test_shell_bang_spawn.py @@ -6,32 +6,47 @@ import pytest import pythinker_code.ui.shell as shell_module +from pythinker_code.ui.shell.command_runner import ShellCommandResult @pytest.mark.asyncio -async def test_bang_command_detaches_windows_console(monkeypatch) -> None: - """`!` foreground commands must not attach to the interactive console on Windows. - - A child sharing the TUI's console can clear it or reset its modes via the - Win32 console API (bypassing the stdio pipes), blanking the shell UI until - the terminal is restarted; CREATE_NO_WINDOW detaches it. +async def test_bang_command_runs_through_host_exec_runner(monkeypatch) -> None: + """`!` foreground commands must execute via ShellCommandRunner/Host.exec. + + Host.exec applies CREATE_NO_WINDOW on Windows (covered by + packages/pythinker-host/tests/test_local_host.py), which keeps a child off + the interactive console: console-API writes from a child sharing the TUI's + console bypass the stdio pipes and can blank the shell UI until terminal + restart. The shell must therefore never fall back to + asyncio.create_subprocess_shell for `!` commands. """ shell = object.__new__(shell_module.Shell) - captured: dict[str, object] = {} + ran: dict[str, object] = {} + + class _RecordingRunner: + async def run(self, command: str) -> ShellCommandResult: + ran["command"] = command + return ShellCommandResult( + stdout="hi\n", + stderr="", + returncode=0, + truncated_stdout=False, + truncated_stderr=False, + ) - async def _capture_and_abort(command: str, **kwargs): - captured.update(kwargs) - raise RuntimeError("spawn intercepted by test") + async def _forbidden_shell_spawn(command: str, **kwargs): + raise AssertionError("`!` commands must not use create_subprocess_shell") - monkeypatch.setattr(asyncio, "create_subprocess_shell", _capture_and_abort) + monkeypatch.setattr(asyncio, "create_subprocess_shell", _forbidden_shell_spawn) + monkeypatch.setattr(shell_module, "ShellCommandRunner", _RecordingRunner) monkeypatch.setattr(shell_module, "install_sigint_handler", lambda loop, handler: lambda: None) - # Swap the module's `os` reference rather than mutating the global - # `os.name`, which corrupts pathlib/pytest path handling mid-run. - monkeypatch.setattr(shell_module, "os", SimpleNamespace(name="nt")) + printed: list[object] = [] + monkeypatch.setattr( + shell_module, "console", SimpleNamespace(print=lambda value=None: printed.append(value)) + ) await shell._run_shell_command("echo hi") - flags = captured["creationflags"] - assert isinstance(flags, int) - assert flags & 0x08000000 # CREATE_NO_WINDOW + assert ran["command"] == "echo hi" + assert printed diff --git a/tests/ui_and_conv/test_shell_command_runner.py b/tests/ui_and_conv/test_shell_command_runner.py new file mode 100644 index 00000000..f7d126e9 --- /dev/null +++ b/tests/ui_and_conv/test_shell_command_runner.py @@ -0,0 +1,182 @@ +from __future__ import annotations + +import asyncio +import os +import shlex +from pathlib import Path +from types import SimpleNamespace +from typing import cast + +import pytest +from pythinker_host import get_current_host +from pythinker_host.path import HostPath + +import pythinker_code.ui.shell as shell_module +from pythinker_code.soul import Soul +from pythinker_code.ui.shell.command_runner import ShellCommandRunner +from pythinker_code.ui.shell.components.render_utils import render_plain +from pythinker_code.utils.environment import Environment + + +def _environment(shell_path: str = "/bin/sh") -> Environment: + return Environment( + os_kind="macOS", + os_arch="test", + os_version="test", + shell_name="sh", + shell_path=HostPath(shell_path), + ) + + +def _runner( + *, shell_path: str = "/bin/sh", output_limit_bytes: int = 1024 * 1024 +) -> ShellCommandRunner: + return ShellCommandRunner( + get_current_host(), + environment=_environment(shell_path), + output_limit_bytes=output_limit_bytes, + ) + + +@pytest.mark.asyncio +@pytest.mark.skipif(os.name == "nt", reason="requires a POSIX shell") +async def test_runner_preserves_quoting_and_pipelines() -> None: + result = await _runner().run("printf '%s\\n' 'hello world' | tr 'a-z' 'A-Z'") + + assert result.stdout == "HELLO WORLD\n" + assert result.stderr == "" + assert result.returncode == 0 + + +@pytest.mark.asyncio +@pytest.mark.skipif(os.name == "nt", reason="requires a POSIX shell") +async def test_runner_supports_redirection(tmp_path: Path) -> None: + destination = tmp_path / "redirected output.txt" + command = f"printf redirected > {shlex.quote(str(destination))}" + + result = await _runner().run(command) + encoding = "utf-8" + + assert destination.read_text(encoding=encoding) == "redirected" + assert result.stdout == "" + assert result.returncode == 0 + + +@pytest.mark.asyncio +@pytest.mark.skipif(os.name == "nt", reason="requires a POSIX shell") +async def test_runner_captures_stderr_and_nonzero_exit() -> None: + result = await _runner().run("printf problem >&2; exit 7") + + assert result.stdout == "" + assert result.stderr == "problem" + assert result.returncode == 7 + rendered = shell_module._format_local_shell_output( + stdout=result.stdout, + stderr=result.stderr, + returncode=result.returncode, + ) + assert rendered is not None + assert render_plain(rendered).splitlines() == ["problem", "exit 7"] + + +@pytest.mark.asyncio +@pytest.mark.skipif(os.name == "nt", reason="requires a POSIX shell") +async def test_runner_reports_no_output() -> None: + result = await _runner().run(":") + + assert result.stdout == "" + assert result.stderr == "" + assert result.returncode == 0 + rendered = shell_module._format_local_shell_output( + stdout=result.stdout, + stderr=result.stderr, + returncode=result.returncode, + ) + assert rendered is not None + assert render_plain(rendered).strip() == "(No output)" + + +@pytest.mark.asyncio +@pytest.mark.skipif(os.name == "nt", reason="requires a POSIX shell") +async def test_runner_truncates_stdout_and_stderr_separately() -> None: + command = "head -c 2048 /dev/zero | tr '\\0' o; head -c 2048 /dev/zero | tr '\\0' e >&2" + result = await _runner(output_limit_bytes=1024).run(command) + + assert result.truncated_stdout is True + assert result.truncated_stderr is True + assert result.stdout.startswith("o" * 1024) + assert result.stderr.startswith("e" * 1024) + assert result.stdout.endswith("\n... output truncated ...\n") + assert result.stderr.endswith("\n... output truncated ...\n") + + +def test_runner_builds_windows_powershell_argv() -> None: + argv = ShellCommandRunner.build_argv( + "Write-Output 'hello'", + shell_name="Windows PowerShell", + shell_path=r"C:\Windows\System32\WindowsPowerShell\v1.0\powershell.exe", + ) + + assert argv == ( + r"C:\Windows\System32\WindowsPowerShell\v1.0\powershell.exe", + "-command", + "Write-Output 'hello'", + ) + + +@pytest.mark.asyncio +@pytest.mark.skipif(os.name == "nt", reason="requires POSIX process inspection") +async def test_runner_cancellation_terminates_and_reaps_process(tmp_path: Path) -> None: + pid_path = tmp_path / "shell.pid" + task = asyncio.create_task( + _runner().run(f"echo $$ > {shlex.quote(str(pid_path))}; exec sleep 30") + ) + for _ in range(100): + if pid_path.exists(): + break + await asyncio.sleep(0.01) + assert pid_path.exists() + + encoding = "utf-8" + pid = int(pid_path.read_text(encoding=encoding).strip()) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + with pytest.raises(ProcessLookupError): + os.kill(pid, 0) + + +@pytest.mark.asyncio +async def test_invalid_shell_path_surfaces_existing_failure_message(monkeypatch) -> None: + class _InvalidRunner(ShellCommandRunner): + def __init__(self) -> None: + super().__init__( + get_current_host(), + environment=_environment("/definitely/missing/pythinker-shell"), + ) + + printed: list[object] = [] + fake_console = SimpleNamespace(print=lambda value=None: printed.append(value)) + monkeypatch.setattr(shell_module, "ShellCommandRunner", _InvalidRunner) + monkeypatch.setattr(shell_module, "console", fake_console) + monkeypatch.setattr(shell_module, "install_sigint_handler", lambda loop, handler: lambda: None) + shell = shell_module.Shell( + cast(Soul, SimpleNamespace(available_slash_commands=[], name="test")), + None, + ) + + await shell._run_shell_command("printf hello") + + assert "Failed to run shell command:" in str(printed[-1]) + + +def test_shell_output_rendering_sanitizes_control_characters() -> None: + rendered = shell_module._format_local_shell_output( + stdout="safe\x1b[2J\x01text\n", + stderr="", + returncode=0, + ) + + assert rendered is not None + assert render_plain(rendered).strip() == "safetext" diff --git a/tests/ui_and_conv/test_tool_call_block.py b/tests/ui_and_conv/test_tool_call_block.py index 86b4c5b1..ea846a09 100644 --- a/tests/ui_and_conv/test_tool_call_block.py +++ b/tests/ui_and_conv/test_tool_call_block.py @@ -14,6 +14,7 @@ get_tool_renderer, register_builtin_renderers, ) +from pythinker_code.ui.shell.visualize import _blocks as blocks_module from pythinker_code.ui.shell.visualize import _ToolCallBlock, _worklog from pythinker_code.wire.types import ToolResult @@ -390,6 +391,33 @@ def test_streaming_args_never_flash_invalid_badge( assert "" not in rendered, f"flashed after streaming {ch!r}" +def test_streaming_large_arguments_bounds_label_extraction_scans(monkeypatch) -> None: + original = blocks_module.extract_key_argument + extraction_calls = 0 + first_non_empty_call: int | None = None + + def _recording_extract(arguments: str, tool_name: str) -> str | None: + nonlocal extraction_calls, first_non_empty_call + extraction_calls += 1 + result = original(arguments, tool_name) + if result and first_non_empty_call is None: + first_non_empty_call = extraction_calls + return result + + monkeypatch.setattr(blocks_module, "extract_key_argument", _recording_extract) + block = _ToolCallBlock(_tool_call("Shell", "")) + arguments = json.dumps({"command": "x" * (64 * 1024)}) + encoding = "utf-8" + for offset in range(0, len(arguments), 64): + block.append_args_part(arguments[offset : offset + 64]) + block.mark_execution_started() + + assert first_non_empty_call is not None + scanned_after_label = extraction_calls - first_non_empty_call + growth_scans = (len(arguments.encode(encoding=encoding)) + 1023) // 1024 + assert scanned_after_label <= growth_scans + 1 + + def test_finished_call_with_non_string_command_still_shows_invalid_badge( _card_style_with_builtin_renderers, ): From 993deb7c270688b7ad3e63fa4fe3bfafb70deb09 Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 21 Jul 2026 01:22:46 -0400 Subject: [PATCH 3/4] docs(changelog): record lossless wire replay and shell-exec hardening --- CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index d972a9e1..cdc0a181 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,9 @@ GitHub Releases page; `0.8.0` is the new starting line. ## Unreleased +- Render image/audio/video and unknown content parts as payload-free labels in the live view, roll nested subagent activity (up to 16 levels) under the correct root tool card, and stop provider-remapped tool results from starting replay user turns. +- **Behavior change:** `!` shell commands now run through the detected configured shell (` -c`, PowerShell `-command` on Windows) instead of the implicit `/bin/sh`/`cmd.exe`, with separate 1 MiB stdout/stderr caps and cancellation cleanup; existing `cmd.exe`-syntax commands on Windows may need updating. +- Distinguish background auto-trigger grace expiry from user input activity with typed prompt events, and bound streamed tool-argument label scans to 1 KiB growth boundaries. - Scope prompt Git status, toasts, history, and clipboard to each shell session with a new `/prompt-history status|clear` command, bounded locked history storage, and cross-session isolation. - Unify the shell footer behind one per-frame view model shared by the legacy and card renderers, and key theme style-resolver caching by terminal capabilities. - Fix a rare queued-follow-up "ghost card": echoing a drained queued command now commits through the scrollback handoff (which hides the input card before the terminal teardown erases the prompt), so the input-card border can no longer fossilize into scrollback above the echoed command under heavy load. From 044f52321731135e15383be2482d027794b1a632 Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 21 Jul 2026 12:57:55 -0400 Subject: [PATCH 4/4] fix(tui): guard shell-exec cleanup and purge stale subagent ancestry Address review findings on the session-scoped shell/live-view work: - ShellCommandRunner.run: move stdin.close() and stream/wait task creation inside the guarded scope so a failure there terminates and reaps the child instead of leaking it; _terminate_and_reap now tolerates a short task tuple. - _LiveView: purge subagent tool-call ancestry entries when their root block is flushed, so nested-event ids don't accumulate for the rest of a turn and a late/duplicate nested event falls back to "missing ancestry" instead of mutating an archived block. - _LiveView: scope the unknown-content-part "log once" set to each view instead of a process-global, keeping the behavior deterministic and tests isolated. Add regression tests for the stdin-close cleanup path, the ancestry purge on root flush, and the negative output-limit guard. --- src/pythinker_code/ui/shell/command_runner.py | 31 ++++++------ .../ui/shell/visualize/_live_view.py | 29 +++++++++-- .../test_nested_subagent_events.py | 31 ++++++++++++ .../ui_and_conv/test_shell_command_runner.py | 48 ++++++++++++++++++- 4 files changed, 120 insertions(+), 19 deletions(-) diff --git a/src/pythinker_code/ui/shell/command_runner.py b/src/pythinker_code/ui/shell/command_runner.py index 832ebaaf..8a03ca08 100644 --- a/src/pythinker_code/ui/shell/command_runner.py +++ b/src/pythinker_code/ui/shell/command_runner.py @@ -2,6 +2,7 @@ import asyncio from dataclasses import dataclass +from typing import Any from pythinker_host import AsyncReadable, Host, HostProcess, get_current_host @@ -52,17 +53,21 @@ async def run(self, command: str) -> ShellCommandResult: shell_path=str(environment.shell_path), ) process = await self._host.exec(*argv, env=get_clean_env()) - process.stdin.close() - - stdout_task = asyncio.create_task( - self._read_stream_limited(process.stdout, self._output_limit_bytes) - ) - stderr_task = asyncio.create_task( - self._read_stream_limited(process.stderr, self._output_limit_bytes) - ) - wait_task = asyncio.create_task(process.wait()) - tasks = (stdout_task, stderr_task, wait_task) + # Everything after a successful spawn runs inside the guarded scope: if + # stdin.close() or task creation raises, the child would otherwise leak + # with nothing left to terminate/reap it. `tasks` starts empty so cleanup + # is safe even when the failure happens before any task is created. + tasks: tuple[asyncio.Task[Any], ...] = () try: + process.stdin.close() + stdout_task = asyncio.create_task( + self._read_stream_limited(process.stdout, self._output_limit_bytes) + ) + stderr_task = asyncio.create_task( + self._read_stream_limited(process.stderr, self._output_limit_bytes) + ) + wait_task = asyncio.create_task(process.wait()) + tasks = (stdout_task, stderr_task, wait_task) await asyncio.gather(*tasks) except asyncio.CancelledError: await self._terminate_and_reap(process, tasks) @@ -107,11 +112,7 @@ async def _read_stream_limited(stream: AsyncReadable, limit: int) -> tuple[bytes @staticmethod async def _terminate_and_reap( process: HostProcess, - tasks: tuple[ - asyncio.Task[tuple[bytes, bool]], - asyncio.Task[tuple[bytes, bool]], - asyncio.Task[int], - ], + tasks: tuple[asyncio.Task[Any], ...], ) -> None: try: await process.kill() diff --git a/src/pythinker_code/ui/shell/visualize/_live_view.py b/src/pythinker_code/ui/shell/visualize/_live_view.py index c0f2f174..9aa95004 100644 --- a/src/pythinker_code/ui/shell/visualize/_live_view.py +++ b/src/pythinker_code/ui/shell/visualize/_live_view.py @@ -152,7 +152,6 @@ _WORKING_TIP_MIN_ELAPSED_S = 4.0 _MAX_PINNED_TODO_ROWS = 5 _MAX_SUBAGENT_EVENT_DEPTH = 16 -_SEEN_UNKNOWN_CONTENT_PART_TYPES: set[str] = set() def _todo_activity_label(label: str) -> str: @@ -263,6 +262,10 @@ def __init__( self._current_content_block: _ContentBlock | None = None self._tool_call_blocks: dict[str, _ToolCallBlock] = {} self._subagent_tool_call_ancestry: dict[str, tuple[_ToolCallBlock, int]] = {} + # Per-view so the "log an unknown content-part type once" guarantee is + # scoped to this session/view rather than the whole process (keeps the + # log-once behavior deterministic and test-isolated). + self._seen_unknown_content_part_types: set[str] = set() self._last_tool_call_block: _ToolCallBlock | None = None self._held_tool_search_block: _ToolCallBlock | None = None self._completed_expandable_tool_blocks = deque[_ToolCallBlock](maxlen=20) @@ -1648,6 +1651,11 @@ def flush_finished_tool_calls(self) -> None: self._archive_completed_tool_card(block) self._tool_call_blocks.pop(tool_call_id) + # Drop ancestry entries for this now-archived root so nested-event + # ids don't accumulate for the rest of the turn and a late/duplicate + # nested event resolves via the "missing ancestry" fallback instead + # of silently mutating a block that is no longer rendered live. + self._purge_ancestry_for_block(block) if self._last_tool_call_block == block: self._last_tool_call_block = None if block.is_tool_search: @@ -1659,6 +1667,21 @@ def flush_finished_tool_calls(self) -> None: self._emit_action_block(block.compose()) self.refresh_soon() + def _purge_ancestry_for_block(self, block: _ToolCallBlock) -> None: + """Drop subagent ancestry entries whose root is *block*. + + All nested tool-call ids under a subagent tree resolve to the same root + ``_ToolCallBlock`` (the ancestry fallback propagates it down), so once + that root is flushed its whole subtree of ancestry entries is stale. + """ + stale_ids = [ + tool_call_id + for tool_call_id, (ancestor, _depth) in self._subagent_tool_call_ancestry.items() + if ancestor is block + ] + for tool_call_id in stale_ids: + del self._subagent_tool_call_ancestry[tool_call_id] + def flush_notifications(self) -> None: """Flush rendered notifications to terminal history.""" self._live_notification_blocks.clear() @@ -1706,8 +1729,8 @@ def append_content(self, part: ContentPart) -> None: self._append_content_label("[video]") case _: part_type = part.type - if part_type not in _SEEN_UNKNOWN_CONTENT_PART_TYPES: - _SEEN_UNKNOWN_CONTENT_PART_TYPES.add(part_type) + if part_type not in self._seen_unknown_content_part_types: + self._seen_unknown_content_part_types.add(part_type) logger.debug( "Rendering unknown content part type in live view: {part_type}", part_type=part_type, diff --git a/tests/ui_and_conv/test_nested_subagent_events.py b/tests/ui_and_conv/test_nested_subagent_events.py index 9fd39064..00923e60 100644 --- a/tests/ui_and_conv/test_nested_subagent_events.py +++ b/tests/ui_and_conv/test_nested_subagent_events.py @@ -192,3 +192,34 @@ def test_recursive_depth_overflow_renders_one_fallback( assert "Nested subagent activity unavailable" in output assert "depth limit reached" in output assert "TOO_DEEP" not in output + + +def test_flushing_root_agent_purges_nested_ancestry() -> None: + view = _view() + # Two nested tool calls under the root Agent populate the ancestry map. + view.dispatch_wire_message( + _nested_event([], _tool_call("nested-1", "Agent", '{"description":"deeper"}')) + ) + view.dispatch_wire_message( + _nested_event(["nested-1"], _tool_call("nested-2", "Read", '{"file_path":"x"}')) + ) + root_block = view._tool_call_blocks["root-agent"] + assert set(view._subagent_tool_call_ancestry) == {"nested-1", "nested-2"} + + # Finishing the root Agent tool call flushes it out of the live area. + view.dispatch_wire_message( + ToolResult(tool_call_id="root-agent", return_value=ToolOk(output="done")) + ) + + assert "root-agent" not in view._tool_call_blocks + # The whole nested subtree's ancestry is purged, not left dangling on the + # now-archived root block. + assert view._subagent_tool_call_ancestry == {} + assert root_block is not None # kept referenced so identity purge is exercised + + # A late nested event for a now-purged parent falls back to "missing + # ancestry" instead of resolving to (and mutating) the archived root block. + view.dispatch_wire_message( + _nested_event(["nested-1"], _tool_call("nested-3", "Read", '{"file_path":"y"}')) + ) + assert "nested-3" not in view._subagent_tool_call_ancestry diff --git a/tests/ui_and_conv/test_shell_command_runner.py b/tests/ui_and_conv/test_shell_command_runner.py index f7d126e9..c15f8806 100644 --- a/tests/ui_and_conv/test_shell_command_runner.py +++ b/tests/ui_and_conv/test_shell_command_runner.py @@ -8,7 +8,7 @@ from typing import cast import pytest -from pythinker_host import get_current_host +from pythinker_host import Host, get_current_host from pythinker_host.path import HostPath import pythinker_code.ui.shell as shell_module @@ -38,6 +38,52 @@ def _runner( ) +def test_runner_rejects_negative_output_limit() -> None: + with pytest.raises(ValueError): + ShellCommandRunner( + get_current_host(), + environment=_environment(), + output_limit_bytes=-1, + ) + + +@pytest.mark.asyncio +async def test_runner_reaps_process_when_stdin_close_fails() -> None: + """A failure after spawn (here stdin.close()) must still terminate/reap the child.""" + + class _FailingStdin: + def close(self) -> None: + raise BrokenPipeError("stdin already closed") + + class _FakeProcess: + def __init__(self) -> None: + self.stdin = _FailingStdin() + self.killed = False + self.reaped = False + + async def kill(self) -> None: + self.killed = True + + async def wait(self) -> int: + self.reaped = True + return 0 + + process = _FakeProcess() + + class _FakeHost: + async def exec(self, *_argv: str, env: dict[str, str] | None = None) -> _FakeProcess: + return process + + runner = ShellCommandRunner(cast(Host, _FakeHost()), environment=_environment()) + + with pytest.raises(BrokenPipeError): + await runner.run("echo hi") + + # The spawned child was terminated and reaped rather than leaked. + assert process.killed is True + assert process.reaped is True + + @pytest.mark.asyncio @pytest.mark.skipif(os.name == "nt", reason="requires a POSIX shell") async def test_runner_preserves_quoting_and_pipelines() -> None: