feat(chat): stop an in-flight Author or Plan and Execute run

PLUGIN_MCP.md

@@ -327,6 +327,11 @@ What you must return:
 The orchestrator reads these into the in-memory run cache; the live
 `/plugin/mcp/status` endpoint surfaces them to the UI.
 
+The user can stop a run from the chat header, which cancels the asyncio task
+your `run()` is executing in. Catch `Exception`, never `BaseException`, and put
+subprocess and session teardown in `async with` blocks so a cancel unwinds
+cleanly. Swallowing `CancelledError` keeps the run alive with no way to stop it.
+
 ### 4.4 Reusing the framework's helpers
 
 Don't reinvent these. Import from the MCP plugin's namespace:

app/mcp_api.py

@@ -77,6 +77,15 @@
 _VALID_PROVIDERS = {"openai_compatible", "ollama"}
 
 
+def _stop_reason(tags) -> str:
+    """Why a KILLED run stopped: "user", "orphaned", or "" for anything else."""
+    if tags.get("mcp.cancelled") == "user":
+        return "user"
+    if tags.get("mcp.reconciled") == "orphaned":
+        return "orphaned"
+    return ""
+
+
 def _is_secret(name) -> bool:
     return bool(_SECRET_KEY_NAME.search(str(name))) and not str(name).endswith("_env")
 
@@ -337,6 +346,35 @@ async def features(self, request):
             self.log.error(f"[MCP] Error listing feature catalog: {e}")
             return web.json_response({"error": str(e)}, status=500)
 
+    async def cancel(self, request):
+        """POST /plugin/mcp/cancel {"run_id": ...}. Stop one in-flight run.
+
+        Always answers 200 for a well-formed body: pressing Stop twice, or
+        after the run already finished, is a no-op reported as
+        cancelling=false. The caller keeps polling /status to see the run
+        reach KILLED, which takes as long as the workflow needs to unwind.
+        """
+        try:
+            data = await request.json()
+        except Exception:
+            return web.json_response({"error": "Body must be JSON"}, status=400)
+        if not isinstance(data, dict):
+            return web.json_response({"error": "Body must be a JSON object"}, status=400)
+
+        run_id = data.get("run_id")
+        # Typed, not just present: run_id is used as a dict key, and an
+        # unhashable one would answer a bad request with a 500.
+        if not run_id or not isinstance(run_id, str):
+            return web.json_response({"error": "Missing run_id"}, status=400)
+
+        cancelling = self.mcp_svc.cancel_run(run_id)
+        snapshot = self.mcp_svc.get_run(run_id) or {}
+        return web.json_response({
+            "run_id": run_id,
+            "cancelling": cancelling,
+            "status": snapshot.get("status", "UNKNOWN"),
+        })
+
     async def status(self, request):
         """Live status of an in-flight or recently-finished run.
 
@@ -476,6 +514,9 @@ async def list_runs(self, request):
                             run_data.tags.get("process_result_summary")
                             or run_data.tags.get("process_result", "")
                         ),
+                        # Both a user stop and the boot sweep land on KILLED;
+                        # only the tags say which, and they read differently.
+                        "stop_reason": _stop_reason(run_data.tags),
                     }
                     all_runs.append(run_record)
 
@@ -516,6 +557,7 @@ async def get_run_detail(self, request):
             response = {
                 "run_id": run_id,
                 "status": run.info.status,
+                "stop_reason": _stop_reason(run.data.tags),
                 "start_time": run.info.start_time,
                 "end_time": run.info.end_time,
                 "run_name": run.data.tags.get("mlflow.runName", "Unnamed Run"),

app/mcp_svc.py

@@ -4,6 +4,7 @@
 import os
 import re
 import tempfile
+import traceback
 from pathlib import Path
 from app.utility.base_service import BaseService
 import asyncio
@@ -12,7 +13,9 @@
 from plugins.mcp.app.mlflow_run import (
     RunTracker,
     reconcile_orphaned_runs as _reconcile_orphaned_runs,
+    register_stage_observer,
     summarize_exception,
+    unregister_stage_observer,
 )
 
 
@@ -67,6 +70,15 @@ def _scrub_secrets(text: str) -> str:
 _SESSION_CACHE_LIMIT = 64
 _SESSION_TURN_CAP = 8
 
+# Shown on a run that was cancelled. Cancelling unwinds our own task; it
+# cannot undo the CALDERA writes the agent already made.
+_CANCELLED_BY_USER = (
+    "Anything already created in CALDERA stays, including a running operation."
+)
+_CANCELLED_BY_SERVER = (
+    "The server shut down. Anything already created in CALDERA stays."
+)
+
 
 class MCPService(BaseService):
     def __init__(self, services, server_registry=None, workflow_registry=None, capability_registry=None):
@@ -88,6 +100,16 @@ def __init__(self, services, server_registry=None, workflow_registry=None, capab
         # OrderedDict gives us O(1) LRU eviction without an extra dependency.
         self._runs: "collections.OrderedDict[str, dict]" = collections.OrderedDict()
 
+        # Live task handles, so a cancel request can name a running workflow.
+        # Kept out of self._runs, which evicts at _RUN_CACHE_LIMIT while long
+        # runs are still executing.
+        self._tasks: dict[str, asyncio.Task] = {}
+
+        # Runs we asked to stop. anyio can repackage a cancel that lands during
+        # MCP session startup as an ExceptionGroup, so the intent is recorded
+        # here rather than inferred from the exception type.
+        self._cancelling: set[str] = set()
+
         # Per-session chat history for workflows whose registration sets
         # supports_chat_history=True. session_id -> list of
         #   {"prompt": str, "response": str}
@@ -119,6 +141,77 @@ def _record_run(self, run_id: str, snapshot: dict) -> None:
         while len(self._runs) > _RUN_CACHE_LIMIT:
             self._runs.popitem(last=False)
 
+    def _update_stage(self, run_id: str, stage: str) -> None:
+        """Mirror a workflow's stage tag into the polled snapshot.
+
+        Workflows report progress with tracker.set_tag("stage", ...), which
+        reaches MLflow only. /status serves this cache, so without the
+        mirror the chat UI shows "initializing" for the whole run. Guarded
+        on RUNNING so a late tag cannot overwrite a terminal stage.
+        """
+        snapshot = self._runs.get(run_id)
+        if snapshot is not None and snapshot.get("status") == "RUNNING":
+            snapshot["stage"] = stage
+
+    def cancel_run(self, run_id: str) -> bool:
+        """Ask the task owning ``run_id`` to stop.
+
+        False when no live task owns it: already finished, never started, or
+        from a previous process. Cancelling is by run id, so one run stopping
+        never touches another, and a repeat press is a no-op.
+        """
+        task = self._tasks.get(run_id)
+        if task is None or task.done():
+            return False
+        self._cancelling.add(run_id)
+        task.cancel()
+        return True
+
+    def _on_task_done(self, task: asyncio.Task, run_id: str, workflow_id: str,
+                      session_id: str, prompt: str) -> None:
+        """Drop a finished run's handle, and close out one that never ran.
+
+        A task cancelled before its body was scheduled leaves no cache
+        snapshot and an MLflow run stuck on RUNNING, which is the state this
+        whole feature exists to avoid. Runs after the task, so it also clears
+        a cancel that raced the run's last await.
+        """
+        self._tasks.pop(run_id, None)
+        if task.cancelled() and self.get_run(run_id) is None:
+            tracker = RunTracker(run_id)
+            self._record_cancelled_run(
+                run_id, tracker, workflow_id, session_id, prompt
+            )
+            tracker.terminate("KILLED")
+        self._cancelling.discard(run_id)
+
+    def _record_cancelled_run(self, run_id: str, tracker, workflow_id: str,
+                              session_id: str, prompt: str) -> None:
+        """Mark a cancelled run in both the cache and MLflow.
+
+        Only a cancel this service asked for is the user's. Everything else
+        is the event loop tearing down in-flight tasks at shutdown, and
+        History must not report that as somebody pressing Stop.
+        """
+        by_user = run_id in self._cancelling
+        self.log.info(
+            f"[MCP] Run {run_id} cancelled by {'user' if by_user else 'server'}"
+        )
+        tracker.set_tag("stage", "stopped by user" if by_user else "cancelled")
+        tracker.set_tag("status", "cancelled")
+        tracker.set_tag("mcp.cancelled", "user" if by_user else "server")
+        self._record_run(run_id, {
+            "status": "KILLED",
+            "stage": "stopped by user" if by_user else "cancelled",
+            "workflow_id": workflow_id,
+            "session_id": session_id,
+            "prompt": prompt,
+            "process_result": "",
+            "reasoning": "",
+            "trajectory": {},
+            "error": _CANCELLED_BY_USER if by_user else _CANCELLED_BY_SERVER,
+        })
+
     async def reconcile_orphaned_runs(self) -> list:
         """Terminate MLflow runs left RUNNING by a process that died.
 
@@ -396,7 +489,7 @@ async def execute(self, focus: str = None, prompt: str = "", model_config: dict
         # MLflow can group their (single-turn) sessions consistently.
         resolved_session_id = session_id or run_id
 
-        asyncio.create_task(self._run_execution(
+        task = asyncio.create_task(self._run_execution(
             workflow=workflow,
             prompt=prompt,
             run_id=run_id,
@@ -408,6 +501,10 @@ async def execute(self, focus: str = None, prompt: str = "", model_config: dict
             workflow_context=workflow_context if isinstance(workflow_context, dict) else {},
             disable_history=disable_history,
         ))
+        self._tasks[run_id] = task
+        task.add_done_callback(lambda t, rid=run_id: self._on_task_done(
+            t, rid, workflow.id, resolved_session_id, prompt
+        ))
         return {
             "run_id": run_id,
             "session_id": resolved_session_id,
@@ -469,6 +566,10 @@ async def _run_execution(self, workflow, prompt, run_id, lm_obj,
         tracker = RunTracker(run_id)
         terminal_status = "FAILED"
         try:
+            # Every tracker bound to this run mirrors its stage tags into
+            # the snapshot /status serves, including the one the workflow
+            # binds. Inside the try so the finally always releases it.
+            register_stage_observer(run_id, lambda v: self._update_stage(run_id, v))
             tracker.set_tag("stage", "initializing")
             tracker.set_tag("workflow_id", workflow.id)
             tracker.set_tag("mcp.session_id", effective_session_id)
@@ -567,12 +668,42 @@ async def _run_execution(self, workflow, prompt, run_id, lm_obj,
                     effective_session_id, prompt, process_result
                 )
 
-        except Exception as e:
+        except asyncio.CancelledError:
+            terminal_status = "KILLED"
+            self._record_cancelled_run(
+                run_id, tracker, workflow.id, effective_session_id, prompt
+            )
+            raise
+
+        except BaseException as e:
+            if run_id in self._cancelling:
+                # anyio hands a cancel that lands during MCP session startup
+                # back wrapped in an exception group, so the type proves
+                # nothing; the recorded intent does. Re-raising as cancelled
+                # also stops asyncio logging the wrapper as unretrieved.
+                terminal_status = "KILLED"
+                self._record_cancelled_run(
+                    run_id, tracker, workflow.id, effective_session_id, prompt
+                )
+                raise asyncio.CancelledError from None
+            if not isinstance(e, Exception):
+                # KeyboardInterrupt and friends are not this run's failure.
+                raise
+
             error_summary = summarize_exception(e)
-            self.log.error(f"[MCP] Execution failed: {error_summary}")
+            # exc_info because the workflow no longer prints the stack and
+            # the traceback param below is best-effort: a dead tracking
+            # store would otherwise leave one summary line and nothing else.
+            self.log.error(f"[MCP] Execution failed: {error_summary}", exc_info=e)
             tracker.set_tag("stage", "error")
             tracker.set_tag("status", "error")
             tracker.log_param("error", error_summary)
+            # Written here rather than in the workflow: params are immutable,
+            # and this is the only layer that knows a cancel from a crash.
+            tracker.log_param(
+                "traceback",
+                "".join(traceback.format_exception(type(e), e, e.__traceback__)),
+            )
             self._record_run(run_id, {
                 "status": "FAILED",
                 "stage": "error",
@@ -586,4 +717,5 @@ async def _run_execution(self, workflow, prompt, run_id, lm_obj,
             })
 
         finally:
+            unregister_stage_observer(run_id)
             tracker.terminate(terminal_status)

app/mlflow_run.py

@@ -35,6 +35,33 @@ def _leaf_exceptions(exc: BaseException) -> list[BaseException]:
     return [exc]
 
 
+# Stage observers by run id. A workflow binds its own RunTracker to the run
+# the service started, so the mirror has to follow the run rather than a
+# tracker instance -- otherwise every stage a workflow reports is invisible
+# to /status. Keyed by run id, so concurrent runs never collide.
+_stage_observers = {}
+
+
+def register_stage_observer(run_id: str, fn) -> None:
+    """Call ``fn(stage)`` whenever a stage tag is written for ``run_id``."""
+    _stage_observers[run_id] = fn
+
+
+def unregister_stage_observer(run_id: str) -> None:
+    _stage_observers.pop(run_id, None)
+
+
+def _notify_stage(run_id: str, value: str) -> None:
+    fn = _stage_observers.get(run_id)
+    if fn is None:
+        return
+    # Progress reporting never fails the run it describes.
+    try:
+        fn(value)
+    except Exception as e:
+        log.debug(f"[MCP] Stage mirror failed for run {run_id}: {e}")
+
+
 def summarize_exception(exc: BaseException) -> str:
     """The innermost cause, scrubbed, rather than the group that wraps it.
 
@@ -141,6 +168,10 @@ def terminated(self) -> bool:
         return self._terminated
 
     def set_tag(self, key: str, value) -> None:
+        # Mirror first: a tracking server that is down must not cost the
+        # chat UI its progress display.
+        if key == "stage":
+            _notify_stage(self.run_id, str(value))
         self._write(self._client.set_tag, key, value, _MAX_TAG_CHARS)
 
     def log_param(self, key: str, value) -> None:

app/workflows/author.py

@@ -6,10 +6,13 @@
 import sys
 import mlflow
 import traceback
+import logging
 from contextlib import AsyncExitStack
 
 from plugins.mcp.app.config import caldera_connection, llm_defaults, mlflow_settings
 from plugins.mcp.app.mlflow_run import RunTracker, summarize_exception
+
+log = logging.getLogger("plugins.mcp")
 from plugins.mcp.app.dspy_env import (
     ENV_API_BASE,
     ENV_API_KEY,
@@ -332,14 +335,21 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, run_
 
     except Exception as e:
         tb = traceback.format_exc()
+        # A run the orchestrator handed us is the orchestrator's to classify.
+        # A cancel landing during MCP session teardown arrives here as an
+        # anyio group carrying BrokenResourceError and no CancelledError
+        # leaf, so the type proves nothing; writing failure state would
+        # stamp an immutable error param on a run the user simply stopped.
+        if not created_local_run:
+            log.debug(f"[MCP] Raising to the orchestrator:\n{tb}")
+            raise
         print("[MCP] Exception occurred:")
         print(tb)
         tracker.set_tag("status", "failed")
         tracker.set_tag("stage", "error")
         tracker.log_param("error", summarize_exception(e))
         tracker.log_param("traceback", tb)
-        if created_local_run:
-            tracker.terminate("FAILED")
+        tracker.terminate("FAILED")
         raise