diff --git a/app/mcp_api.py b/app/mcp_api.py index 4404798..424b1b7 100644 --- a/app/mcp_api.py +++ b/app/mcp_api.py @@ -411,7 +411,12 @@ async def list_runs(self, request): "prompt": run_data.params.get("prompt", ""), "stage": run_data.tags.get("stage", ""), "model": run_data.params.get("model", ""), - "process_result": run_data.params.get("process_result", ""), + # The result is written as a tag, never as a param, + # so params.process_result was always empty. + "process_result": ( + run_data.tags.get("process_result_summary") + or run_data.tags.get("process_result", "") + ), } all_runs.append(run_record) @@ -462,7 +467,13 @@ async def get_run_detail(self, request): "stage": run.data.tags.get("stage"), "reasoning": run.data.tags.get("reasoning"), "prompt": run.data.params.get("prompt"), - "process_result": run.data.params.get("process_result"), + # Workflows tag the full result and log it as the + # result_summary param; params.process_result never exists, + # so the detail view's Result box never rendered. + "process_result": ( + run.data.tags.get("process_result") + or run.data.params.get("result_summary", "") + ), } return web.json_response(response) diff --git a/app/mcp_svc.py b/app/mcp_svc.py index 7914c50..039fa3c 100644 --- a/app/mcp_svc.py +++ b/app/mcp_svc.py @@ -6,10 +6,13 @@ import tempfile from pathlib import Path from app.utility.base_service import BaseService -import mlflow import asyncio from plugins.mcp.app.config import resolve_llm_config +from plugins.mcp.app.mlflow_run import ( + RunTracker, + reconcile_orphaned_runs as _reconcile_orphaned_runs, +) # Where chat session history lives on disk. Co-located with RAG uploads @@ -141,6 +144,40 @@ def _record_run(self, run_id: str, snapshot: dict) -> None: while len(self._runs) > _RUN_CACHE_LIMIT: self._runs.popitem(last=False) + async def reconcile_orphaned_runs(self) -> list: + """Terminate MLflow runs left RUNNING by a process that died. + + Only the task that owns a run writes its terminal status, so a + `kill -9` mid-flight strands the run as RUNNING with a null + end_time and History shows it as running forever. Called once at + plugin enable(), where the live cache is empty and every RUNNING + run in a workflow experiment is residue from a previous process. + Runs still in the cache are skipped, so this is also safe to call + while requests are in flight. + """ + experiments = { + wf.mlflow_experiment for wf in self.workflow_registry.values() + if getattr(wf, "mlflow_experiment", None) + } + if not experiments: + return [] + live = { + run_id for run_id, snapshot in self._runs.items() + if snapshot.get("status") == "RUNNING" + } + try: + reconciled = await asyncio.to_thread( + _reconcile_orphaned_runs, sorted(experiments), live + ) + except Exception as e: + self.log.warning(f"[MCP] Orphaned run reconciliation failed: {e}") + return [] + if reconciled: + self.log.info( + f"[MCP] Reconciled {len(reconciled)} orphaned MLflow run(s) to KILLED" + ) + return reconciled + def _session_turns(self, session_id: str) -> list[dict]: """Return the recorded turns for a session, or [] if unknown.""" return self._sessions.get(session_id) or [] @@ -365,35 +402,14 @@ async def execute(self, focus: str = None, prompt: str = "", model_config: dict rag_cfg["ssl_verify"] = resolved_lm.get("ssl_verify") cap_settings["rag"] = rag_cfg - # mlflow keeps a single "active run" per Python process; firing - # two concurrent /plugin/mcp/execute requests would collide on - # the second one with "Run X is already active". Use the - # MlflowClient to mint a run record WITHOUT making it the - # process-wide active run — the actual work in _run_execution - # below uses `mlflow.start_run(run_id=...)` in its own context - # manager, which is per-task and doesn't conflict. - from mlflow.tracking import MlflowClient as _Mlc - # Mint into the workflow's own experiment. Hardcoding one name here - # broke every run after the first: mlflow.start_run(run_id=...) raises - # when the process-wide active experiment differs from the run's, and - # a workflow that set_experiment a different name flipped it. - # create_experiment not set_experiment: the latter would mutate the - # process-wide active experiment this block avoids. - _exp_name = workflow.mlflow_experiment - _exp = mlflow.get_experiment_by_name(_exp_name) - if _exp is not None: - _exp_id = _exp.experiment_id - else: - try: - _exp_id = _Mlc().create_experiment(_exp_name) - except Exception: - _exp = mlflow.get_experiment_by_name(_exp_name) - _exp_id = _exp.experiment_id if _exp else "0" - run = _Mlc().create_run( - experiment_id=_exp_id, - tags={"mlflow.runName": f"MCP {workflow.display_name}"}, - ) - run_id = run.info.run_id + # Mint the run record without activating it. Every write against it + # goes through a RunTracker bound to this id, so concurrent requests + # never share mlflow's thread-local active-run stack. Minting into + # the workflow's own experiment also keeps History grouped the way + # each workflow declares it. + run_id = RunTracker.start( + workflow.mlflow_experiment, f"MCP {workflow.display_name}" + ).run_id # Resolve the session this run belongs to. The first turn auto-starts # a session keyed by its own run_id; follow-up turns echo the same @@ -436,6 +452,9 @@ async def _run_execution(self, workflow, prompt, run_id, lm_obj, self._runs as the run progresses. MLflow keeps logging tags and params for the History tab and the MLflow UI; it is no longer the source of truth for active runs. + + The run reaches its terminal status exactly once, in the finally + below, written against its own run id. """ self._record_run(run_id, { "status": "RUNNING", @@ -464,108 +483,117 @@ async def _run_execution(self, workflow, prompt, run_id, lm_obj, self._format_session_history(effective_session_id) if supports_history else "" ) + # Every MLflow write below names run_id explicitly. The fluent API + # would route them through a thread-local active-run stack that all + # concurrent tasks share, so an overlapping request could steal the + # terminal status or the tags of this one. + tracker = RunTracker(run_id) + terminal_status = "FAILED" try: - mlflow.end_run() - with mlflow.start_run(run_id=run_id): - mlflow.set_tag("stage", "initializing") - mlflow.set_tag("workflow_id", workflow.id) - mlflow.set_tag("mcp.session_id", effective_session_id) - mlflow.set_tag("mcp.turn_index", prior_turn_count) - if workflow_opts_in and disable_history: - mlflow.set_tag("mcp.history_disabled", "true") - if supports_history: - mlflow.set_tag("mcp.session_history_chars", len(chat_history)) - mlflow.log_param("workflow", workflow.id) - mlflow.log_param("prompt", prompt) - mlflow.log_param("enabled_servers", ",".join(enabled_servers)) - mlflow.log_param("enabled_capabilities", ",".join(enabled_capabilities)) - if isinstance(workflow_context, dict) and workflow_context: - mlflow.log_param( - "workflow_context_keys", - ",".join(sorted(workflow_context.keys())), + tracker.set_tag("stage", "initializing") + tracker.set_tag("workflow_id", workflow.id) + tracker.set_tag("mcp.session_id", effective_session_id) + tracker.set_tag("mcp.turn_index", prior_turn_count) + if workflow_opts_in and disable_history: + tracker.set_tag("mcp.history_disabled", "true") + if supports_history: + tracker.set_tag("mcp.session_history_chars", len(chat_history)) + tracker.log_param("workflow", workflow.id) + tracker.log_param("prompt", prompt) + tracker.log_param("enabled_servers", ",".join(enabled_servers)) + tracker.log_param("enabled_capabilities", ",".join(enabled_capabilities)) + # History's Model column reads params.model; nothing wrote it + # before, so every row rendered as '-'. + if isinstance(lm_obj, dict) and lm_obj.get("model"): + tracker.log_param("model", lm_obj["model"]) + if isinstance(workflow_context, dict) and workflow_context: + tracker.log_param( + "workflow_context_keys", + ",".join(sorted(workflow_context.keys())), + ) + + capability_context = {} + for cap_id in enabled_capabilities: + cap = self.capability_registry[cap_id] + if cap.enrich is None: + continue + settings = capability_settings.get(cap_id, {}) + tracker.set_tag(f"capability_{cap_id}_stage", "running") + try: + self.log.info(f"[MCP] Running capability '{cap_id}' enrich()") + contrib = await cap.enrich(prompt, settings) or {} + capability_context.update(contrib) + tracker.set_tag(f"capability_{cap_id}_stage", "complete") + tracker.set_tag( + f"capability_{cap_id}_fields", + ",".join(sorted(contrib.keys())), ) + except Exception as e: + tracker.set_tag(f"capability_{cap_id}_stage", "error") + self.log.warning(f"[MCP] Capability '{cap_id}' enrich failed: {e}") - capability_context = {} - for cap_id in enabled_capabilities: - cap = self.capability_registry[cap_id] - if cap.enrich is None: - continue - settings = capability_settings.get(cap_id, {}) - mlflow.set_tag(f"capability_{cap_id}_stage", "running") - try: - self.log.info(f"[MCP] Running capability '{cap_id}' enrich()") - contrib = await cap.enrich(prompt, settings) or {} - capability_context.update(contrib) - mlflow.set_tag(f"capability_{cap_id}_stage", "complete") - mlflow.set_tag( - f"capability_{cap_id}_fields", - ",".join(sorted(contrib.keys())), - ) - except Exception as e: - mlflow.set_tag(f"capability_{cap_id}_stage", "error") - self.log.warning(f"[MCP] Capability '{cap_id}' enrich failed: {e}") - - if workflow.run is None: - raise RuntimeError(f"Workflow {workflow.id} has no run() function") - - self.log.info( - f"[MCP] Executing workflow '{workflow.id}' with prompt={prompt!r}" - ) - mlflow.set_tag("stage", "executing workflow") - run_kwargs = dict(capability_context) - if supports_history: - # Workflows that opt in receive accumulated session history - # as a labelled transcript. Workflows that don't opt in - # never see the kwarg, so their signatures stay unchanged. - run_kwargs["chat_history"] = chat_history - if isinstance(workflow_context, dict) and workflow_context: - run_kwargs["workflow_context"] = workflow_context - result = await workflow.run( - prompt, - lm_obj, - run_id=run_id, - enabled_servers=enabled_servers, - server_registry=self.server_registry, - **run_kwargs, - ) + if workflow.run is None: + raise RuntimeError(f"Workflow {workflow.id} has no run() function") - mlflow.set_tag("stage", "complete") - mlflow.set_tag("status", "success") - result_dict = result or {} - process_result = _strip_dspy_markers( - str(result_dict.get("process_result", "")) - ) - reasoning = _strip_dspy_markers( - str(result_dict.get("reasoning", "")) + self.log.info( + f"[MCP] Executing workflow '{workflow.id}' with prompt={prompt!r}" + ) + tracker.set_tag("stage", "executing workflow") + run_kwargs = dict(capability_context) + if supports_history: + # Workflows that opt in receive accumulated session history + # as a labelled transcript. Workflows that don't opt in + # never see the kwarg, so their signatures stay unchanged. + run_kwargs["chat_history"] = chat_history + if isinstance(workflow_context, dict) and workflow_context: + run_kwargs["workflow_context"] = workflow_context + result = await workflow.run( + prompt, + lm_obj, + run_id=run_id, + enabled_servers=enabled_servers, + server_registry=self.server_registry, + **run_kwargs, + ) + + tracker.set_tag("stage", "complete") + tracker.set_tag("status", "success") + terminal_status = "FINISHED" + result_dict = result or {} + process_result = _strip_dspy_markers( + str(result_dict.get("process_result", "")) + ) + reasoning = _strip_dspy_markers( + str(result_dict.get("reasoning", "")) + ) + if process_result: + tracker.set_tag("process_result_summary", process_result[:250]) + + self._record_run(run_id, { + "status": "FINISHED", + "stage": "complete", + "workflow_id": workflow.id, + "session_id": effective_session_id, + "prompt": prompt, + "process_result": process_result, + "reasoning": reasoning, + "trajectory": result_dict.get("trajectory") or {}, + }) + + # Append this turn to the session bucket only on success and + # only for opt-in workflows. Failed runs and single-shot + # workflows leave the session store untouched. + if supports_history and process_result: + self._record_session_turn( + effective_session_id, prompt, process_result ) - if process_result: - mlflow.set_tag("process_result_summary", process_result[:250]) - - self._record_run(run_id, { - "status": "FINISHED", - "stage": "complete", - "workflow_id": workflow.id, - "session_id": effective_session_id, - "prompt": prompt, - "process_result": process_result, - "reasoning": reasoning, - "trajectory": result_dict.get("trajectory") or {}, - }) - - # Append this turn to the session bucket only on success and - # only for opt-in workflows. Failed runs and single-shot - # workflows leave the session store untouched. - if supports_history and process_result: - self._record_session_turn( - effective_session_id, prompt, process_result - ) except Exception as e: error_summary = _summarize_exception(e) self.log.error(f"[MCP] Execution failed: {error_summary}") - mlflow.set_tag("stage", "error") - mlflow.set_tag("status", "error") - mlflow.log_param("error", error_summary) + tracker.set_tag("stage", "error") + tracker.set_tag("status", "error") + tracker.log_param("error", error_summary) self._record_run(run_id, { "status": "FAILED", "stage": "error", @@ -579,4 +607,4 @@ async def _run_execution(self, workflow, prompt, run_id, lm_obj, }) finally: - mlflow.end_run() + tracker.terminate(terminal_status) diff --git a/app/mlflow_run.py b/app/mlflow_run.py new file mode 100644 index 0000000..bc3929f --- /dev/null +++ b/app/mlflow_run.py @@ -0,0 +1,137 @@ +"""Run-bound MLflow logging. + +The fluent API (``mlflow.set_tag``, ``mlflow.end_run``) writes to whichever +run sits on ``mlflow.tracking.fluent._active_run_stack``, and that stack is +a thread-local. Every ``/plugin/mcp/execute`` request is an asyncio task on +the single aiohttp event-loop thread, so concurrent runs share one pointer: +a bare ``end_run()`` terminates whichever run is on top, which during +overlap belongs to another request, and a ``set_tag()`` with an empty stack +silently mints a phantom run in the active experiment. + +``RunTracker`` binds an ``MlflowClient`` to one run id so every write names +the run it belongs to and no process-wide state is involved. +""" + +import logging + +from mlflow.tracking import MlflowClient + +log = logging.getLogger("plugins.mcp") + +# mlflow rejects an oversized value outright. Truncating keeps a completed +# run from being recorded as failed just because its reasoning was long. +_MAX_TAG_CHARS = 5000 +_MAX_PARAM_CHARS = 6000 + + +def resolve_experiment_id(name: str, client: MlflowClient = None) -> str: + """Experiment id for ``name``, creating the experiment if absent. + + Deliberately not ``mlflow.set_experiment``: that mutates the + process-wide active experiment this module exists to stay clear of. + """ + client = client or MlflowClient() + experiment = client.get_experiment_by_name(name) + if experiment is not None: + return experiment.experiment_id + try: + return client.create_experiment(name) + except Exception: + # Lost the race with a concurrent creator; re-read rather than fail. + experiment = client.get_experiment_by_name(name) + return experiment.experiment_id if experiment else "0" + + +class RunTracker: + """MLflow writes bound to a single run id.""" + + def __init__(self, run_id: str, client: MlflowClient = None): + self.run_id = run_id + self._client = client or MlflowClient() + self._terminated = False + + @classmethod + def start(cls, experiment_name: str, run_name: str, + client: MlflowClient = None) -> "RunTracker": + """Mint a run in ``experiment_name`` without activating it.""" + client = client or MlflowClient() + run = client.create_run( + experiment_id=resolve_experiment_id(experiment_name, client), + tags={"mlflow.runName": run_name}, + ) + return cls(run.info.run_id, client) + + @classmethod + def bind(cls, run_id: str, experiment_name: str, run_name: str, + client: MlflowClient = None) -> "RunTracker": + """Track ``run_id``, or mint a run when the caller has none. + + Workflows are normally handed a run id by the orchestrator, which + owns terminating it. Invoked directly (tests, scripts) they mint + and terminate their own. + """ + if run_id: + return cls(run_id, client) + return cls.start(experiment_name, run_name, client) + + @property + def terminated(self) -> bool: + return self._terminated + + def set_tag(self, key: str, value) -> None: + self._write(self._client.set_tag, key, value, _MAX_TAG_CHARS) + + def log_param(self, key: str, value) -> None: + self._write(self._client.log_param, key, value, _MAX_PARAM_CHARS) + + def terminate(self, status: str = "FINISHED") -> None: + """Write the terminal status once; later calls are no-ops.""" + if self._terminated: + return + self._terminated = True + try: + self._client.set_terminated(self.run_id, status) + except Exception as e: + log.warning(f"[MCP] Could not terminate run {self.run_id}: {e}") + + def _write(self, fn, key: str, value, limit: int) -> None: + # Observability writes never fail the run they describe. + try: + fn(self.run_id, key, str(value)[:limit]) + except Exception as e: + log.debug(f"[MCP] Dropped {key} for run {self.run_id}: {e}") + + +def reconcile_orphaned_runs(experiment_names, live_run_ids, + client: MlflowClient = None) -> list: + """Terminate RUNNING runs that no live task owns, returning their ids. + + A run only reaches a terminal status from the task that owns it, so a + process killed mid-flight leaves its run RUNNING with a null end_time + and History renders it as running forever. Called at boot, where the + live cache is empty and every RUNNING row is by definition residue. + """ + client = client or MlflowClient() + experiment_ids = [] + for name in experiment_names: + experiment = client.get_experiment_by_name(name) + if experiment is not None: + experiment_ids.append(experiment.experiment_id) + if not experiment_ids: + return [] + + reconciled = [] + for run in client.search_runs( + experiment_ids=experiment_ids, + filter_string="attributes.status = 'RUNNING'", + ): + run_id = run.info.run_id + if run_id in live_run_ids: + continue + try: + client.set_tag(run_id, "mcp.reconciled", "orphaned") + client.set_terminated(run_id, "KILLED") + reconciled.append(run_id) + except Exception as e: + log.warning(f"[MCP] Could not reconcile orphaned run {run_id}: {e}") + return reconciled diff --git a/app/workflows/author.py b/app/workflows/author.py index d9192d4..705a912 100644 --- a/app/workflows/author.py +++ b/app/workflows/author.py @@ -6,11 +6,11 @@ import sys import mlflow import traceback -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 from plugins.mcp.app.dspy_env import ( ENV_API_BASE, ENV_API_KEY, @@ -141,7 +141,6 @@ class DSPyCalderaFactoryClientWithRAG(dspy.Signature): # Factory function to create tool functions with proper closure def create_tool_function(session, tool_name, tool_description): async def tool_function(**kwargs): - mlflow.set_tag("stage", f"Tool.{tool_name}") result = await session.call_tool(tool_name, kwargs) return result tool_function.__doc__ = tool_description @@ -183,6 +182,14 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, run_ lm_settings = dict(lm_obj) if lm_obj else llm_defaults() max_tool_calls = lm_settings.get("max_tool_calls") or 5 + # Every MLflow write below is addressed to run_id. The fluent API + # routes through a thread-local active-run stack that concurrent + # requests share, so it retagged and terminated each other's runs and + # minted phantom ones whenever the stack was empty. + created_local_run = not run_id + tracker = RunTracker.bind(run_id, _MLFLOW_EXPERIMENT, "MCP Ability Factory") + run_id = tracker.run_id + # Both credentials are checked here rather than at dspy.LM() so the # failure lands before the AsyncExitStack spawns MCP subprocesses. missing = next( @@ -191,26 +198,17 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, run_ if missing: error_msg = f"{missing} is required but not provided. Please set it in the Global Model Configuration." print(f"[MCP] ERROR: {error_msg}") - if not run_id: - run = mlflow.start_run(run_name="MCP Ability Factory") - run_id = run.info.run_id - mlflow.set_tag("status", "failed") - mlflow.set_tag("stage", "error") - mlflow.log_param("error", error_msg) - mlflow.log_param("prompt", adversary_emulation_task) - mlflow.end_run() + tracker.set_tag("status", "failed") + tracker.set_tag("stage", "error") + tracker.log_param("error", error_msg) + tracker.log_param("prompt", adversary_emulation_task) + if created_local_run: + tracker.terminate("FAILED") raise ValueError(error_msg) - # Use the passed-in run_id to continue the MLflow run if provided - created_local_run = False - if not run_id: - run = mlflow.start_run(run_name="MCP Ability Factory") - run_id = run.info.run_id - created_local_run = True - - mlflow.set_tag("status", "running") - mlflow.set_tag("stage", "initializing") - mlflow.log_param("prompt", adversary_emulation_task) + tracker.set_tag("status", "running") + tracker.set_tag("stage", "initializing") + tracker.log_param("prompt", adversary_emulation_task) # Resolve which MCP servers to spawn if not enabled_servers: @@ -225,12 +223,12 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, run_ } } - mlflow.log_param("enabled_servers", ",".join(enabled_servers)) + tracker.log_param("enabled_servers", ",".join(enabled_servers)) # Bump max iters when non-core servers are in the mix (realistic runs need more) if any(name != "caldera_core" for name in enabled_servers) and max_tool_calls < 10: max_tool_calls = 10 - mlflow.log_param("max_tool_calls", max_tool_calls) + tracker.log_param("max_tool_calls", max_tool_calls) try: async with AsyncExitStack() as stack: @@ -244,14 +242,14 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, run_ args=[str(info["path"])], env=get_env(lm_settings), ) - mlflow.set_tag("stage", f"initializing MCP session: {server_name}") + tracker.set_tag("stage", f"initializing MCP session: {server_name}") read, write = await stack.enter_async_context(stdio_client(params)) session = await stack.enter_async_context(ClientSession(read, write)) await session.initialize() sessions.append(session) # Merge tools across all sessions with collision detection - mlflow.set_tag("stage", "listing tools") + tracker.set_tag("stage", "listing tools") seen = {} dspy_tools = [] for server_name, session in zip(enabled_servers, sessions): @@ -265,12 +263,12 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, run_ ) seen[tool.name] = server_name dspy_tools.append(dspy.Tool.from_mcp_tool(session, tool)) - mlflow.log_param("tool_count", len(dspy_tools)) + tracker.log_param("tool_count", len(dspy_tools)) # Use context to set LM for this task/run lm_kwargs = dspy_lm_kwargs_from_settings(lm_settings) with dspy.context(lm=dspy.LM(**lm_kwargs)): - mlflow.set_tag("stage", "creating DSPy ReAct instance") + tracker.set_tag("stage", "creating DSPy ReAct instance") # Resolve CTI context: prefer the orchestrator-supplied string, # fall back to formatting the legacy structured dict. @@ -280,15 +278,15 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, run_ if resolved_cti: signature = DSPyCalderaFactoryClientWithRAG - mlflow.log_param("cti_context_preview", resolved_cti[:1000]) - mlflow.set_tag("cti_context_length", len(resolved_cti)) + tracker.log_param("cti_context_preview", resolved_cti[:1000]) + tracker.set_tag("cti_context_length", len(resolved_cti)) if rag_context: - mlflow.set_tag("cti_search_results_count", len(rag_context.get("search_results", []))) - mlflow.set_tag("cti_detailed_context_count", len(rag_context.get("detailed_context", []))) + tracker.set_tag("cti_search_results_count", len(rag_context.get("search_results", []))) + tracker.set_tag("cti_detailed_context_count", len(rag_context.get("detailed_context", []))) print(f"[MCP] Passing CTI context to LLM ({len(resolved_cti)} chars)") react = dspy.ReAct(signature, tools=dspy_tools, max_iters=max_tool_calls) - mlflow.set_tag("stage", "executing DSPy ReAct with RAG") + tracker.set_tag("stage", "executing DSPy ReAct with RAG") result = await safe_react_acall( react, adversary_emulation_task=adversary_emulation_task, @@ -297,7 +295,7 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, run_ else: signature = DSPyCalderaFactoryClient react = dspy.ReAct(signature, tools=dspy_tools, max_iters=max_tool_calls) - mlflow.set_tag("stage", "executing DSPy ReAct") + tracker.set_tag("stage", "executing DSPy ReAct") result = await safe_react_acall( react, adversary_emulation_task=adversary_emulation_task, @@ -308,21 +306,22 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, run_ # writes are observability for the MLflow UI and the History # tab (which still scrapes tags to reconstruct trajectories # for past runs that have aged out of the live cache). - mlflow.set_tag("stage", "completed") - mlflow.set_tag("status", "complete") - mlflow.set_tag("reasoning", result.reasoning) - mlflow.set_tag("process_result", result.process_result) + tracker.set_tag("stage", "completed") + tracker.set_tag("status", "complete") + tracker.set_tag("reasoning", result.reasoning) + tracker.set_tag("process_result", result.process_result) for k, v in result.trajectory.items(): - mlflow.set_tag(k, json.dumps(v) if isinstance(v, (dict, list)) else str(v)) + tracker.set_tag(k, json.dumps(v) if isinstance(v, (dict, list)) else str(v)) - mlflow.log_param("result_summary", result.process_result) + tracker.log_param("result_summary", result.process_result) print(json.dumps(result.toDict(), indent=4)) - # End the run only if we created it locally + # The orchestrator terminates the run it handed us; a run we + # minted here is ours to close. if created_local_run: - mlflow.end_run() + tracker.terminate("FINISHED") return { "process_result": result.process_result, @@ -334,12 +333,12 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, run_ tb = traceback.format_exc() print("[MCP] Exception occurred:") print(tb) - mlflow.set_tag("status", "failed") - mlflow.set_tag("stage", "error") - mlflow.log_param("error", str(e)) - mlflow.log_param("traceback", tb) + tracker.set_tag("status", "failed") + tracker.set_tag("stage", "error") + tracker.log_param("error", str(e)) + tracker.log_param("traceback", tb) if created_local_run: - mlflow.end_run() + tracker.terminate("FAILED") raise diff --git a/app/workflows/base.py b/app/workflows/base.py index 54caddb..019bb68 100644 --- a/app/workflows/base.py +++ b/app/workflows/base.py @@ -52,8 +52,8 @@ class Workflow: accepted_capabilities: list[str] = field(default_factory=list) # MLflow experiment this workflow's runs belong to. mcp_svc mints the run - # here, and the workflow must set_experiment the same name: mlflow refuses - # start_run(run_id=...) when the active experiment differs from the run's. + # here, and the boot-time orphan sweep only reconciles runs in the + # experiments the registered workflows declare. mlflow_experiment: str = "caldera-mcp-client-1" # Optional plan-then-execute mode. When set, the orchestrator treats the diff --git a/app/workflows/plan_execute.py b/app/workflows/plan_execute.py index 0cea98a..ca8d65e 100644 --- a/app/workflows/plan_execute.py +++ b/app/workflows/plan_execute.py @@ -11,6 +11,7 @@ from contextlib import AsyncExitStack from plugins.mcp.app.config import caldera_connection, llm_defaults, mlflow_settings +from plugins.mcp.app.mlflow_run import RunTracker from plugins.mcp.app.dspy_env import ( ENV_API_BASE, ENV_API_KEY, @@ -174,17 +175,17 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, lm_instance = _build_lm_from_settings(lm_settings) max_tool_calls = lm_settings.get("max_tool_calls") or 5 - # Start or resume MLflow run - if run_id: - mlflow.end_run() # Ensure no active run - mlflow.start_run(run_id=run_id) - else: - run = mlflow.start_run(run_name="MCP Planner Run") - run_id = run.info.run_id + # Every MLflow write below is addressed to run_id. The fluent API + # routes through a thread-local active-run stack that concurrent + # requests share, so it retagged and terminated each other's runs and + # minted phantom ones whenever the stack was empty. + created_local_run = not run_id + tracker = RunTracker.bind(run_id, _MLFLOW_EXPERIMENT, "MCP Planner Run") + run_id = tracker.run_id - mlflow.set_tag("status", "running") - mlflow.set_tag("stage", "initializing") - mlflow.log_param("prompt", adversary_emulation_task) + tracker.set_tag("status", "running") + tracker.set_tag("stage", "initializing") + tracker.log_param("prompt", adversary_emulation_task) # Resolve which MCP servers to spawn if not enabled_servers: @@ -228,12 +229,12 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, "cti_pipeline) in the server registry" ) - mlflow.log_param("enabled_servers", ",".join(enabled_servers)) + tracker.log_param("enabled_servers", ",".join(enabled_servers)) # Bump max iters when non-core servers are in the mix (realistic runs need more) if any(name != "caldera_core" for name in enabled_servers) and max_tool_calls < 10: max_tool_calls = 10 - mlflow.log_param("max_tool_calls", max_tool_calls) + tracker.log_param("max_tool_calls", max_tool_calls) try: async with AsyncExitStack() as stack: @@ -247,14 +248,14 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, args=[str(info["path"])], env=get_env(lm_settings), ) - mlflow.set_tag("stage", f"initializing MCP session: {server_name}") + tracker.set_tag("stage", f"initializing MCP session: {server_name}") read, write = await stack.enter_async_context(stdio_client(params)) session = await stack.enter_async_context(ClientSession(read, write)) await session.initialize() sessions.append(session) # Merge tools across all sessions with collision detection - mlflow.set_tag("stage", "listing tools") + tracker.set_tag("stage", "listing tools") seen = {} dspy_tools = [] for server_name, session in zip(enabled_servers, sessions): @@ -268,11 +269,11 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, ) seen[tool.name] = server_name dspy_tools.append(dspy.Tool.from_mcp_tool(session, tool)) - mlflow.log_param("tool_count", len(dspy_tools)) + tracker.log_param("tool_count", len(dspy_tools)) # Use per-call LM context, honoring lm_obj if provided with dspy.context(lm=lm_instance): - mlflow.set_tag("stage", "creating DSPy ReAct instance") + tracker.set_tag("stage", "creating DSPy ReAct instance") # Resolve CTI context: prefer the orchestrator-supplied string, # fall back to formatting the legacy structured dict. resolved_cti = cti_context @@ -280,20 +281,20 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, resolved_cti = format_rag_context(rag_context) operation_context = format_plan_execute_context(workflow_context) if operation_context: - mlflow.log_param("operation_context_preview", operation_context[:1000]) - mlflow.set_tag("operation_context_length", len(operation_context)) + tracker.log_param("operation_context_preview", operation_context[:1000]) + tracker.set_tag("operation_context_length", len(operation_context)) if resolved_cti: signature = DSPyCalderaPlannerClientWithRAG - mlflow.log_param("cti_context_preview", resolved_cti[:1000]) - mlflow.set_tag("cti_context_length", len(resolved_cti)) + tracker.log_param("cti_context_preview", resolved_cti[:1000]) + tracker.set_tag("cti_context_length", len(resolved_cti)) if rag_context: - mlflow.set_tag("cti_search_results_count", len(rag_context.get("search_results", []))) - mlflow.set_tag("cti_detailed_context_count", len(rag_context.get("detailed_context", []))) + tracker.set_tag("cti_search_results_count", len(rag_context.get("search_results", []))) + tracker.set_tag("cti_detailed_context_count", len(rag_context.get("detailed_context", []))) print(f"[MCP] Passing CTI context to LLM ({len(resolved_cti)} chars)") react = dspy.ReAct(signature, tools=dspy_tools, max_iters=max_tool_calls) - mlflow.set_tag("stage", "executing DSPy ReAct with RAG") + tracker.set_tag("stage", "executing DSPy ReAct with RAG") result = await safe_react_acall( react, adversary_emulation_task=adversary_emulation_task, @@ -304,7 +305,7 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, else: signature = DSPyCalderaPlannerClient react = dspy.ReAct(signature, tools=dspy_tools, max_iters=max_tool_calls) - mlflow.set_tag("stage", "executing DSPy ReAct") + tracker.set_tag("stage", "executing DSPy ReAct") result = await safe_react_acall( react, adversary_emulation_task=adversary_emulation_task, @@ -313,22 +314,25 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, ) if chat_history: - mlflow.set_tag("chat_history_length", len(chat_history)) + tracker.set_tag("chat_history_length", len(chat_history)) # Log outputs and trajectory. The live /status endpoint reads # from mcp_svc's in-memory run cache, not from MLflow; these # writes are observability for the MLflow UI and the History # tab (which still scrapes tags to reconstruct trajectories # for past runs that have aged out of the live cache). - mlflow.set_tag("stage", "completed") - mlflow.set_tag("status", "complete") - mlflow.set_tag("reasoning", result.reasoning) - mlflow.set_tag("process_result", result.process_result) + tracker.set_tag("stage", "completed") + tracker.set_tag("status", "complete") + tracker.set_tag("reasoning", result.reasoning) + tracker.set_tag("process_result", result.process_result) for k, v in result.trajectory.items(): - mlflow.set_tag(k, json.dumps(v) if isinstance(v, (dict, list)) else str(v)) + tracker.set_tag(k, json.dumps(v) if isinstance(v, (dict, list)) else str(v)) - mlflow.log_param("result_summary", result.process_result) - mlflow.end_run() + tracker.log_param("result_summary", result.process_result) + # The orchestrator terminates the run it handed us; a run we + # minted here is ours to close. + if created_local_run: + tracker.terminate("FINISHED") print(json.dumps(result.toDict(), indent=4)) return { "process_result": result.process_result, @@ -340,11 +344,12 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, tb = traceback.format_exc() print("[MCP] Exception occurred:") print(tb) - mlflow.set_tag("status", "failed") - mlflow.set_tag("stage", "error") - mlflow.log_param("error", str(e)) - mlflow.log_param("traceback", tb) - mlflow.end_run() + tracker.set_tag("status", "failed") + tracker.set_tag("stage", "error") + tracker.log_param("error", str(e)) + tracker.log_param("traceback", tb) + if created_local_run: + tracker.terminate("FAILED") raise # Optional streaming updates (if desired for parity) diff --git a/gui/views/mcp_history.vue b/gui/views/mcp_history.vue index 86accdc..a8167c0 100644 --- a/gui/views/mcp_history.vue +++ b/gui/views/mcp_history.vue @@ -69,7 +69,7 @@