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..30ae224 100644 --- a/app/mcp_svc.py +++ b/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) diff --git a/app/mlflow_run.py b/app/mlflow_run.py index ccf1edb..8f7b242 100644 --- a/app/mlflow_run.py +++ b/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: diff --git a/app/workflows/author.py b/app/workflows/author.py index 4f7c352..e6d7894 100644 --- a/app/workflows/author.py +++ b/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 diff --git a/app/workflows/plan_execute.py b/app/workflows/plan_execute.py index 086e43a..f9c119d 100644 --- a/app/workflows/plan_execute.py +++ b/app/workflows/plan_execute.py @@ -6,12 +6,15 @@ import sys import mlflow import traceback +import logging from mlflow.tracking import MlflowClient import asyncio 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, @@ -343,14 +346,21 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, 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 # Optional streaming updates (if desired for parity) diff --git a/gui/views/chat/ChatMessage.vue b/gui/views/chat/ChatMessage.vue index d5cc4d6..3e42593 100644 --- a/gui/views/chat/ChatMessage.vue +++ b/gui/views/chat/ChatMessage.vue @@ -30,8 +30,14 @@
{{ message.errorMessage }}
- -{{ message.errorMessage || STOPPED_NOTE }}
+No result returned.
+Status: - {{ runDetail.status }} + {{ statusLabel(runDetail) }}