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 }}

- -
+ +
+ Run stopped. +

{{ message.errorMessage || STOPPED_NOTE }}

+
+ + +
No result returned.

+ +

No result returned.

+ import { computed } from 'vue' import { formatProcessResult } from '../format_result.js' +import { isDeliveredResponse } from './messageState.js' import ChatLoadingState from './ChatLoadingState.vue' import ChatThoughts from './ChatThoughts.vue' @@ -68,7 +78,13 @@ const props = defineProps({ isInjectedSentence: { type: Function, default: () => 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 stays.' + const role = computed(() => props.message.role || 'assistant') +// Shared with the header's response count so the two cannot disagree. +const isDelivered = computed(() => isDeliveredResponse(props.message)) const formattedResult = computed(() => formatProcessResult(props.message.finalResult || '')) const formattedTime = computed(() => { @@ -141,6 +157,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..7bca982 100644 --- a/gui/views/chat/ChatWorkflow.vue +++ b/gui/views/chat/ChatWorkflow.vue @@ -34,17 +34,27 @@ Back - - Working - - - {{ messages.filter(m => m.role === 'assistant').length }} response{{ - messages.filter(m => m.role === 'assistant').length === 1 ? '' : 's' - }} + + {{ run.isStopping.value ? 'Stopping…' : 'Working' }} + + + + {{ responseCount }} response{{ responseCount === 1 ? '' : 's' }}
@@ -67,7 +77,7 @@ @click="clearTranscript" type="button" :title="run.isRunning.value - ? 'Start a new chat. The run in progress keeps going on the server and stays in the History tab.' + ? 'Start a new chat. The run in progress keeps going; stop it from the Working badge, or wait for it to finish.' : 'Start a new chat'" > @@ -108,13 +118,14 @@ @@ -545,6 +584,31 @@ function clearTranscript() { border-radius: 999px; border: 1px solid transparent; } +/* Status and action in one pinned element: the badge says a run is live and + is also the thing that ends it, so the two can never disagree. */ +.run-control { + font-family: inherit; + cursor: pointer; + transition: color 0.15s ease, background-color 0.15s ease, border-color 0.15s ease; +} +.run-control:hover:not(.is-stopping) { + color: #ffffff; + background-color: rgba(239, 108, 68, 0.22); + border-color: #ef6c44; +} +.run-control.is-stopping { + cursor: default; + opacity: 0.75; +} +.badge-sep { color: rgba(245, 158, 11, 0.55); } +.badge-stop { + display: inline-flex; + align-items: center; + gap: 0.3rem; + font-weight: 600; + color: #ffd7c2; +} +.run-control:hover .badge-stop { color: #ffffff; } .header-status.running { color: #ffe7b8; background-color: rgba(245, 158, 11, 0.13); diff --git a/gui/views/chat/composables/useMcpRun.js b/gui/views/chat/composables/useMcpRun.js index c3c4e6e..374f11f 100644 --- a/gui/views/chat/composables/useMcpRun.js +++ b/gui/views/chat/composables/useMcpRun.js @@ -1,20 +1,26 @@ // One MCP run: POST /plugin/mcp/execute, then poll /plugin/mcp/status until -// FINISHED or FAILED. Returns reactive state plus start() and attach(). +// it reaches a terminal status. Returns reactive state plus start(), +// attach() and requestStop(). // // The run belongs to the server, so it outlives the browser. attach() is the // counterpart to start() for one already in flight, which is how a remounted -// view picks its run back up. +// view picks its run back up. requestStop() is the only call that ends the +// server-side run; stop() below just detaches this page from it. import { ref, computed } from 'vue' import { SESSION_EXPIRED } from '../../../composables/request.js' const POLL_INTERVAL_MS = 1000 +// Anything the server will not move again. KILLED is a run the user stopped, +// or one the boot sweep reconciled; polling it forever locks the page. +const TERMINAL_STATUSES = new Set(['FINISHED', 'FAILED', 'KILLED']) + // A single failed status GET is usually a blip, not a dead run. const MAX_CONSECUTIVE_POLL_ERRORS = 3 export function useMcpRun($api) { - const status = ref('idle') // 'idle' | 'RUNNING' | 'FINISHED' | 'FAILED' + const status = ref('idle') // 'idle' | 'RUNNING' | 'FINISHED' | 'FAILED' | 'KILLED' const stage = ref('') const runId = ref(null) const prompt = ref('') @@ -22,10 +28,15 @@ export function useMcpRun($api) { const finalResult = ref('') const trajectory = ref({}) const errorMessage = ref('') + // Set from the moment Stop is pressed until the run reaches a terminal + // status. Unwinding an agent loop takes a few seconds, and the server + // still reports RUNNING throughout. + const stopping = ref(false) const isRunning = computed(() => status.value === 'RUNNING') const isFinished = computed(() => status.value === 'FINISHED') const isFailed = computed(() => status.value === 'FAILED') + const isStopping = computed(() => stopping.value && isRunning.value) let pollTimer = null let consecutivePollErrors = 0 @@ -35,6 +46,11 @@ export function useMcpRun($api) { // Bumped by reset(). runId cannot mark supersession on its own: it is null // for the whole /execute POST, which is the window "New chat" opens. let generation = 0 + // Stop pressed before the run had an id. That window is not small: /execute + // mints the MLflow run before it answers, so a slow tracking server holds it + // open for as long as its retries last, which is exactly when a user reaches + // for Stop. start() sends the cancel once the id arrives. + let stopPending = false function _clearTimer() { if (pollTimer) { @@ -52,6 +68,8 @@ export function useMcpRun($api) { finalResult.value = '' trajectory.value = {} errorMessage.value = '' + stopping.value = false + stopPending = false consecutivePollErrors = 0 detached = false generation += 1 @@ -73,6 +91,10 @@ export function useMcpRun($api) { const response = await $api.post('/plugin/mcp/execute', payload) if (gen !== generation) return null runId.value = response.data.run_id + if (stopPending) { + stopPending = false + _sendCancel(runId.value) + } if (!detached) _beginPolling(runId.value) // The caller needs fields above the per-run scope, like session_id. return response.data @@ -131,7 +153,7 @@ export function useMcpRun($api) { } consecutivePollErrors = 0 _applySnapshot(res.data) - return status.value !== 'FINISHED' && status.value !== 'FAILED' + return !TERMINAL_STATUSES.has(status.value) } catch (err) { if (_superseded(id)) return false // 404 means the run left the live cache: server restart, or LRU @@ -174,6 +196,35 @@ export function useMcpRun($api) { pollTimer = timer } + async function _sendCancel(id) { + const gen = generation + try { + await $api.post('/plugin/mcp/cancel', { run_id: id }) + } catch { + // The run is still whatever it was; let the poller report it. Guarded + // like every other post-await write here: a cancel that rejects only + // after a newer run started would otherwise clear that run's badge + // back to "Working" while its own stop is still in flight. + if (gen === generation && runId.value === id) stopping.value = false + } + } + + /** + * Ask the server to cancel this run. Usable for the whole life of the run, + * including before /execute has answered with an id. Polling carries on: the + * run reaches KILLED only once its workflow has unwound, which is what the + * caller waits for. Pressing twice is a harmless no-op. + */ + async function requestStop() { + if (!isRunning.value) return + stopping.value = true + if (!runId.value) { + stopPending = true + return + } + await _sendCancel(runId.value) + } + /** Stop polling without touching run state; the server run keeps going. */ function stop() { detached = true @@ -182,7 +233,7 @@ export function useMcpRun($api) { return { status, stage, runId, prompt, reasoning, finalResult, trajectory, - errorMessage, isRunning, isFinished, isFailed, - start, attach, stop, reset, + errorMessage, isRunning, isFinished, isFailed, isStopping, + start, attach, requestStop, stop, reset, } } diff --git a/gui/views/chat/messageState.js b/gui/views/chat/messageState.js new file mode 100644 index 0000000..c8cc666 --- /dev/null +++ b/gui/views/chat/messageState.js @@ -0,0 +1,24 @@ +// One definition of "this assistant bubble carries an answer the user can +// read", shared by the transcript that renders it and the header that counts +// it. The two used to decide separately: the header allow-listed FINISHED +// while ChatMessage rendered a result for anything that was not RUNNING, +// FAILED or KILLED. A snapshot arriving without a status becomes 'unknown' +// (useMcpRun._applySnapshot), which showed an answer the header refused to +// count. + +/** + * Statuses an assistant bubble renders as something other than an answer. + * ChatMessage carries one branch per member, each with its own styling. + */ +export const UNDELIVERED_STATUSES = new Set(['RUNNING', 'FAILED', 'KILLED']) + +/** True when this message is an assistant answer, whatever its exact status. */ +export function isDeliveredResponse(message) { + // Missing role defaults to assistant, matching how ChatMessage decides + // which bubble to draw. Disagreeing here is the bug this module exists + // to prevent. + return ( + (message?.role || 'assistant') === 'assistant' + && !UNDELIVERED_STATUSES.has(message?.status) + ) +} diff --git a/gui/views/mcp_history.vue b/gui/views/mcp_history.vue index a8167c0..b98b7a6 100644 --- a/gui/views/mcp_history.vue +++ b/gui/views/mcp_history.vue @@ -59,7 +59,7 @@ {{ formatDate(run.start_time) }} - {{ run.status }} + {{ statusLabel(run) }} @@ -109,7 +109,7 @@

Status: - {{ runDetail.status }} + {{ statusLabel(runDetail) }}

@@ -241,7 +241,7 @@ const filteredRuns = computed(() => { return ( run.prompt?.toLowerCase().includes(query) || run.run_id?.toLowerCase().includes(query) || - run.status?.toLowerCase().includes(query) || + statusLabel(run).toLowerCase().includes(query) || run.stage?.toLowerCase().includes(query) ) }) @@ -315,9 +315,10 @@ function formatDuration(run) { const { start_time: startTime, end_time: endTime, status } = run if (!startTime) return '-' if (!endTime) return 'Running...' - // A run reconciled at boot was terminated by the sweep, not by itself, - // so its end_time says when we noticed rather than when it stopped. - if (status?.toUpperCase() === 'KILLED') return '-' + // A run reconciled at boot was terminated by the sweep, not by itself, so + // its end_time says when we noticed rather than when it stopped. A run the + // user stopped ended when they said so, and its duration is real. + if (status?.toUpperCase() === 'KILLED' && run.stop_reason !== 'user') return '-' try { const duration = endTime - startTime @@ -337,6 +338,14 @@ function formatDuration(run) { } } +// MLflow has one terminal status for "did not finish on its own", so the tag +// text is what separates a run the user stopped from one the boot sweep +// closed out. +function statusLabel(run) { + if (run?.stop_reason === 'user') return 'STOPPED' + return run?.status || '' +} + function getStatusClass(status) { switch (status?.toUpperCase()) { case 'FINISHED': diff --git a/hook.py b/hook.py index c4997ab..2d35efb 100644 --- a/hook.py +++ b/hook.py @@ -179,6 +179,7 @@ async def enable(services): mcp_api = McpAPI(services) app.router.add_route('POST', '/plugin/mcp/execute', mcp_api.execute) app.router.add_route('GET', '/plugin/mcp/status', mcp_api.status) + app.router.add_route('POST', '/plugin/mcp/cancel', mcp_api.cancel) app.router.add_route('POST', "/plugin/mcp/rag/upload", mcp_api.upload_rag) app.router.add_route('GET', "/plugin/mcp/rag/list", mcp_api.list_rag) app.router.add_route('GET', '/plugin/mcp/servers', mcp_api.list_servers) diff --git a/tests/test_run_cancellation.py b/tests/test_run_cancellation.py new file mode 100644 index 0000000..f6722f1 --- /dev/null +++ b/tests/test_run_cancellation.py @@ -0,0 +1,338 @@ +"""Stopping one in-flight run leaves it KILLED and leaves the others alone. + +A bare task.cancel() is worse than no Stop button: CancelledError is a +BaseException, so the `except Exception` arm misses it while the finally still +writes FAILED, and the run cache stays RUNNING forever with the page locked +behind it. anyio makes it worse by repackaging a cancel that lands during MCP +session startup as an ExceptionGroup, which `except Exception` does catch. +Both shapes are driven here. +""" +import asyncio +import json +import types + +import mlflow +import pytest +from mlflow.tracking import MlflowClient + +from plugins.mcp.app import mcp_svc as mcp_svc_module +from plugins.mcp.app.mcp_svc import MCPService + +EXPERIMENT = "test-run-cancellation" + + +def _workflow(workflow_id, run): + return types.SimpleNamespace( + id=workflow_id, + display_name=workflow_id.title(), + mlflow_experiment=EXPERIMENT, + required_servers=[], + optional_servers=[], + accepted_capabilities=[], + supports_chat_history=False, + run=run, + ) + + +@pytest.fixture +def svc(tmp_path, monkeypatch): + previous_uri = mlflow.get_tracking_uri() + mlflow.set_tracking_uri(f"sqlite:///{tmp_path}/mlflow.db") + monkeypatch.setattr(mcp_svc_module, "_SESSIONS_FILE", tmp_path / "sessions.json") + monkeypatch.setattr( + mcp_svc_module, "resolve_llm_config", lambda cfg: {"model": "gpt-4o-mini"} + ) + service = MCPService(services={}) + yield service + if mlflow.active_run(): + mlflow.end_run() + mlflow.set_tracking_uri(previous_uri) + + +async def _forever(prompt, lm_obj, run_id=None, **kwargs): + await asyncio.sleep(3600) + + +async def _settle(): + """Let every background run task finish unwinding.""" + await asyncio.gather( + *(asyncio.all_tasks() - {asyncio.current_task()}), return_exceptions=True + ) + + +def test_cancel_leaves_the_run_killed_not_failed(svc): + """The defect: FAILED in MLflow, RUNNING in the cache, page locked.""" + svc.workflow_registry = {"hang": _workflow("hang", _forever)} + + async def start_then_stop(): + handle = await svc.execute(workflow_id="hang", prompt="p") + await asyncio.sleep(0) + assert svc.cancel_run(handle["run_id"]) is True + await _settle() + return handle + + handle = asyncio.run(start_then_stop()) + run_id = handle["run_id"] + + snapshot = svc.get_run(run_id) + assert snapshot["status"] == "KILLED" + assert snapshot["stage"] == "stopped by user" + # The bubble heading already says it stopped; this line says what remains. + assert snapshot["error"].startswith("Anything already created in CALDERA stays") + assert "operation" in snapshot["error"] + + run = MlflowClient().get_run(run_id) + assert run.info.status == "KILLED" + assert run.info.end_time is not None + assert run.data.tags["mcp.cancelled"] == "user" + + +def test_a_stopped_run_is_not_left_carrying_an_error_param(svc): + """The service must not route a cancel through its failure branch. + + Covers the service half only; the workflow half of this bug lives in + tests/test_workflow_failure_ownership.py. The exception here is the + shape anyio really produced when Stop tore down the MCP stdio + transport: an ExceptionGroup wrapping BrokenResourceError with no + CancelledError leaf, so the type alone cannot identify the cancel. + """ + async def transport_torn_down(prompt, lm_obj, run_id=None, **kwargs): + try: + await asyncio.sleep(3600) + except asyncio.CancelledError: + # No CancelledError leaf: exactly what anyio hands back. + raise ExceptionGroup( + "unhandled errors in a TaskGroup", [RuntimeError("BrokenResourceError")] + ) + + svc.workflow_registry = {"torn": _workflow("torn", transport_torn_down)} + + async def start_then_stop(): + handle = await svc.execute(workflow_id="torn", prompt="p") + await asyncio.sleep(0) + svc.cancel_run(handle["run_id"]) + await _settle() + return handle + + handle = asyncio.run(start_then_stop()) + run = MlflowClient().get_run(handle["run_id"]) + + assert run.info.status == "KILLED" + assert run.data.tags["mcp.cancelled"] == "user" + # The whole point: History must not show this as an error. + assert "error" not in run.data.params + assert "traceback" not in run.data.params + + +def test_a_genuine_failure_still_records_its_traceback(svc): + """Moving the param must not cost real failures their diagnosis.""" + async def boom(prompt, lm_obj, run_id=None, **kwargs): + raise RuntimeError("workflow exploded") + + svc.workflow_registry = {"boom": _workflow("boom", boom)} + + async def once(): + handle = await svc.execute(workflow_id="boom", prompt="p") + await _settle() + return handle + + handle = asyncio.run(once()) + run = MlflowClient().get_run(handle["run_id"]) + + assert run.info.status == "FAILED" + assert run.data.params["error"] == "workflow exploded" + assert "RuntimeError: workflow exploded" in run.data.params["traceback"] + + +def test_cancel_wrapped_in_an_exception_group_is_still_a_stop(svc): + """anyio hands an early cancel back as an ExceptionGroup, an Exception.""" + async def anyio_shaped(prompt, lm_obj, run_id=None, **kwargs): + try: + await asyncio.sleep(3600) + except asyncio.CancelledError: + raise ExceptionGroup("unhandled errors in a TaskGroup", [RuntimeError("closed")]) + + svc.workflow_registry = {"wrapped": _workflow("wrapped", anyio_shaped)} + + async def start_then_stop(): + handle = await svc.execute(workflow_id="wrapped", prompt="p") + await asyncio.sleep(0) + svc.cancel_run(handle["run_id"]) + await _settle() + return handle + + handle = asyncio.run(start_then_stop()) + + assert svc.get_run(handle["run_id"])["status"] == "KILLED" + run = MlflowClient().get_run(handle["run_id"]) + assert run.info.status == "KILLED" + assert run.data.tags["mcp.cancelled"] == "user" + + +def test_cancel_names_one_run_and_spares_the_others(svc): + svc.workflow_registry = {"hang": _workflow("hang", _forever)} + + async def start_two_stop_one(): + doomed = await svc.execute(workflow_id="hang", prompt="doomed") + spared = await svc.execute(workflow_id="hang", prompt="spared") + await asyncio.sleep(0) + svc.cancel_run(doomed["run_id"]) + # Only the cancelled task should unwind; the other is still sleeping. + await asyncio.sleep(0.05) + # Read inside the loop: asyncio.run() cancels what is left on exit. + return svc.get_run(doomed["run_id"]), svc.get_run(spared["run_id"]) + + doomed, spared = asyncio.run(start_two_stop_one()) + + assert doomed["status"] == "KILLED" + assert spared["status"] == "RUNNING" + + +def test_a_shutdown_cancel_is_not_reported_as_a_user_stop(svc): + """Loop teardown cancels in-flight tasks; nobody pressed Stop.""" + svc.workflow_registry = {"hang": _workflow("hang", _forever)} + + async def start_and_walk_away(): + handle = await svc.execute(workflow_id="hang", prompt="p") + await asyncio.sleep(0) + return handle + + # asyncio.run() cancels the still-sleeping run task as it closes the loop. + handle = asyncio.run(start_and_walk_away()) + + assert svc.get_run(handle["run_id"])["status"] == "KILLED" + run = MlflowClient().get_run(handle["run_id"]) + assert run.info.status == "KILLED" + assert run.data.tags["mcp.cancelled"] == "server" + + +def test_a_cancel_before_the_task_body_runs_still_closes_the_run(svc): + """Nothing recorded it, so nothing else would ever move it off RUNNING.""" + svc.workflow_registry = {"hang": _workflow("hang", _forever)} + + async def stop_before_it_starts(): + handle = await svc.execute(workflow_id="hang", prompt="p") + # No sleep(0): the task exists but has never been scheduled. + assert svc.cancel_run(handle["run_id"]) is True + await _settle() + return handle + + handle = asyncio.run(stop_before_it_starts()) + + assert svc.get_run(handle["run_id"])["status"] == "KILLED" + run = MlflowClient().get_run(handle["run_id"]) + assert run.info.status == "KILLED" + assert run.data.tags["mcp.cancelled"] == "user" + + +def test_cancelling_twice_or_after_the_run_ended_is_harmless(svc): + async def quick(prompt, lm_obj, run_id=None, **kwargs): + return {"process_result": "done"} + + svc.workflow_registry = { + "hang": _workflow("hang", _forever), + "quick": _workflow("quick", quick), + } + + async def scenario(): + hung = await svc.execute(workflow_id="hang", prompt="p") + await asyncio.sleep(0) + first = svc.cancel_run(hung["run_id"]) + second = svc.cancel_run(hung["run_id"]) + await _settle() + third = svc.cancel_run(hung["run_id"]) + + done = await svc.execute(workflow_id="quick", prompt="p") + await _settle() + after_finish = svc.cancel_run(done["run_id"]) + return first, second, third, after_finish, done + + first, second, third, after_finish, done = asyncio.run(scenario()) + + assert first is True + # Already cancelling, already over, and never ours: all no-ops. + assert (second, third, after_finish) == (True, False, False) + assert svc.get_run(done["run_id"])["status"] == "FINISHED" + assert MlflowClient().get_run(done["run_id"]).info.status == "FINISHED" + + +def test_the_task_registry_empties_itself(svc): + """It is not the LRU run cache; nothing else evicts it.""" + async def quick(prompt, lm_obj, run_id=None, **kwargs): + return {"process_result": "done"} + + svc.workflow_registry = {"quick": _workflow("quick", quick)} + + async def once(): + await svc.execute(workflow_id="quick", prompt="p") + await _settle() + + asyncio.run(once()) + assert svc._tasks == {} + + +class _FakeRequest: + """Not a web_request.Request, so check_authorization stands aside.""" + + def __init__(self, payload): + self._payload = payload + + async def json(self): + if self._payload is _NOT_JSON: + raise ValueError("not json") + return self._payload + + +_NOT_JSON = object() + + +class _StubService: + def __init__(self, cancellable, snapshot=None): + self._cancellable = cancellable + self._snapshot = snapshot + self.asked = [] + + def cancel_run(self, run_id): + self.asked.append(run_id) + return self._cancellable + + def get_run(self, run_id): + return self._snapshot + + +def _api(svc): + from plugins.mcp.app import mcp_api + return mcp_api.McpAPI({"mcp_svc": svc}) + + +async def _cancel(svc, payload): + resp = await _api(svc).cancel(_FakeRequest(payload)) + return resp.status, json.loads(resp.body.decode()) + + +class TestCancelEndpoint: + @pytest.mark.asyncio + @pytest.mark.parametrize( + "payload", + [_NOT_JSON, ["run_id"], {}, {"run_id": ""}, {"run_id": ["a"]}], + ) + async def test_a_body_without_a_run_id_is_refused(self, payload): + svc = _StubService(cancellable=True) + status, _ = await _cancel(svc, payload) + assert status == 400 + assert svc.asked == [] + + @pytest.mark.asyncio + async def test_a_live_run_is_asked_to_stop(self): + svc = _StubService(cancellable=True, snapshot={"status": "RUNNING"}) + status, body = await _cancel(svc, {"run_id": "abc"}) + assert (status, svc.asked) == (200, ["abc"]) + assert body == {"run_id": "abc", "cancelling": True, "status": "RUNNING"} + + @pytest.mark.asyncio + async def test_a_run_nothing_owns_answers_200_not_an_error(self): + """Pressing Stop twice, or after the run ended, must not look broken.""" + svc = _StubService(cancellable=False) + status, body = await _cancel(svc, {"run_id": "abc"}) + assert status == 200 + assert body == {"run_id": "abc", "cancelling": False, "status": "UNKNOWN"} diff --git a/tests/test_run_stage_progress.py b/tests/test_run_stage_progress.py new file mode 100644 index 0000000..db443d9 --- /dev/null +++ b/tests/test_run_stage_progress.py @@ -0,0 +1,157 @@ +"""Stage transitions a workflow reports have to reach the polling UI. + +The defect: workflows report progress with tracker.set_tag("stage", ...), +which writes to MLflow only, while /status serves MCPService._runs. The +cached snapshot was written twice per run -- "initializing" at the top of +_run_execution and a terminal value at the bottom -- so the chat UI's stage +line read "initializing" for the whole run no matter how many stages the +workflow announced. + +A workflow also binds its OWN RunTracker to the run id the service started, +so the mirror has to follow the run rather than one tracker instance. +""" +import asyncio +import types + +import mlflow +import pytest + +from plugins.mcp.app import mcp_svc as mcp_svc_module +from plugins.mcp.app.mcp_svc import MCPService +from plugins.mcp.app.mlflow_run import RunTracker, _stage_observers + +EXPERIMENT = "test-run-stage-progress" + + +def _workflow(workflow_id, run): + return types.SimpleNamespace( + id=workflow_id, + display_name=workflow_id.title(), + mlflow_experiment=EXPERIMENT, + required_servers=[], + optional_servers=[], + accepted_capabilities=[], + supports_chat_history=False, + run=run, + ) + + +@pytest.fixture +def svc(tmp_path, monkeypatch): + previous_uri = mlflow.get_tracking_uri() + mlflow.set_tracking_uri(f"sqlite:///{tmp_path}/mlflow.db") + monkeypatch.setattr(mcp_svc_module, "_SESSIONS_FILE", tmp_path / "sessions.json") + monkeypatch.setattr( + mcp_svc_module, "resolve_llm_config", lambda cfg: {"model": "gpt-4o-mini"} + ) + service = MCPService(services={}) + yield service + if mlflow.active_run(): + mlflow.end_run() + mlflow.set_tracking_uri(previous_uri) + + +def _drive(svc, workflow_id, prompt="p"): + async def once(): + handle = await svc.execute(workflow_id=workflow_id, prompt=prompt) + await asyncio.gather(*asyncio.all_tasks() - {asyncio.current_task()}) + return handle + + return asyncio.run(once()) + + +def test_a_stage_a_workflow_reports_reaches_the_polled_snapshot(svc): + """The bug: /status served "initializing" for the life of the run.""" + seen = [] + + async def staged(prompt, lm_obj, run_id=None, **kwargs): + # Exactly what the real workflows do: bind a fresh tracker to the + # run the service already started, then tag progress through it. + tracker = RunTracker.bind(run_id, EXPERIMENT, "staged") + for stage in ("listing tools", "creating DSPy ReAct instance"): + tracker.set_tag("stage", stage) + seen.append(svc.get_run(run_id)["stage"]) + return {"process_result": "ok", "reasoning": "", "trajectory": {}} + + svc.workflow_registry = {"staged": _workflow("staged", staged)} + _drive(svc, "staged") + + assert seen == ["listing tools", "creating DSPy ReAct instance"] + + +def test_the_terminal_stage_still_wins(svc): + """A run that finished must not be left showing its last live stage.""" + async def staged(prompt, lm_obj, run_id=None, **kwargs): + RunTracker.bind(run_id, EXPERIMENT, "staged").set_tag("stage", "executing") + return {"process_result": "ok", "reasoning": "", "trajectory": {}} + + svc.workflow_registry = {"staged": _workflow("staged", staged)} + handle = _drive(svc, "staged") + + assert svc.get_run(handle["run_id"])["stage"] == "complete" + + +def test_a_late_stage_tag_cannot_reopen_a_terminal_snapshot(svc): + """Guarded on RUNNING: a stray tag must not overwrite a final stage.""" + async def noop(prompt, lm_obj, run_id=None, **kwargs): + return {"process_result": "ok", "reasoning": "", "trajectory": {}} + + svc.workflow_registry = {"noop": _workflow("noop", noop)} + handle = _drive(svc, "noop") + run_id = handle["run_id"] + + svc._update_stage(run_id, "listing tools") + + assert svc.get_run(run_id)["stage"] == "complete" + + +def test_the_observer_is_released_when_the_run_ends(svc): + """Registered per run; leaking one would pin the service per request.""" + async def noop(prompt, lm_obj, run_id=None, **kwargs): + return {"process_result": "ok", "reasoning": "", "trajectory": {}} + + svc.workflow_registry = {"noop": _workflow("noop", noop)} + handle = _drive(svc, "noop") + + assert handle["run_id"] not in _stage_observers + + +def test_a_failed_run_releases_its_observer_too(svc): + """The finally has to run on the error path, not just the happy one.""" + async def boom(prompt, lm_obj, run_id=None, **kwargs): + raise RuntimeError("workflow exploded") + + svc.workflow_registry = {"boom": _workflow("boom", boom)} + handle = _drive(svc, "boom") + + assert handle["run_id"] not in _stage_observers + assert svc.get_run(handle["run_id"])["stage"] == "error" + + +def test_concurrent_runs_do_not_cross_stages(svc): + """Keyed by run id, so one run's progress cannot land on another's.""" + async def staged(prompt, lm_obj, run_id=None, **kwargs): + tracker = RunTracker.bind(run_id, EXPERIMENT, "staged") + tracker.set_tag("stage", f"working on {prompt}") + # Hold both runs open at once so the observers genuinely overlap. + await asyncio.sleep(0.05) + return {"process_result": "ok", "reasoning": "", "trajectory": {}} + + svc.workflow_registry = {"staged": _workflow("staged", staged)} + + async def two_at_once(): + handles = [await svc.execute(workflow_id="staged", prompt=f"p{i}") + for i in range(2)] + # Sampled mid-flight: after this sleep both workflows have tagged + # their stage but neither has returned. + await asyncio.sleep(0.02) + sampled = [svc.get_run(h["run_id"])["stage"] for h in handles] + await asyncio.gather(*asyncio.all_tasks() - {asyncio.current_task()}) + return sampled + + assert asyncio.run(two_at_once()) == ["working on p0", "working on p1"] + + +def test_a_tracker_for_an_unregistered_run_is_harmless(svc): + """Trackers outlive runs (the cancel path builds one); no KeyError.""" + RunTracker(run_id="never-registered").set_tag("stage", "orphan") diff --git a/tests/test_workflow_failure_ownership.py b/tests/test_workflow_failure_ownership.py new file mode 100644 index 0000000..b1e9601 --- /dev/null +++ b/tests/test_workflow_failure_ownership.py @@ -0,0 +1,190 @@ +"""Whoever owns a run's lifecycle owns writing its terminal state. + +Observed in the field: pressing Stop on a Plan-and-Execute run tore down the +MCP stdio transport, and anyio surfaced that as + + ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception) + +-- anyio.BrokenResourceError + +with no CancelledError leaf anywhere in it. A workflow cannot tell that from +a genuine crash, but its `except Exception` ran full failure bookkeeping +before re-raising -- printing a traceback that read as a server fault and, +worse, calling log_param("error"/"traceback"). MLflow params are immutable, +so the service's later KILLED tags could not take them back and History +showed a deliberate stop as an error. + +The service knows the difference, because it records the intent in +_cancelling before it cancels. So a workflow handed a run by the orchestrator +now re-raises and lets the orchestrator classify; only a workflow that minted +its own run writes that run's failure state. +""" +import asyncio +import logging +import types + +import mlflow +import pytest +from mlflow.tracking import MlflowClient + +from plugins.mcp.app.workflows import author, plan_execute +from plugins.mcp.app.mlflow_run import RunTracker + +# The shape anyio actually produced, reproduced exactly: an Exception (so the +# workflow's `except Exception` catches it) carrying no cancellation leaf. +def _torn_down_transport(): + return ExceptionGroup( + "unhandled errors in a TaskGroup", + [RuntimeError("anyio.BrokenResourceError")], + ) + + +class _ExplodingStack: + """Stands in for AsyncExitStack, the first thing inside the try block.""" + + async def __aenter__(self): + raise _torn_down_transport() + + async def __aexit__(self, *_exc): + return False + + +@pytest.fixture(params=[plan_execute, author], ids=["plan_execute", "author"]) +def workflow(request, tmp_path, monkeypatch): + """Both workflows carry the same handler, so both are held to it.""" + previous_uri = mlflow.get_tracking_uri() + mlflow.set_tracking_uri(f"sqlite:///{tmp_path}/mlflow.db") + monkeypatch.setattr( + request.param, "AsyncExitStack", lambda *a, **k: _ExplodingStack() + ) + yield request.param + if mlflow.active_run(): + mlflow.end_run() + mlflow.set_tracking_uri(previous_uri) + + +# Enough config to clear each workflow's credential guard, which fires +# before the try block and is a genuine failure, not the case under test. +_LM = { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + "api_base": "http://localhost:1", +} + + +def _service_workflow(workflow_id, run): + """Registry entry shaped the way MCPService expects.""" + return types.SimpleNamespace( + id=workflow_id, display_name=workflow_id.title(), + mlflow_experiment="test-failure-logging", required_servers=[], + optional_servers=[], accepted_capabilities=[], + supports_chat_history=False, run=run, + ) + + +@pytest.fixture +def svc_logs(tmp_path, monkeypatch): + """MCPService wired to a temp store, plus a helper that drives one run.""" + from plugins.mcp.app import mcp_svc as mcp_svc_module + from plugins.mcp.app.mcp_svc import MCPService + + previous_uri = mlflow.get_tracking_uri() + mlflow.set_tracking_uri(f"sqlite:///{tmp_path}/logs.db") + monkeypatch.setattr(mcp_svc_module, "_SESSIONS_FILE", tmp_path / "sessions.json") + monkeypatch.setattr( + mcp_svc_module, "resolve_llm_config", lambda cfg: {"model": "gpt-4o-mini"} + ) + service = MCPService(services={}) + + def drive(workflow_id): + async def once(): + await service.execute(workflow_id=workflow_id, prompt="p") + await asyncio.gather(*asyncio.all_tasks() - {asyncio.current_task()}) + asyncio.run(once()) + + yield service, drive + if mlflow.active_run(): + mlflow.end_run() + mlflow.set_tracking_uri(previous_uri) + + +def _drive(module, run_id): + with pytest.raises(BaseException): + asyncio.run(module.run("test", lm_obj=dict(_LM), run_id=run_id)) + + +def test_an_orchestrated_run_is_left_for_the_orchestrator_to_classify(workflow): + """The bug: a stopped run kept an immutable error param naming a crash.""" + owner = RunTracker.start(workflow._MLFLOW_EXPERIMENT, "orchestrated") + + _drive(workflow, owner.run_id) + + run = MlflowClient().get_run(owner.run_id) + assert "error" not in run.data.params + assert "traceback" not in run.data.params + # Still open: the orchestrator terminates the run it handed out, and it + # is the layer that decides whether this was KILLED or FAILED. + assert run.info.status == "RUNNING" + + +def test_a_run_the_workflow_minted_is_still_the_workflow_to_close(workflow): + """Standalone use has no orchestrator, so the handler must stay.""" + before = {r.info.run_id for r in _experiment_runs(workflow)} + + _drive(workflow, None) + + minted = [r for r in _experiment_runs(workflow) if r.info.run_id not in before] + assert len(minted) == 1 + run = minted[0] + assert run.info.status == "FAILED" + assert "BrokenResourceError" in run.data.params["error"] + assert "traceback" in run.data.params + + +def _experiment_runs(module): + client = MlflowClient() + experiment = client.get_experiment_by_name(module._MLFLOW_EXPERIMENT) + return client.search_runs([experiment.experiment_id]) if experiment else [] + + +def test_a_genuine_failure_puts_its_stack_on_the_console(svc_logs, caplog): + """Workflows stopped printing it, so the service has to. + + The MLflow traceback param is best-effort (RunTracker._write swallows a + dead tracking store), so exc_info here is the only copy an operator is + guaranteed to see. + """ + svc, drive = svc_logs + + async def boom(prompt, lm_obj, run_id=None, **kwargs): + raise RuntimeError("workflow exploded") + + svc.workflow_registry = {"boom": _service_workflow("boom", boom)} + with caplog.at_level(logging.ERROR, logger="plugins.mcp"): + drive("boom") + + failures = [r for r in caplog.records if "Execution failed" in r.getMessage()] + assert failures, "no failure logged at ERROR" + assert failures[0].exc_info is not None + assert "RuntimeError: workflow exploded" in caplog.text + + +def test_a_stopped_run_logs_no_stack(svc_logs, caplog): + """The cancel path returns above the failure branch; keep it that way.""" + svc, _ = svc_logs + + async def slow(prompt, lm_obj, run_id=None, **kwargs): + await asyncio.sleep(3600) + + svc.workflow_registry = {"slow": _service_workflow("slow", slow)} + + async def start_then_stop(): + handle = await svc.execute(workflow_id="slow", prompt="p") + await asyncio.sleep(0) + svc.cancel_run(handle["run_id"]) + for _ in range(6): + await asyncio.sleep(0) + + with caplog.at_level(logging.ERROR, logger="plugins.mcp"): + asyncio.run(start_then_stop()) + + assert not [r for r in caplog.records if "Execution failed" in r.getMessage()]