From e3cdf76472acf247d41fe789484fe52df55267b5 Mon Sep 17 00:00:00 2001 From: HackedRico Date: Tue, 1 Sep 2026 14:33:29 -0400 Subject: [PATCH 1/9] feat(chat): add a Stop button for in-flight runs Keep each run's task handle so POST /plugin/mcp/cancel can name one. The run reaches KILLED in the cache and in MLflow, which the poller treats as terminal and History labels as stopped. Covers the anyio exception-group shape and a cancel landing before the task body runs, both of which left the run RUNNING forever. Closes #34 --- PLUGIN_MCP.md | 5 + app/mcp_api.py | 42 ++++ app/mcp_svc.py | 110 +++++++++- gui/views/chat/ChatMessage.vue | 18 ++ gui/views/chat/ChatWorkflow.vue | 46 +++- gui/views/chat/composables/useMcpRun.js | 41 +++- gui/views/mcp_history.vue | 21 +- hook.py | 1 + tests/test_run_cancellation.py | 281 ++++++++++++++++++++++++ 9 files changed, 549 insertions(+), 16 deletions(-) create mode 100644 tests/test_run_cancellation.py diff --git a/PLUGIN_MCP.md b/PLUGIN_MCP.md index 88a0a87..4282ba6 100644 --- a/PLUGIN_MCP.md +++ b/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: diff --git a/app/mcp_api.py b/app/mcp_api.py index cbb087a..f108cb3 100644 --- a/app/mcp_api.py +++ b/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"), diff --git a/app/mcp_svc.py b/app/mcp_svc.py index 24bd70a..8e155e4 100644 --- a/app/mcp_svc.py +++ b/app/mcp_svc.py @@ -67,6 +67,18 @@ 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 = ( + "Stopped by user. Anything already created in CALDERA remains, including " + "an operation that was started, which may still be tasking agents. This " + "cannot be rolled back." +) +_CANCELLED_BY_SERVER = ( + "Stopped when the server shut down. Anything already created in CALDERA " + "remains." +) + 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,65 @@ def _record_run(self, run_id: str, snapshot: dict) -> None: while len(self._runs) > _RUN_CACHE_LIMIT: self._runs.popitem(last=False) + 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 +477,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 +489,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, @@ -567,7 +652,28 @@ 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}") tracker.set_tag("stage", "error") diff --git a/gui/views/chat/ChatMessage.vue b/gui/views/chat/ChatMessage.vue index d5cc4d6..9e94c1a 100644 --- a/gui/views/chat/ChatMessage.vue +++ b/gui/views/chat/ChatMessage.vue @@ -30,6 +30,12 @@

{{ message.errorMessage }}

+ +
+ Run stopped. +

{{ message.errorMessage || STOPPED_NOTE }}

+
+
false }, }) +// Fallback for a stopped bubble whose cache entry aged out before the page +// read the reason. +const STOPPED_NOTE = + 'Anything already created in CALDERA remains. This cannot be rolled back.' + const role = computed(() => props.message.role || 'assistant') const formattedResult = computed(() => formatProcessResult(props.message.finalResult || '')) @@ -141,6 +152,13 @@ const formattedTime = computed(() => { padding: 0.7rem 0.9rem; border-radius: 4px; } +.stopped-text { + color: #ffe0b8; + background-color: rgba(245, 158, 11, 0.08); + border-left: 3px solid #f59e0b; + padding: 0.7rem 0.9rem; + border-radius: 4px; +} .muted { color: #888888; font-style: italic; } .result-content { diff --git a/gui/views/chat/ChatWorkflow.vue b/gui/views/chat/ChatWorkflow.vue index 5760afc..0819b39 100644 --- a/gui/views/chat/ChatWorkflow.vue +++ b/gui/views/chat/ChatWorkflow.vue @@ -46,6 +46,17 @@ messages.filter(m => m.role === 'assistant').length === 1 ? '' : 's' }} +