feat(agent): support user-initiated cancel for in-flight agent runs

agent/protocol/__init__.py

@@ -3,6 +3,11 @@ from .agent_stream import AgentStreamExecutor
 from .task import Task, TaskType, TaskStatus
 from .result import AgentResult, AgentAction, AgentActionType, ToolResult
 from .models import LLMModel, LLMRequest, ModelFactory
+from .cancel import (
+    AgentCancelledError,
+    CancelTokenRegistry,
+    get_cancel_registry,
+)
 
 __all__ = [
     'Agent', 
@@ -16,5 +21,8 @@ __all__ = [
     'ToolResult',
     'LLMModel',
     'LLMRequest', 
-    'ModelFactory'
-]
+    'ModelFactory',
+    'AgentCancelledError',
+    'CancelTokenRegistry',
+    'get_cancel_registry',
+]

agent/protocol/agent.py

@@ -365,7 +365,8 @@ class Agent:
 
         return action
 
-    def run_stream(self, user_message: str, on_event=None, clear_history: bool = False, skill_filter=None) -> str:
+    def run_stream(self, user_message: str, on_event=None, clear_history: bool = False,
+                   skill_filter=None, cancel_event=None) -> str:
         """
         Execute single agent task with streaming (based on tool-call)
 
@@ -374,6 +375,7 @@ class Agent:
         - Multi-turn reasoning based on tool-call
         - Event callbacks
         - Persistent conversation history across calls
+        - User-initiated cancellation via ``cancel_event``
 
         Args:
             user_message: User message
@@ -381,6 +383,11 @@ class Agent:
                      event = {"type": str, "timestamp": float, "data": dict}
             clear_history: If True, clear conversation history before this call (default: False)
             skill_filter: Optional list of skill names to include in this run
+            cancel_event: Optional threading.Event polled at agent checkpoints.
+                When set, the loop exits at the next safe point, injects a
+                "[Interrupted by user]" assistant note, and returns the
+                partial response. ``messages`` stays in a valid state
+                (tool_use/tool_result pairs preserved).
 
         Returns:
             Final response text
@@ -424,7 +431,8 @@ class Agent:
             max_turns=self.max_steps,
             on_event=on_event,
             messages=messages_copy,  # Pass copied message history
-            max_context_turns=max_context_turns
+            max_context_turns=max_context_turns,
+            cancel_event=cancel_event,
         )
 
         # Execute

agent/protocol/agent_stream.py

@@ -7,6 +7,7 @@ import json
 import time
 from typing import List, Dict, Any, Optional, Callable, Tuple
 
+from agent.protocol.cancel import AgentCancelledError
 from agent.protocol.models import LLMRequest, LLMModel
 from agent.protocol.message_utils import sanitize_claude_messages, compress_turn_to_text_only
 from agent.tools.base_tool import BaseTool, ToolResult
@@ -64,7 +65,8 @@ class AgentStreamExecutor:
             max_turns: int = 50,
             on_event: Optional[Callable] = None,
             messages: Optional[List[Dict]] = None,
-            max_context_turns: int = 30
+            max_context_turns: int = 30,
+            cancel_event=None,
     ):
         """
         Initialize stream executor
@@ -78,6 +80,10 @@ class AgentStreamExecutor:
             on_event: Event callback function
             messages: Optional existing message history (for persistent conversations)
             max_context_turns: Maximum number of conversation turns to keep in context
+            cancel_event: Optional threading.Event used to signal user cancel.
+                Checked at every safe point (turn boundary, before tool execution,
+                during LLM streaming). When set, raises AgentCancelledError which
+                run_stream catches to gracefully wind down.
         """
         self.agent = agent
         self.model = model
@@ -87,6 +93,7 @@ class AgentStreamExecutor:
         self.max_turns = max_turns
         self.on_event = on_event
         self.max_context_turns = max_context_turns
+        self.cancel_event = cancel_event
 
         # Message history - use provided messages or create new list
         self.messages = messages if messages is not None else []
@@ -97,6 +104,73 @@ class AgentStreamExecutor:
         # Track files to send (populated by read tool)
         self.files_to_send = []  # List of file metadata dicts
 
+    def _check_cancelled(self) -> None:
+        """Raise AgentCancelledError if the user requested cancellation.
+
+        Called at safe points (turn start, between tool calls, between LLM
+        chunks). Cheap to call: just an Event.is_set() probe.
+        """
+        if self.cancel_event is not None and self.cancel_event.is_set():
+            raise AgentCancelledError("agent cancelled by user")
+
+    def _handle_cancelled(self, partial_response: str) -> None:
+        """Wind down ``self.messages`` after a user-initiated cancel.
+
+        The messages list may be in any of these states when we get here:
+          (a) Last message is an assistant message containing tool_use
+              blocks but the matching tool_result has not been appended yet.
+          (b) Last message is an assistant text-only reply (cancel happened
+              right before the next turn started).
+          (c) Last message is a user tool_result message and we cancelled
+              between turns.
+
+        For (a) we MUST synthesise tool_result blocks, otherwise the next
+        request will fail Claude/OpenAI's strict pairing validation. For
+        (b)/(c) the state is already valid and we just append a small
+        cancellation note so the user/LLM both see the boundary clearly.
+        """
+        try:
+            # Step 1: close any orphaned tool_use in the trailing assistant
+            # message by injecting matching tool_result blocks.
+            if self.messages and isinstance(self.messages[-1], dict) \
+                    and self.messages[-1].get("role") == "assistant":
+                last = self.messages[-1]
+                content = last.get("content")
+                if isinstance(content, list):
+                    pending_tool_use_ids = [
+                        block.get("id")
+                        for block in content
+                        if isinstance(block, dict) and block.get("type") == "tool_use"
+                    ]
+                    pending_tool_use_ids = [tid for tid in pending_tool_use_ids if tid]
+                    if pending_tool_use_ids:
+                        tool_result_blocks = [
+                            {
+                                "type": "tool_result",
+                                "tool_use_id": tid,
+                                "content": "Cancelled by user before this tool finished.",
+                                "is_error": True,
+                            }
+                            for tid in pending_tool_use_ids
+                        ]
+                        self.messages.append({
+                            "role": "user",
+                            "content": tool_result_blocks,
+                        })
+                        logger.info(
+                            f"[Agent] Injected {len(tool_result_blocks)} cancellation "
+                            f"tool_result blocks to keep message history valid"
+                        )
+
+            # Step 2: append a stable "interrupted" marker so the LLM sees a
+            # clear stop boundary on the next turn.
+            self.messages.append({
+                "role": "assistant",
+                "content": [{"type": "text", "text": "_(Cancelled by user)_"}],
+            })
+        except Exception as e:
+            logger.warning(f"[Agent] _handle_cancelled cleanup failed: {e}")
+
     def _emit_event(self, event_type: str, data: dict = None):
         """Emit event"""
         if self.on_event:
@@ -270,8 +344,13 @@ class AgentStreamExecutor:
         final_response = ""
         turn = 0
 
+        cancelled = False
         try:
             while turn < self.max_turns:
+                # Check at the very top of every turn so a cancel arriving
+                # between turns short-circuits cleanly.
+                self._check_cancelled()
+
                 turn += 1
                 logger.info(f"[Agent] 第 {turn} 轮")
                 self._emit_event("turn_start", {"turn": turn})
@@ -375,6 +454,8 @@ class AgentStreamExecutor:
 
                 try:
                     for tool_call in tool_calls:
+                        # Honour cancel between tool invocations within the same turn
+                        self._check_cancelled()
                         result = self._execute_tool(tool_call)
                         tool_results.append(result)
                         
@@ -557,6 +638,15 @@ class AgentStreamExecutor:
                         self.messages.pop(prompt_insert_idx)
                         logger.debug("[Agent] Removed injected max-steps prompt from message history")
 
+        except AgentCancelledError:
+            # User-initiated stop: wind down message history cleanly so the
+            # next turn is unaffected; channels emit a "cancelled" UI event.
+            cancelled = True
+            logger.info(f"[Agent] 🛑 已被用户中止 (第 {turn} 轮)")
+            self._handle_cancelled(final_response)
+            if not final_response or not final_response.strip():
+                final_response = "_(Cancelled)_"
+
         except Exception as e:
             logger.error(f"❌ Agent执行错误: {e}")
             self._emit_event("error", {"error": str(e)})
@@ -564,8 +654,11 @@ class AgentStreamExecutor:
 
         finally:
             final_response = final_response.strip() if final_response else final_response
-            logger.info(f"[Agent] 🏁 完成 ({turn}轮)")
-            self._emit_event("agent_end", {"final_response": final_response})
+            if cancelled:
+                # Emit before agent_end so channels can mark UI as cancelled
+                self._emit_event("agent_cancelled", {"final_response": final_response})
+            logger.info(f"[Agent] 🏁 完成 ({turn}轮)" + (" [cancelled]" if cancelled else ""))
+            self._emit_event("agent_end", {"final_response": final_response, "cancelled": cancelled})
 
         return final_response
 
@@ -644,7 +737,32 @@ class AgentStreamExecutor:
         try:
             stream = self.model.call_stream(request)
 
+            # Probe cancel every N chunks to bound reaction time without
+            # checking on every token.
+            _cancel_probe_counter = 0
+            _CANCEL_PROBE_EVERY = 8
+
             for chunk in stream:
+                _cancel_probe_counter += 1
+                if _cancel_probe_counter >= _CANCEL_PROBE_EVERY:
+                    _cancel_probe_counter = 0
+                    if self.cancel_event is not None and self.cancel_event.is_set():
+                        # Persist partial text only; tool_use args may be
+                        # truncated mid-stream and would fail validation.
+                        logger.info("[Agent] cancel detected mid-stream, aborting LLM call")
+                        if full_content:
+                            partial_msg = {
+                                "role": "assistant",
+                                "content": [{"type": "text", "text": full_content}],
+                            }
+                            self.messages.append(partial_msg)
+                        self._emit_event("message_end", {
+                            "content": full_content,
+                            "tool_calls": [],
+                            "cancelled": True,
+                        })
+                        raise AgentCancelledError("cancelled during LLM streaming")
+
                 # Check for errors
                 if isinstance(chunk, dict) and chunk.get("error"):
                     # Extract error message from nested structure
@@ -738,6 +856,10 @@ class AgentStreamExecutor:
                     elif isinstance(choice, dict) and choice.get("_gemini_raw_parts"):
                         gemini_raw_parts = choice["_gemini_raw_parts"]
 
+        except AgentCancelledError:
+            # Must propagate untouched; never treat as a retryable error.
+            raise
+
         except Exception as e:
             error_str = str(e)
             error_str_lower = error_str.lower()

agent/protocol/cancel.py

@@ -0,0 +1,121 @@
+"""
+Cancel token registry for aborting in-flight agent runs.
+
+A user cancel (web Cancel button, /cancel command) sets a threading.Event
+that the agent loop polls at safe checkpoints. Tokens are keyed by
+request_id (preferred) and tracked under session_id as a fallback. Entries
+are released after the run completes to keep the registry bounded.
+
+No project deps — importable from any layer without circular imports.
+"""
+
+from __future__ import annotations
+
+import threading
+from typing import Dict, Optional
+
+
+class AgentCancelledError(Exception):
+    """Raised inside the agent loop when a stop has been requested.
+
+    The agent stream executor catches this, injects a "[Interrupted]" note
+    into the message history (preserving tool_use/tool_result integrity)
+    and returns a partial response to the caller.
+    """
+
+
+class _CancelEntry:
+    __slots__ = ("event", "session_id")
+
+    def __init__(self, session_id: Optional[str]):
+        self.event = threading.Event()
+        self.session_id = session_id
+
+
+class CancelTokenRegistry:
+    """In-process registry mapping request_id -> cancel Event.
+
+    Thread-safe. Singleton via module-level ``_registry``.
+    """
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._by_request: Dict[str, _CancelEntry] = {}
+        # session_id -> set of request_ids currently in flight (usually 1).
+        self._by_session: Dict[str, set] = {}
+
+    def register(self, request_id: str, session_id: Optional[str] = None) -> threading.Event:
+        """Create (or return existing) cancel event for a request.
+
+        Returns the threading.Event the caller should poll via ``is_set()``.
+        """
+        if not request_id:
+            return threading.Event()
+        with self._lock:
+            entry = self._by_request.get(request_id)
+            if entry is None:
+                entry = _CancelEntry(session_id)
+                self._by_request[request_id] = entry
+                if session_id:
+                    self._by_session.setdefault(session_id, set()).add(request_id)
+            return entry.event
+
+    def get_event(self, request_id: str) -> Optional[threading.Event]:
+        if not request_id:
+            return None
+        with self._lock:
+            entry = self._by_request.get(request_id)
+            return entry.event if entry else None
+
+    def cancel_request(self, request_id: str) -> bool:
+        """Trigger cancel for a specific request. Returns True when matched."""
+        if not request_id:
+            return False
+        with self._lock:
+            entry = self._by_request.get(request_id)
+        if entry is None:
+            return False
+        entry.event.set()
+        return True
+
+    def cancel_session(self, session_id: str) -> int:
+        """Trigger cancel for every in-flight request of a session.
+
+        Returns the number of requests cancelled (0 when nothing was running).
+        """
+        if not session_id:
+            return 0
+        with self._lock:
+            request_ids = list(self._by_session.get(session_id, ()))
+            entries = [self._by_request[r] for r in request_ids if r in self._by_request]
+        for entry in entries:
+            entry.event.set()
+        return len(entries)
+
+    def unregister(self, request_id: str) -> None:
+        """Remove an entry once the agent run is done. Safe to call twice."""
+        if not request_id:
+            return
+        with self._lock:
+            entry = self._by_request.pop(request_id, None)
+            if entry and entry.session_id:
+                bucket = self._by_session.get(entry.session_id)
+                if bucket is not None:
+                    bucket.discard(request_id)
+                    if not bucket:
+                        self._by_session.pop(entry.session_id, None)
+
+    def has_active(self, session_id: str) -> bool:
+        if not session_id:
+            return False
+        with self._lock:
+            bucket = self._by_session.get(session_id)
+            return bool(bucket)
+
+
+_registry = CancelTokenRegistry()
+
+
+def get_cancel_registry() -> CancelTokenRegistry:
+    """Module-level accessor for the singleton registry."""
+    return _registry