From f6ad97274df160ce76aad336e4260cd0e2af7a0c Mon Sep 17 00:00:00 2001 From: Rohit C Prasad Date: Tue, 21 Jul 2026 23:16:40 -0700 Subject: [PATCH] engine: Stop works in every state, not just between iterations Stream drops between chunks (+ the pre-first-token wait); pending approvals/questions/plans resolve as interrupted. Running shell commands die via an executor interrupt hook; skipped tool calls still get results (no orphans). Also fixes delete_session calling a nonexistent engine.interrupt(). --- coworker/agent.py | 2 + coworker/engine.py | 119 ++++++++++++++++++++++--- coworker/server/manager.py | 4 +- coworker/tools/shell.py | 23 ++++- tests/test_engine_stop.py | 174 +++++++++++++++++++++++++++++++++++++ 5 files changed, 307 insertions(+), 15 deletions(-) create mode 100644 tests/test_engine_stop.py diff --git a/coworker/agent.py b/coworker/agent.py index 4f322b8d..d0f82a0e 100644 --- a/coworker/agent.py +++ b/coworker/agent.py @@ -308,6 +308,8 @@ def build_engine( model=model, instructions=instructions, approver=approver, + # Stop kills the in-flight foreground shell command, not just the loop. + interrupt_hooks=[executor.interrupt_now] if executor is not None else None, max_iterations=( max_iterations if max_iterations is not None else config.max_iterations ), diff --git a/coworker/engine.py b/coworker/engine.py index ac3bab15..2e24b8bb 100644 --- a/coworker/engine.py +++ b/coworker/engine.py @@ -73,6 +73,9 @@ class TurnEngine: question_asker: Optional[ Callable[[dict[str, Any]], "Awaitable[dict[str, Any]]"] ] = None, + # Called (thread-safe, best-effort) when the user stops the turn — e.g. the + # executor's kill for a running shell command. + interrupt_hooks: Optional[list[Callable[[], None]]] = None, ) -> None: self.provider = provider self.registry = registry @@ -111,10 +114,38 @@ class TurnEngine: # tool_call.id → the standing rule that auto-allowed it ("tool → target"), so the # TOOL_FINISHED event can carry the note to the tool card (§25). self._standing_notes: dict[str, str] = {} + self._interrupt_hooks: list[Callable[[], None]] = list(interrupt_hooks or []) # -- external controls ------------------------------------------------------ def request_interrupt(self) -> None: + """Stop the turn as soon as possible, from ANY state: mid-stream (the producer + thread drops the stream between chunks), mid-tool (interrupt hooks kill the + running command), awaiting an approval/question/plan (the await resolves as + interrupted), or between iterations (the loop checkpoint). Every pending + tool_call still gets a tool-error result so the history never carries orphans + (hosted templates reject them, and durable-resume would re-prompt them).""" self._cancel.set() + for hook in self._interrupt_hooks: + try: + hook() + except Exception: + pass # best-effort: a dead executor must not block the stop + + async def _interruptible(self, coro: Any, interrupted: Any) -> Any: + """Await `coro`, but resolve early with `interrupted` if the user stops the + turn. The pending task is cancelled so an answered-later Inbox card no-ops.""" + task = asyncio.ensure_future(coro) + cancel_wait = asyncio.ensure_future(self._cancel.wait()) + try: + done, _ = await asyncio.wait( + {task, cancel_wait}, return_when=asyncio.FIRST_COMPLETED + ) + if task in done: + return task.result() + task.cancel() + return interrupted + finally: + cancel_wait.cancel() def queue_steering( self, text: str, source: Optional[dict[str, Any]] = None @@ -202,9 +233,11 @@ class TurnEngine: iterations += 1 turn: Optional[AssistantTurn] = None + streamed: list[str] = [] try: async for chunk in self._astream(): if chunk.text_delta: + streamed.append(chunk.text_delta) yield Event( EventType.ASSISTANT_DELTA, {"text": chunk.text_delta} ) @@ -220,6 +253,16 @@ class TurnEngine: payload["raw"] = str(exc) yield Event(EventType.ERROR, payload) return + if self._cancel.is_set() and turn is None: + # Stopped mid-stream: persist exactly what the user watched arrive — + # the partial text, NO tool calls (any half-formed calls would either + # orphan or execute against the user's explicit stop). + if streamed: + self.messages.append( + _assistant_message(AssistantTurn(text="".join(streamed))) + ) + yield Event(EventType.INTERRUPTED, {"iterations": iterations}) + return if turn is None: turn = AssistantTurn() @@ -269,6 +312,10 @@ class TurnEngine: for chunk in provider.stream( model=model, messages=messages, tools=tools, **settings ): + # User pressed Stop: drop the stream between chunks (reading the + # asyncio.Event's flag from a thread is safe; we only read). + if self._cancel.is_set(): + break loop.call_soon_threadsafe(queue.put_nowait, ("chunk", chunk)) except Exception as exc: # surfaced to the awaiting consumer loop.call_soon_threadsafe(queue.put_nowait, ("error", exc)) @@ -277,7 +324,18 @@ class TurnEngine: loop.run_in_executor(None, produce) while True: - kind, payload = await queue.get() + # Race the queue against Stop so a stalled stream (no chunks arriving — + # the pre-first-token wait, a wedged connection) can't hold the turn. + get_task = asyncio.ensure_future(queue.get()) + cancel_task = asyncio.ensure_future(self._cancel.wait()) + done, _ = await asyncio.wait( + {get_task, cancel_task}, return_when=asyncio.FIRST_COMPLETED + ) + cancel_task.cancel() + if get_task not in done: + get_task.cancel() + return # interrupted — the producer exits on its own next chunk + kind, payload = get_task.result() if kind == "chunk": yield payload elif kind == "error": @@ -293,6 +351,10 @@ class TurnEngine: run concurrently; everything else runs one at a time in call order.""" cleared: list[ToolCall] = [] for tool_call in tool_calls: + if self._cancel.is_set(): + # Stopped: every remaining call still gets an answer (no orphans). + yield self._interrupted_tool(tool_call) + continue yield Event( EventType.TOOL_PROPOSED, {"name": tool_call.name, "arguments": tool_call.arguments}, @@ -340,11 +402,27 @@ class TurnEngine: yield self._record_result(tool_call, result, status) for tool_call in serial: + if self._cancel.is_set(): + yield self._interrupted_tool(tool_call) + continue yield Event(EventType.TOOL_STARTED, {"name": tool_call.name}) self._audit(tool_call, stage="started") result, status = await asyncio.to_thread(self._execute_sync, tool_call) yield self._record_result(tool_call, result, status) + def _interrupted_tool(self, tool_call: ToolCall) -> Event: + """The stop-path answer for a call that will not run: a tool-error result in the + history (hosted chat templates reject orphaned tool_calls, and durable-resume + would otherwise re-prompt it) + the finished event for the tool card.""" + self.messages.append(_tool_error_message(tool_call, "interrupted by user")) + self._audit( + tool_call, stage="finished", status="interrupted", reason="user stop" + ) + return Event( + EventType.TOOL_FINISHED, + {"name": tool_call.name, "status": "interrupted", "reason": "stopped"}, + ) + def _parallel_safe(self, tool_call: ToolCall) -> bool: # Only metadata-declared low-risk tools (reads, searches, git queries) run # concurrently; writes, shell, and anything unannotated stay strictly ordered. @@ -398,17 +476,23 @@ class TurnEngine: }, ) self._audit(tool_call, stage="approval_requested", reason=decision.reason) - outcome = await self.approver( - PermissionRequest( - tool_name=tool_call.name, - arguments=tool_call.arguments, - metadata=metadata, - reason=decision.reason, - tool_call_id=tool_call.id, - ) + outcome = await self._interruptible( + self.approver( + PermissionRequest( + tool_name=tool_call.name, + arguments=tool_call.arguments, + metadata=metadata, + reason=decision.reason, + tool_call_id=tool_call.id, + ) + ), + interrupted=ApprovalOutcome.DENY, ) if outcome is ApprovalOutcome.DENY: - allowed, reason = False, "denied by user" + allowed, reason = ( + False, + "interrupted by user" if self._cancel.is_set() else "denied by user", + ) self._audit( tool_call, stage="approval_resolved", @@ -554,7 +638,10 @@ class TurnEngine: else: yield Event(EventType.PLAN_PROPOSED, {"plan": plan}) self._audit(tool_call, stage="plan_proposed") - result = await self.plan_approver(dict(args), tool_call.id) or { + result = await self._interruptible( + self.plan_approver(dict(args), tool_call.id), + interrupted={"approved": False, "error": "interrupted by user"}, + ) or { "approved": False, "error": "no response", } @@ -615,7 +702,10 @@ class TurnEngine: stage="directory_requested", reason=str(args.get("reason", "")), ) - result = await self.directory_requester(dict(args), tool_call.id) or { + result = await self._interruptible( + self.directory_requester(dict(args), tool_call.id), + interrupted={"granted": False, "error": "interrupted by user"}, + ) or { "granted": False, "error": "no response", } @@ -656,7 +746,10 @@ class TurnEngine: # The asker is mode-aware (attended → live inline prompt; unattended → Inbox), so it # owns surfacing the question. The engine just awaits the answer. self._audit(tool_call, stage="question_requested", reason=question) - result = await self.question_asker(dict(args), tool_call.id) or { + result = await self._interruptible( + self.question_asker(dict(args), tool_call.id), + interrupted={"answer": "", "error": "interrupted by user"}, + ) or { "answer": "", "error": "no response", } diff --git a/coworker/server/manager.py b/coworker/server/manager.py index b54086f8..fd7b7622 100644 --- a/coworker/server/manager.py +++ b/coworker/server/manager.py @@ -3276,7 +3276,9 @@ class SessionManager: engine = self._engines.pop(session_id, None) if engine is not None: try: - engine.interrupt() + # (was engine.interrupt() — a method that never existed; the AttributeError + # was silently swallowed, so deleting a running session never stopped it.) + engine.request_interrupt() except Exception: pass record = self.session_store.load(session_id) diff --git a/coworker/tools/shell.py b/coworker/tools/shell.py index 5fe6afc0..b30ea78f 100644 --- a/coworker/tools/shell.py +++ b/coworker/tools/shell.py @@ -151,6 +151,9 @@ class LocalExecutor(Executor): self._is_windows = _IS_WINDOWS self._bg_tasks: dict[str, _BackgroundTask] = {} self._bg_counter = 0 + # Set by interrupt_now() (user Stop) — run()'s read loop treats it like an + # early deadline, so the in-flight foreground command dies within one tick. + self._abort = threading.Event() # Pick a native shell per-OS. POSIX drives bash line-by-line; Windows drives # PowerShell in `-Command -` mode, which is a true stdin REPL (executes @@ -224,6 +227,7 @@ class LocalExecutor(Executor): ) timeout = timeout or self.default_timeout + self._abort.clear() # Run the command, then emit a marker line with exit code + cwd. self._proc.stdin.write(command + "\n") self._proc.stdin.write(self._trailer()) @@ -232,10 +236,16 @@ class LocalExecutor(Executor): deadline = time.monotonic() + timeout interrupted = False timed_out = False + aborted = False exit_code: Optional[int] = None lines: list[str] = [] while True: + if self._abort.is_set(): + # User Stop: reuse the deadline path this tick (interrupt-and-resync on + # POSIX, decisive shell kill on Windows) instead of waiting out the timeout. + aborted = True + deadline = time.monotonic() remaining = deadline - time.monotonic() if remaining <= 0: if self._is_windows: @@ -279,9 +289,20 @@ class LocalExecutor(Executor): # Keep the TAIL: builds and test runners put the verdict at the end. output = output[-self.max_output_chars :] return self._result( - command, exit_code, output, timed_out=timed_out, truncated=truncated + command, + exit_code, + output, + timed_out=timed_out, + truncated=truncated, + error="interrupted by user" if aborted else None, ) + def interrupt_now(self) -> None: + """User Stop: make an in-flight foreground `run()` bail on its next read tick + (≤0.5s). Thread-safe; a no-op when nothing is running. Background tasks are + left alone — they're explicitly fire-and-forget.""" + self._abort.set() + # -- background tasks --------------------------------------------------------- def run_background(self, command: str) -> dict[str, Any]: self._bg_counter += 1 diff --git a/tests/test_engine_stop.py b/tests/test_engine_stop.py new file mode 100644 index 00000000..85b24043 --- /dev/null +++ b/tests/test_engine_stop.py @@ -0,0 +1,174 @@ +"""Stop-button semantics: interrupt must bite in EVERY engine state, not just the +between-iterations checkpoint (v0.1.4 shipped with that as the only one — ledgered +2026-07-21). History invariant throughout: every tool_call gets a tool result, since +hosted chat templates reject orphans and durable-resume re-prompts them.""" + +from __future__ import annotations + +import asyncio +import time + +from coworker.engine import ApprovalOutcome, TurnEngine +from coworker.events import EventType +from coworker.permissions import PermissionEngine +from coworker.providers import ( + AssistantTurn, + ModelCapabilities, + ProviderClient, + StreamChunk, + ToolCall, +) +from coworker.tools import ToolRegistry + + +class EndlessStreamProvider(ProviderClient): + """Streams deltas ~forever (bounded so a regression fails instead of hanging).""" + + def __init__(self): + self.chunks_produced = 0 + + def complete(self, **kwargs): # pragma: no cover + raise NotImplementedError + + def capabilities(self, model): + return ModelCapabilities() + + def stream(self, *, model, messages, tools=None, **settings): + for i in range(200): + self.chunks_produced += 1 + yield StreamChunk(text_delta=f"w{i} ") + time.sleep(0.01) + yield StreamChunk(turn=AssistantTurn(text="full", finish_reason="stop")) + + +def _tool_turn(calls): + return AssistantTurn( + tool_calls=[ToolCall(id=f"c{i}", name=n, arguments=a) for i, (n, a) in enumerate(calls)], + finish_reason="tool_calls", + ) + + +class OneTurnProvider(ProviderClient): + def __init__(self, turn): + self._turn = turn + self.calls = 0 + + def complete(self, **kwargs): + self.calls += 1 + return self._turn + + def capabilities(self, model): + return ModelCapabilities() + + +def _tool_results(engine): + return [m for m in engine.messages if m.get("role") == "tool"] + + +def test_stop_mid_stream_keeps_partial_text(tmp_path): + provider = EndlessStreamProvider() + engine = TurnEngine( + provider=provider, + registry=ToolRegistry(), + permissions=PermissionEngine(workspace_root=tmp_path), + model="gpt-5.5", + ) + + async def run(): + events = [] + async for ev in engine.run("go"): + events.append(ev) + if ev.type == EventType.ASSISTANT_DELTA and len(events) > 3: + engine.request_interrupt() + return events + + events = asyncio.run(run()) + assert events[-1].type == EventType.INTERRUPTED + # Far fewer than the full 200 chunks were consumed… + assert provider.chunks_produced < 100 + # …and the partial text the user watched is persisted, with no tool calls. + last = engine.messages[-1] + assert last["role"] == "assistant" and last["content"].startswith("w0 ") + assert "tool_calls" not in last + + +def test_stop_while_awaiting_approval(tmp_path): + async def never_answers(_req): + await asyncio.Event().wait() # a pending approval card nobody answers + + registry = ToolRegistry() + + def write_file(path: str, content: str): # pragma: no cover — never approved + raise AssertionError("executed while awaiting approval") + + registry.register(write_file) + engine = TurnEngine( + provider=OneTurnProvider(_tool_turn([("write_file", {"path": "x", "content": "y"})])), + registry=registry, + permissions=PermissionEngine(workspace_root=tmp_path), + model="gpt-5.5", + approver=never_answers, + ) + + async def run(): + events = [] + async for ev in engine.run("go"): + events.append(ev) + if ev.type == EventType.PERMISSION_REQUIRED: + engine.request_interrupt() + return events + + events = asyncio.run(run()) + assert events[-1].type == EventType.INTERRUPTED + (result,) = _tool_results(engine) + assert "interrupted by user" in result["content"] + + +def test_stop_skips_remaining_tool_calls(tmp_path): + registry = ToolRegistry() + holder = {} + + def first_tool(): + """Runs, then the user hits Stop while it holds the turn.""" + holder["engine"].request_interrupt() + return {"ok": True} + + def second_tool(): # pragma: no cover — must never run + raise AssertionError("second tool executed after stop") + + registry.register(first_tool) + registry.register(second_tool) + + async def approve(_req): + return ApprovalOutcome.ONCE + + engine = TurnEngine( + provider=OneTurnProvider(_tool_turn([("first_tool", {}), ("second_tool", {})])), + registry=registry, + permissions=PermissionEngine(workspace_root=tmp_path), + model="gpt-5.5", + approver=approve, + ) + holder["engine"] = engine + + async def run(): + return [ev async for ev in engine.run("go")] + + events = asyncio.run(run()) + assert events[-1].type == EventType.INTERRUPTED + results = _tool_results(engine) + assert len(results) == 2 # both calls answered — no orphans + assert "interrupted by user" in results[1]["content"] + + +def test_interrupt_hook_fires(tmp_path): + fired = [] + engine = TurnEngine( + provider=OneTurnProvider(AssistantTurn(text="hi", finish_reason="stop")), + registry=ToolRegistry(), + permissions=PermissionEngine(workspace_root=tmp_path), + model="gpt-5.5", + interrupt_hooks=[lambda: fired.append(True)], + ) + engine.request_interrupt() + assert fired == [True]