From 8917ae3529ea13359266009a286b32099b3fb929 Mon Sep 17 00:00:00 2001 From: Eric Lee Date: Fri, 15 May 2026 01:15:52 -0700 Subject: [PATCH] refactor(context): make ToolContext.abort_controller non-optional MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous ``Any | None = None`` default papered over a hazard class that masked the ESC-into-subagent regression PR #135 fixed: every caller that owns the per-run cancellation lifecycle (TUI bridge, REPL engine) had to remember to plumb the controller onto the context, and a forgotten plumbing call silently degraded ESC propagation rather than failing loudly. Use ``field(default_factory=AbortController)`` so every fresh context carries a real (untripped) controller. Downstream readers (subagent inheritance, streaming executor, Bash supervisor, tool execution, tool hooks, tasks_v2 abort fast-path) drop the defensive ``or AbortController()`` and ``if abort_ctrl and …`` guards that existed only to compensate for the impossible-now ``None`` state. The TUI bridge's ``_finish()`` now installs a fresh controller instead of clearing to ``None`` so the dataclass invariant holds across runs. Three new contract tests pin the invariants: default exists, default is per-context (not shared), explicit override still wins. Co-Authored-By: Claude Opus 4.7 --- src/agent/run_agent.py | 12 +++-- src/agent/subagent_context.py | 11 ++-- .../tool_execution/streaming_executor.py | 18 +++++-- src/services/tool_execution/tool_execution.py | 6 ++- src/services/tool_execution/tool_hooks.py | 6 ++- src/tool_system/context.py | 11 +++- src/tool_system/tools/bash/bash_tool.py | 8 +-- src/tool_system/tools/tasks_v2.py | 7 +-- src/tui/agent_bridge.py | 23 ++++---- tests/test_esc_cancel_propagation.py | 38 ++++++++----- tests/test_tool_context_abort_default.py | 54 +++++++++++++++++++ 11 files changed, 149 insertions(+), 45 deletions(-) create mode 100644 tests/test_tool_context_abort_default.py diff --git a/src/agent/run_agent.py b/src/agent/run_agent.py index d42606d1b..970e8ed3e 100644 --- a/src/agent/run_agent.py +++ b/src/agent/run_agent.py @@ -269,18 +269,22 @@ async def run_agent(params: RunAgentParams) -> AsyncGenerator[Message, None]: or get_agent_system_prompt(agent_def, params.parent_system_prompt) ) - # Determine abort controller + # Determine abort controller. + # ``params.parent_context.abort_controller`` is now non-optional on + # the ``ToolContext`` dataclass, so the legacy "parent has no + # controller → mint a fresh one" branch is gone. The remaining + # priority order is: explicit caller override → fresh controller + # for async (so background agents survive parent cancel) → + # share with parent for sync (so parent ESC propagates). if params.abort_controller is not None: abort_controller = params.abort_controller elif params.is_async: # Async agents run independently in the background and should survive # parent cancellation events. abort_controller = AbortController() - elif params.parent_context.abort_controller is not None: + else: # Sync agents share abort with parent abort_controller = params.parent_context.abort_controller - else: - abort_controller = AbortController() # Build permission context perm_context = _build_permission_context( diff --git a/src/agent/subagent_context.py b/src/agent/subagent_context.py index 18ed2af01..59a0879c7 100644 --- a/src/agent/subagent_context.py +++ b/src/agent/subagent_context.py @@ -74,15 +74,16 @@ def create_subagent_context( overrides = SubagentContextOverrides() # --- Abort controller --- - # Priority: explicit override > share parent's > new child linked to parent + # Priority: explicit override > share parent's > new child linked to parent. + # ``parent_context.abort_controller`` is now non-optional on the + # ``ToolContext`` dataclass, so the legacy "parent has no controller" + # branch is gone — every parent context carries a real controller. if overrides.abort_controller is not None: abort_controller = overrides.abort_controller - elif overrides.share_abort_controller and parent_context.abort_controller is not None: + elif overrides.share_abort_controller: abort_controller = parent_context.abort_controller - elif parent_context.abort_controller is not None: - abort_controller = create_child_abort_controller(parent_context.abort_controller) else: - abort_controller = AbortController() + abort_controller = create_child_abort_controller(parent_context.abort_controller) # --- Permission context --- # If sharing abort controller, it's interactive and can show UI. diff --git a/src/services/tool_execution/streaming_executor.py b/src/services/tool_execution/streaming_executor.py index 504436e7a..7ad7ca30c 100644 --- a/src/services/tool_execution/streaming_executor.py +++ b/src/services/tool_execution/streaming_executor.py @@ -24,7 +24,6 @@ create_user_message, ) from src.utils.abort_controller import ( - AbortController, AbortError, create_child_abort_controller, ) @@ -79,8 +78,14 @@ def __init__( self._tool_use_context = tool_use_context self._has_errored = False self._errored_tool_description = "" + # ``tool_use_context.abort_controller`` is non-optional on + # ``ToolContext`` — the previous defensive ``or AbortController()`` + # papered over the "field is None" hazard class that broke ESC + # propagation into subagents. Now that the field is guaranteed, + # the sibling controller always parents on the real per-run + # controller and ESC reaches every executing tool. self._sibling_abort_controller = create_child_abort_controller( - tool_use_context.abort_controller or AbortController() + tool_use_context.abort_controller ) self._discarded = False self._progress_available_event = asyncio.Event() @@ -207,8 +212,10 @@ def _get_abort_reason( return "streaming_fallback" if self._has_errored: return "sibling_error" + # ``abort_controller`` is non-optional on ``ToolContext`` — no + # need to guard truthiness before reading ``.signal.aborted``. ctx_abort = self._tool_use_context.abort_controller - if ctx_abort and ctx_abort.signal.aborted: + if ctx_abort.signal.aborted: if ctx_abort.signal.reason == "interrupt": behavior = self._get_tool_interrupt_behavior(tool) return "user_interrupted" if behavior == "cancel" else None @@ -307,9 +314,12 @@ async def collect_results() -> None: ) def _on_tool_abort() -> None: + # ``abort_controller`` is non-optional on the context, so + # the historical truthiness check is gone — we only need + # to guard against re-entering the parent's already-fired + # abort and against discarded executors. if ( tool_abort_controller.signal.reason != "sibling_error" - and self._tool_use_context.abort_controller and not self._tool_use_context.abort_controller.signal.aborted and not self._discarded ): diff --git a/src/services/tool_execution/tool_execution.py b/src/services/tool_execution/tool_execution.py index 9145b1ad4..cca94c006 100644 --- a/src/services/tool_execution/tool_execution.py +++ b/src/services/tool_execution/tool_execution.py @@ -81,8 +81,10 @@ async def run_tool_use( tool_input = tool_use.input if isinstance(tool_use.input, dict) else {} try: - abort_ctrl = tool_use_context.abort_controller - if abort_ctrl and abort_ctrl.signal.aborted: + # ``abort_controller`` is non-optional on ``ToolContext`` — the + # historical ``if abort_ctrl and …`` guard masked the field-is-None + # hazard class that broke ESC propagation into subagents. + if tool_use_context.abort_controller.signal.aborted: content = _create_tool_result_stop(tool_use.id) yield MessageUpdateLazy( message=create_user_message( diff --git a/src/services/tool_execution/tool_hooks.py b/src/services/tool_execution/tool_hooks.py index 6cea74668..2fecfd854 100644 --- a/src/services/tool_execution/tool_hooks.py +++ b/src/services/tool_execution/tool_hooks.py @@ -130,8 +130,10 @@ async def run_pre_tool_use_hooks( if result.get("message"): yield {"type": "message", "message": {"message": result["message"]}} - abort_ctrl = tool_use_context.abort_controller - if abort_ctrl and abort_ctrl.signal.aborted: + # ``abort_controller`` is non-optional on ``ToolContext``; + # the truthiness guard used to paper over the field-is-None + # hazard class. + if tool_use_context.abort_controller.signal.aborted: yield { "type": "message", "message": { diff --git a/src/tool_system/context.py b/src/tool_system/context.py index 822fc5717..e5401c51b 100644 --- a/src/tool_system/context.py +++ b/src/tool_system/context.py @@ -10,6 +10,7 @@ from src.permissions.types import ToolPermissionContext from src.services.swarm.agent_name_registry import AgentNameRegistry from src.task_registry import RuntimeTaskRegistry +from src.utils.abort_controller import AbortController def _resolve_path(p: str | Path) -> Path: @@ -143,7 +144,15 @@ class ToolContext: permission_handler: Callable[[str, str, Optional[str]], tuple[bool, bool]] | None = None options: ToolUseOptions = field(default_factory=ToolUseOptions) - abort_controller: Any | None = None + # Always present; callers that own the per-run cancellation lifecycle + # (TUI bridge, REPL engine) overwrite this with their own controller + # in ``submit()`` / ``__init__`` so tools, hooks, and subagents see + # the same signal the UI trips. The default factory keeps the field + # non-``None`` for unit tests and SDK callers that never explicitly + # set it — readers can drop the historical ``if ctrl and …`` / + # ``or AbortController()`` defensive checks that masked the "field + # is None" hazard class. + abort_controller: AbortController = field(default_factory=AbortController) messages: list[Any] = field(default_factory=list) set_response_length: Callable[[Callable[[int], int]], None] | None = None set_in_progress_tool_use_ids: Callable[[Callable[[set[str]], set[str]]], None] | None = None diff --git a/src/tool_system/tools/bash/bash_tool.py b/src/tool_system/tools/bash/bash_tool.py index 140206da6..67a9d1356 100644 --- a/src/tool_system/tools/bash/bash_tool.py +++ b/src/tool_system/tools/bash/bash_tool.py @@ -35,9 +35,11 @@ class _BashRunResult: timed_out: bool = False -def _get_abort_signal(context: ToolContext) -> Any | None: - controller = getattr(context, "abort_controller", None) - return getattr(controller, "signal", None) if controller else None +def _get_abort_signal(context: ToolContext) -> Any: + # ``abort_controller`` is non-optional on ``ToolContext``; the + # ``getattr(..., None)`` indirection used to paper over the + # historical "field is None" hazard class. + return context.abort_controller.signal def _kill_process_group(pid: int, sig: int) -> None: diff --git a/src/tool_system/tools/tasks_v2.py b/src/tool_system/tools/tasks_v2.py index 597b50f6c..5cb8197cb 100644 --- a/src/tool_system/tools/tasks_v2.py +++ b/src/tool_system/tools/tasks_v2.py @@ -648,9 +648,10 @@ async def _poll_runtime_until_terminal( # Abort fast-path — if the parent's controller is signalled, # exit with the current snapshot rather than waiting out the - # remaining timeout. - abort_signal = getattr(context.abort_controller, "signal", None) if context.abort_controller else None - if abort_signal is not None and getattr(abort_signal, "aborted", False): + # remaining timeout. ``abort_controller`` is non-optional on + # ``ToolContext`` so the historical truthiness/getattr indirection + # is gone. + if context.abort_controller.signal.aborted: return _runtime_task_to_output(task_id, runtime, context) remaining = deadline - time.monotonic() diff --git a/src/tui/agent_bridge.py b/src/tui/agent_bridge.py index 39f6b99a6..b70c75481 100644 --- a/src/tui/agent_bridge.py +++ b/src/tui/agent_bridge.py @@ -226,21 +226,26 @@ def _finish(self) -> None: self._state.clear_streaming_text() with self._busy_lock: self._busy = False - # Drop the per-run controller from the shared tool context so - # the next ``submit()`` starts from a clean state. Leaving an - # aborted controller in place would cause the next prompt's - # first tool dispatch to see ``signal.aborted == True`` and - # short-circuit before the user has even pressed ESC. + # Replace the per-run controller on the shared tool context + # with a fresh one so the next ``submit()`` starts from a + # clean state. Leaving an aborted controller in place would + # cause the next prompt's first tool dispatch to see + # ``signal.aborted == True`` and short-circuit before the + # user has even pressed ESC. + # + # The dataclass field is non-optional, so we can't simply + # clear to ``None`` — we install an untripped controller + # that mirrors the dataclass default. # # Safety note: ``_finish`` runs on the worker thread *after* # ``run_agent_loop`` has returned or raised, so no in-flight # tool can be reading ``context.abort_controller`` from this - # thread at the moment we clear it. Detached background + # thread at the moment we replace it. Detached background # processes (e.g. ``spawn_background_bash``) capture their # own controller reference at spawn time rather than re- - # reading the context field, so clearing here doesn't orphan - # them either. - self._tool_context.abort_controller = None + # reading the context field, so replacing here doesn't + # orphan them either. + self._tool_context.abort_controller = AbortController() # ---- permission bridge ---- def _permission_handler( diff --git a/tests/test_esc_cancel_propagation.py b/tests/test_esc_cancel_propagation.py index 90887a126..4a7ecfe91 100644 --- a/tests/test_esc_cancel_propagation.py +++ b/tests/test_esc_cancel_propagation.py @@ -282,8 +282,13 @@ def test_query_engine_plumbs_abort_controller_into_tool_context( from src.query.engine import QueryEngine, QueryEngineConfig context = ToolContext(workspace_root=tmp_path) - # Pre-condition: a freshly constructed context has no abort controller. - assert context.abort_controller is None + # Pre-condition: a freshly constructed context now carries a default + # (untripped) controller from the dataclass factory. The engine must + # OVERWRITE this with its own controller — otherwise the engine's + # ``interrupt()`` would trip a controller no tool can see. + default_ctrl = context.abort_controller + assert default_ctrl is not None + assert default_ctrl.signal.aborted is False cfg = QueryEngineConfig( cwd=tmp_path, @@ -294,12 +299,12 @@ def test_query_engine_plumbs_abort_controller_into_tool_context( ) engine = QueryEngine(cfg) - # Post-fix: the engine's controller is now visible on the context. + # Post-fix: the engine's controller has replaced the dataclass default. assert context.abort_controller is engine._abort_controller + assert context.abort_controller is not default_ctrl # Interrupting the engine trips the context-visible signal. engine.interrupt() - assert context.abort_controller is not None assert context.abort_controller.signal.aborted is True # Reset replaces the controller on both the engine and the context. @@ -334,7 +339,11 @@ def _no_run_worker(*_args: Any, **_kwargs: Any) -> None: return None context = ToolContext(workspace_root=tmp_path) - assert context.abort_controller is None + # The dataclass factory always installs a default (untripped) + # controller; the bridge replaces it on each ``submit()``. + default_ctrl = context.abort_controller + assert default_ctrl is not None + assert default_ctrl.signal.aborted is False bridge = AgentBridge( post_message=_post, @@ -351,16 +360,21 @@ def _no_run_worker(*_args: Any, **_kwargs: Any) -> None: # The controller created inside submit() must be visible on the # shared tool context — this is the wiring that lets subagents - # honour ESC. + # honour ESC. It must REPLACE the dataclass default (otherwise the + # bridge's ``cancel()`` would trip a controller no tool can see). assert context.abort_controller is bridge._abort_controller + assert context.abort_controller is not default_ctrl # Cancelling trips both objects (they are the same controller). + aborted_ctrl = context.abort_controller assert bridge.cancel() is True - assert context.abort_controller is not None - assert context.abort_controller.signal.aborted is True + assert aborted_ctrl.signal.aborted is True - # Finishing a run clears the per-run controller from the shared - # context so the next prompt doesn't start with a stale aborted - # signal that would short-circuit every tool dispatch. + # Finishing a run swaps in a FRESH (untripped) controller so the + # next prompt doesn't start with a stale aborted signal that would + # short-circuit every tool dispatch. The field is non-optional, so + # we install a new controller rather than clearing to ``None``. bridge._finish() - assert context.abort_controller is None + assert context.abort_controller is not None + assert context.abort_controller is not aborted_ctrl + assert context.abort_controller.signal.aborted is False diff --git a/tests/test_tool_context_abort_default.py b/tests/test_tool_context_abort_default.py new file mode 100644 index 000000000..0666a1fb7 --- /dev/null +++ b/tests/test_tool_context_abort_default.py @@ -0,0 +1,54 @@ +"""Pin the ``ToolContext.abort_controller`` non-optional contract. + +Previously the field defaulted to ``None`` and every reader had to +defensively guard with ``getattr(..., None)`` or ``ctrl and …``. The +"field is None" hazard class is what allowed the original ESC-into- +subagent bug to slip through: the bridge was supposed to plumb the +field but forgot, and silently every downstream tool ran a fresh +disconnected controller. + +The dataclass factory now installs an untripped ``AbortController`` on +every fresh context, so readers can drop the defensive checks and a +forgotten plumbing call no longer regresses ESC propagation into a +silent "the controller is None" landmine. +""" +from __future__ import annotations + +from pathlib import Path + +from src.tool_system.context import ToolContext +from src.utils.abort_controller import AbortController + + +def test_fresh_context_has_default_abort_controller(tmp_path: Path) -> None: + ctx = ToolContext(workspace_root=tmp_path) + + # The field is populated by the dataclass factory, not None. + assert ctx.abort_controller is not None + assert isinstance(ctx.abort_controller, AbortController) + # Default controller is untripped — readers can dispatch tools + # without worrying about an accidental "already aborted" state. + assert ctx.abort_controller.signal.aborted is False + + +def test_each_context_gets_its_own_controller(tmp_path: Path) -> None: + """Two contexts must NOT share the same default controller instance. + + A shared default would let one context's abort silently cascade into + another — a class of action-at-a-distance bug we deliberately avoid + by using ``default_factory`` instead of ``default=AbortController()``. + """ + a = ToolContext(workspace_root=tmp_path) + b = ToolContext(workspace_root=tmp_path) + assert a.abort_controller is not b.abort_controller + + a.abort_controller.abort("a-only") + assert a.abort_controller.signal.aborted is True + assert b.abort_controller.signal.aborted is False + + +def test_explicit_controller_overrides_default(tmp_path: Path) -> None: + """Callers that pass their own controller still win.""" + explicit = AbortController() + ctx = ToolContext(workspace_root=tmp_path, abort_controller=explicit) + assert ctx.abort_controller is explicit